fix(sqlite): atomic event queue claim

This commit is contained in:
2026-10-25 23:39:04 +01:00
parent b7e69b69dc
commit d073922308
2 changed files with 124 additions and 35 deletions

View File

@@ -127,51 +127,34 @@ struct QueueRow {
id: i64,
payload: String,
attempts: i32,
next_attempt_at: String,
}
/// Claims a batch in one atomic statement. A read-then-write transaction can
/// fail with SQLITE_BUSY_SNAPSHOT (517) when another connection commits in
/// between, which `busy_timeout` cannot retry.
async fn claim_batch(
pool: &SqlitePool,
config: &DbEventQueueConfig,
) -> Result<Vec<QueueRow>, DomainError> {
let mut tx = pool
.begin()
.await
.map_err(|e| DomainError::InfrastructureError(format!("begin tx: {e}")))?;
let rows = sqlx::query_as::<_, QueueRow>(
"SELECT id, payload, attempts FROM event_queue
WHERE status = 'pending'
AND next_attempt_at <= strftime('%Y-%m-%dT%H:%M:%SZ', 'now')
ORDER BY next_attempt_at ASC
LIMIT ?",
let mut rows = sqlx::query_as::<_, QueueRow>(
"UPDATE event_queue SET status = 'processing'
WHERE id IN (
SELECT id FROM event_queue
WHERE status = 'pending'
AND next_attempt_at <= strftime('%Y-%m-%dT%H:%M:%SZ', 'now')
ORDER BY next_attempt_at ASC
LIMIT ?
)
RETURNING id, payload, attempts, next_attempt_at",
)
.bind(config.batch_size)
.fetch_all(&mut *tx)
.fetch_all(pool)
.await
.map_err(|e| DomainError::InfrastructureError(format!("select pending: {e}")))?;
if rows.is_empty() {
tx.rollback().await.ok();
return Ok(vec![]);
}
let placeholders = rows.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let sql = format!(
"UPDATE event_queue SET status = 'processing' WHERE id IN ({})",
placeholders
);
let mut q = sqlx::query(&sql);
for r in &rows {
q = q.bind(r.id);
}
q.execute(&mut *tx)
.await
.map_err(|e| DomainError::InfrastructureError(format!("mark processing: {e}")))?;
tx.commit()
.await
.map_err(|e| DomainError::InfrastructureError(format!("commit claim: {e}")))?;
.map_err(|e| DomainError::InfrastructureError(format!("claim batch: {e}")))?;
// RETURNING order is unspecified; deliver oldest first.
rows.sort_by(|a, b| (&a.next_attempt_at, a.id).cmp(&(&b.next_attempt_at, b.id)));
Ok(rows)
}
@@ -243,3 +226,6 @@ fn backoff_seconds(attempts: i32) -> i64 {
let base: i64 = 5 * (1i64 << attempts.min(6));
base.min(300)
}
#[cfg(test)]
mod tests;

View File

@@ -0,0 +1,103 @@
//! Inline (not under `tests/`) because `claim_batch` is private and the race
//! only shows on its claim path; two pools on a file-backed WAL database are
//! needed, as in production.
use std::str::FromStr;
use sqlx::SqlitePool;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
use super::{DbEventQueueConfig, claim_batch};
async fn pool_on(path: &std::path::Path) -> SqlitePool {
let opts = SqliteConnectOptions::from_str(&format!("sqlite://{}", path.display()))
.unwrap()
.create_if_missing(true)
.journal_mode(SqliteJournalMode::Wal)
.busy_timeout(std::time::Duration::from_secs(5));
SqlitePoolOptions::new()
.max_connections(4)
.connect_with(opts)
.await
.unwrap()
}
fn config(batch_size: i64) -> DbEventQueueConfig {
DbEventQueueConfig {
poll_interval_ms: 10,
batch_size,
max_attempts: 3,
}
}
async fn insert(pool: &SqlitePool, n: usize) {
for i in 0..n {
sqlx::query("INSERT INTO event_queue (event_type, payload) VALUES ('t', ?)")
.bind(format!("p{i}"))
.execute(pool)
.await
.unwrap();
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn claiming_while_another_connection_writes_does_not_fail() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("q.db");
let claimer = pool_on(&path).await;
crate::migrate(&claimer).await.unwrap();
let writer = pool_on(&path).await;
insert(&claimer, 50).await;
let writer_task = {
let writer = writer.clone();
tokio::spawn(async move {
for i in 0..2000 {
sqlx::query("INSERT INTO event_queue (event_type, payload) VALUES ('t', ?)")
.bind(format!("w{i}"))
.execute(&writer)
.await
.unwrap();
}
})
};
let cfg = config(5);
let mut claimed = 0;
while !writer_task.is_finished() {
claimed += claim_batch(&claimer, &cfg)
.await
.expect("claim must not fail while another connection writes")
.len();
}
writer_task.await.unwrap();
assert!(claimed > 0);
}
#[tokio::test]
async fn claimed_rows_are_marked_processing_and_not_claimed_twice() {
let dir = tempfile::tempdir().unwrap();
let pool = pool_on(&dir.path().join("q.db")).await;
crate::migrate(&pool).await.unwrap();
insert(&pool, 5).await;
let cfg = config(3);
let first = claim_batch(&pool, &cfg).await.unwrap();
assert_eq!(first.len(), 3);
let second = claim_batch(&pool, &cfg).await.unwrap();
assert_eq!(second.len(), 2);
assert!(claim_batch(&pool, &cfg).await.unwrap().is_empty());
let mut ids: Vec<i64> = first.iter().chain(&second).map(|r| r.id).collect();
ids.sort();
ids.dedup();
assert_eq!(ids.len(), 5, "no row claimed twice");
let (processing,): (i64,) =
sqlx::query_as("SELECT COUNT(*) FROM event_queue WHERE status = 'processing'")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(processing, 5);
let payloads: Vec<&str> = first.iter().map(|r| r.payload.as_str()).collect();
assert_eq!(payloads, ["p0", "p1", "p2"], "oldest first");
}