| 1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
|
| 6 |
use super::{now_ms, Pool}; |
| 7 |
use sqlx::FromRow; |
| 8 |
|
| 9 |
|
| 10 |
pub const MAX_ATTEMPTS: i64 = 8; |
| 11 |
|
| 12 |
#[derive(Clone, Debug, FromRow)] |
| 13 |
pub struct Item { |
| 14 |
pub id: i64, |
| 15 |
pub kind: String, |
| 16 |
pub payload: String, |
| 17 |
pub dedupe: String, |
| 18 |
pub run_after: i64, |
| 19 |
pub attempts: i64, |
| 20 |
pub last_error: String, |
| 21 |
pub claimed_by: String, |
| 22 |
pub claimed_until: i64, |
| 23 |
pub dead: bool, |
| 24 |
pub created_ms: i64, |
| 25 |
} |
| 26 |
|
| 27 |
|
| 28 |
|
| 29 |
pub async fn enqueue(pool: &Pool, kind: &str, payload: &serde_json::Value, dedupe: &str, delay_ms: i64) -> sqlx::Result<bool> { |
| 30 |
let now = now_ms(); |
| 31 |
let mut tx = super::begin_write(pool).await?; |
| 32 |
if !dedupe.is_empty() { |
| 33 |
let live: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE kind = $1 AND dedupe = $2 AND dead = $3") |
| 34 |
.bind(kind) |
| 35 |
.bind(dedupe) |
| 36 |
.bind(false) |
| 37 |
.fetch_one(&mut *tx) |
| 38 |
.await?; |
| 39 |
if live > 0 { |
| 40 |
return Ok(false); |
| 41 |
} |
| 42 |
} |
| 43 |
sqlx::query("INSERT INTO work_queue (kind, payload, dedupe, run_after, created_ms) VALUES ($1, $2, $3, $4, $5)") |
| 44 |
.bind(kind) |
| 45 |
.bind(payload.to_string()) |
| 46 |
.bind(dedupe) |
| 47 |
.bind(now + delay_ms.max(0)) |
| 48 |
.bind(now) |
| 49 |
.execute(&mut *tx) |
| 50 |
.await?; |
| 51 |
tx.commit().await?; |
| 52 |
Ok(true) |
| 53 |
} |
| 54 |
|
| 55 |
|
| 56 |
|
| 57 |
pub async fn claim(pool: &Pool, worker: &str, limit: i64, lease_ms: i64) -> sqlx::Result<Vec<Item>> { |
| 58 |
let now = now_ms(); |
| 59 |
#[cfg(feature = "postgres")] |
| 60 |
let pick = "SELECT id FROM work_queue WHERE dead = $3 AND run_after <= $4 AND claimed_until <= $4 ORDER BY run_after, id LIMIT $5 FOR UPDATE SKIP LOCKED"; |
| 61 |
#[cfg(not(feature = "postgres"))] |
| 62 |
let pick = "SELECT id FROM work_queue WHERE dead = $3 AND run_after <= $4 AND claimed_until <= $4 ORDER BY run_after, id LIMIT $5"; |
| 63 |
let sql = |
| 64 |
format!("UPDATE work_queue SET claimed_by = $1, claimed_until = $2, attempts = attempts + 1 WHERE id IN ({pick}) RETURNING *"); |
| 65 |
let mut tx = pool.begin().await?; |
| 66 |
let items: Vec<Item> = |
| 67 |
sqlx::query_as(&sql).bind(worker).bind(now + lease_ms).bind(false).bind(now).bind(limit).fetch_all(&mut *tx).await?; |
| 68 |
tx.commit().await?; |
| 69 |
Ok(items) |
| 70 |
} |
| 71 |
|
| 72 |
|
| 73 |
pub async fn done(pool: &Pool, id: i64) -> sqlx::Result<()> { |
| 74 |
sqlx::query("DELETE FROM work_queue WHERE id = $1").bind(id).execute(pool).await?; |
| 75 |
Ok(()) |
| 76 |
} |
| 77 |
|
| 78 |
|
| 79 |
|
| 80 |
pub async fn failed(pool: &Pool, item: &Item, error: &str) -> sqlx::Result<()> { |
| 81 |
let dead = item.attempts >= MAX_ATTEMPTS; |
| 82 |
let backoff = (30_000i64 << (item.attempts.clamp(1, 8) - 1)).min(3_600_000); |
| 83 |
sqlx::query("UPDATE work_queue SET last_error = $1, dead = $2, run_after = $3, claimed_by = '', claimed_until = 0 WHERE id = $4") |
| 84 |
.bind(error) |
| 85 |
.bind(dead) |
| 86 |
.bind(now_ms() + backoff) |
| 87 |
.bind(item.id) |
| 88 |
.execute(pool) |
| 89 |
.await?; |
| 90 |
Ok(()) |
| 91 |
} |
| 92 |
|
| 93 |
|
| 94 |
pub async fn counts(pool: &Pool) -> sqlx::Result<(i64, i64)> { |
| 95 |
let waiting: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(false).fetch_one(pool).await?; |
| 96 |
let dead: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(true).fetch_one(pool).await?; |
| 97 |
Ok((waiting, dead)) |
| 98 |
} |