fix(sqlite): atomic event queue claim
This commit is contained in:
@@ -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;
|
||||
|
||||
103
crates/adapters/sqlite/src/event_queue/tests.rs
Normal file
103
crates/adapters/sqlite/src/event_queue/tests.rs
Normal 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");
|
||||
}
|
||||
Reference in New Issue
Block a user