queue.rs 100.0%

62 of 62 lines run, 10 of 20 functions · run #520 on main · run missed

//! The durable work queue (oxygenhub#276, docs/roles.md): items any worker
//! process claims for a lease, retried with backoff, dead after the last
//! attempt. On PostgreSQL a claim skips rows another worker has locked
//! (`FOR UPDATE SKIP LOCKED`); on SQLite the one writer serialises claims.

use super::{now_ms, Pool};
use sqlx::FromRow;

/// Attempts before an item is kept as dead.
pub const MAX_ATTEMPTS: i64 = 8;

#[derive(Clone, Debug, FromRow)]
pub struct Item {
    pub id: i64,
    pub kind: String,
    pub payload: String,
    pub dedupe: String,
    pub run_after: i64,
    pub attempts: i64,
    pub last_error: String,
    pub claimed_by: String,
    pub claimed_until: i64,
    pub dead: bool,
    pub created_ms: i64,
}

/// Queue `kind` with its payload, due after `delay_ms`. With a `dedupe`
/// key, nothing is added while a live item of that kind and key waits.
pub async fn enqueue(pool: &Pool, kind: &str, payload: &serde_json::Value, dedupe: &str, delay_ms: i64) -> sqlx::Result<bool> {
    let now = now_ms();
    let mut tx = super::begin_write(pool).await?;
    if !dedupe.is_empty() {
        let live: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE kind = $1 AND dedupe = $2 AND dead = $3")
            .bind(kind)
            .bind(dedupe)
            .bind(false)
            .fetch_one(&mut *tx)
            .await?;
        if live > 0 {
            return Ok(false);
        }
    }
    sqlx::query("INSERT INTO work_queue (kind, payload, dedupe, run_after, created_ms) VALUES ($1, $2, $3, $4, $5)")
        .bind(kind)
        .bind(payload.to_string())
        .bind(dedupe)
        .bind(now + delay_ms.max(0))
        .bind(now)
        .execute(&mut *tx)
        .await?;
    tx.commit().await?;
    Ok(true)
}

/// Claim up to `limit` due items for `worker`, each for `lease_ms`: an item
/// whose worker died is due again when its lease ends.
pub async fn claim(pool: &Pool, worker: &str, limit: i64, lease_ms: i64) -> sqlx::Result<Vec<Item>> {
    let now = now_ms();
    #[cfg(feature = "postgres")]
    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";
    #[cfg(not(feature = "postgres"))]
    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";
    let sql =
        format!("UPDATE work_queue SET claimed_by = $1, claimed_until = $2, attempts = attempts + 1 WHERE id IN ({pick}) RETURNING *");
    let mut tx = pool.begin().await?;
    let items: Vec<Item> =
        sqlx::query_as(&sql).bind(worker).bind(now + lease_ms).bind(false).bind(now).bind(limit).fetch_all(&mut *tx).await?;
    tx.commit().await?;
    Ok(items)
}

/// The item is done: it leaves the queue.
pub async fn done(pool: &Pool, id: i64) -> sqlx::Result<()> {
    sqlx::query("DELETE FROM work_queue WHERE id = $1").bind(id).execute(pool).await?;
    Ok(())
}

/// The item failed: due again after a backoff that doubles each attempt
/// (30 s, 1 min, 2 min, … at most an hour), or dead after the last.
pub async fn failed(pool: &Pool, item: &Item, error: &str) -> sqlx::Result<()> {
    let dead = item.attempts >= MAX_ATTEMPTS;
    let backoff = (30_000i64 << (item.attempts.clamp(1, 8) - 1)).min(3_600_000);
    sqlx::query("UPDATE work_queue SET last_error = $1, dead = $2, run_after = $3, claimed_by = '', claimed_until = 0 WHERE id = $4")
        .bind(error)
        .bind(dead)
        .bind(now_ms() + backoff)
        .bind(item.id)
        .execute(pool)
        .await?;
    Ok(())
}

/// Items waiting (claimed or not) and items dead.
pub async fn counts(pool: &Pool) -> sqlx::Result<(i64, i64)> {
    let waiting: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(false).fetch_one(pool).await?;
    let dead: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(true).fetch_one(pool).await?;
    Ok((waiting, dead))
}

1		//! The durable work queue (oxygenhub#276, docs/roles.md): items any worker
2		//! process claims for a lease, retried with backoff, dead after the last
3		//! attempt. On PostgreSQL a claim skips rows another worker has locked
4		//! (`FOR UPDATE SKIP LOCKED`); on SQLite the one writer serialises claims.
5
6		use super::{now_ms, Pool};
7		use sqlx::FromRow;
8
9		/// Attempts before an item is kept as dead.
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		/// Queue `kind` with its payload, due after `delay_ms`. With a `dedupe`
28		/// key, nothing is added while a live item of that kind and key waits.
29	4×	pub async fn enqueue(pool: &Pool, kind: &str, payload: &serde_json::Value, dedupe: &str, delay_ms: i64) -> sqlx::Result<bool> {
30	4×	let now = now_ms();
31	4×	let mut tx = super::begin_write(pool).await?;
32	4×	if !dedupe.is_empty() {
33	3×	let live: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE kind = $1 AND dedupe = $2 AND dead = $3")
34	3×	.bind(kind)
35	3×	.bind(dedupe)
36	3×	.bind(false)
37	3×	.fetch_one(&mut *tx)
38	3×	.await?;
39	3×	if live > 0 {
40	1×	return Ok(false);
41	2×	}
42	1×	}
43	3×	sqlx::query("INSERT INTO work_queue (kind, payload, dedupe, run_after, created_ms) VALUES ($1, $2, $3, $4, $5)")
44	3×	.bind(kind)
45	3×	.bind(payload.to_string())
46	3×	.bind(dedupe)
47	3×	.bind(now + delay_ms.max(0))
48	3×	.bind(now)
49	3×	.execute(&mut *tx)
50	3×	.await?;
51	3×	tx.commit().await?;
52	3×	Ok(true)
53	4×	}
54
55		/// Claim up to `limit` due items for `worker`, each for `lease_ms`: an item
56		/// whose worker died is due again when its lease ends.
57	254×	pub async fn claim(pool: &Pool, worker: &str, limit: i64, lease_ms: i64) -> sqlx::Result<Vec<Item>> {
58	254×	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	254×	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	254×	let sql =
64	254×	format!("UPDATE work_queue SET claimed_by = $1, claimed_until = $2, attempts = attempts + 1 WHERE id IN ({pick}) RETURNING *");
65	254×	let mut tx = pool.begin().await?;
66	254×	let items: Vec<Item> =
67	254×	sqlx::query_as(&sql).bind(worker).bind(now + lease_ms).bind(false).bind(now).bind(limit).fetch_all(&mut *tx).await?;
68	254×	tx.commit().await?;
69	254×	Ok(items)
70	254×	}
71
72		/// The item is done: it leaves the queue.
73	1×	pub async fn done(pool: &Pool, id: i64) -> sqlx::Result<()> {
74	1×	sqlx::query("DELETE FROM work_queue WHERE id = $1").bind(id).execute(pool).await?;
75	1×	Ok(())
76	1×	}
77
78		/// The item failed: due again after a backoff that doubles each attempt
79		/// (30 s, 1 min, 2 min, … at most an hour), or dead after the last.
80	3×	pub async fn failed(pool: &Pool, item: &Item, error: &str) -> sqlx::Result<()> {
81	3×	let dead = item.attempts >= MAX_ATTEMPTS;
82	3×	let backoff = (30_000i64 << (item.attempts.clamp(1, 8) - 1)).min(3_600_000);
83	3×	sqlx::query("UPDATE work_queue SET last_error = $1, dead = $2, run_after = $3, claimed_by = '', claimed_until = 0 WHERE id = $4")
84	3×	.bind(error)
85	3×	.bind(dead)
86	3×	.bind(now_ms() + backoff)
87	3×	.bind(item.id)
88	3×	.execute(pool)
89	3×	.await?;
90	3×	Ok(())
91	3×	}
92
93		/// Items waiting (claimed or not) and items dead.
94	7×	pub async fn counts(pool: &Pool) -> sqlx::Result<(i64, i64)> {
95	7×	let waiting: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(false).fetch_one(pool).await?;
96	7×	let dead: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM work_queue WHERE dead = $1").bind(true).fetch_one(pool).await?;
97	7×	Ok((waiting, dead))
98	7×	}

`/`	Search
`g` `c`	Code
`g` `h`	History
`g` `b`	Bookmarks
`g` `t`	Trust
`y`	Permalink to this state (files and trees)
`?`	This sheet

Coverage / crates/hub/src/db/queue.rs 100.0%

Keyboard