`claim` pre-increments `attempt` before the handler runs, so a real execution failure legitimately counts toward `max_attempts`. But the two TRANSIENT release paths — gate saturation (server at capacity) and script-disabled-at- fire-time — re-queue the message WITHOUT executing it, and previously used the same `nack` that leaves the pre-increment in place. Under sustained overload a message could therefore be dead-lettered after N gate-rejections having run zero times (each rejection burning a retry). Add a non-counting `release` to both `QueueRepo` and `GroupQueueRepo` (re-queue + `attempt = GREATEST(attempt - 1, 0)`, mirroring `nack` otherwise) and route the two transient paths through a new `Dispatcher::q_release`. A real handler failure still uses `nack` and counts. No schema change. Pinned by a new `queue_release` DB test (release undoes the claim increment; nack keeps it) and the updated `disabled_queue_consumer_...` dispatcher test (now asserts a release, not a nack). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
514 lines
18 KiB
Rust
514 lines
18 KiB
Rust
//! `QueueRepo` — CRUD over the `queue_messages` table.
|
|
//!
|
|
//! Producer path: `enqueue(...)` writes one row. Dispatcher hot path:
|
|
//! `claim(...)` does the single-round-trip atomic claim
|
|
//! (`UPDATE ... WHERE id = (SELECT ... FOR UPDATE SKIP LOCKED)`).
|
|
//! Outcome: `ack` deletes the row, `nack` clears the claim with a
|
|
//! `deliver_after` retry delay, `dead_letter` moves the payload to
|
|
//! `dead_letters` and deletes the queue row.
|
|
//!
|
|
//! `reclaim_visibility_timeouts()` is the periodic safety net — a
|
|
//! crashed dispatcher (or one whose handler hung past the visibility
|
|
//! timeout) loses its claim and the message becomes claimable again.
|
|
//!
|
|
//! All methods take `app_id` explicitly; the service layer derives it
|
|
//! from `cx.app_id` and passes it in — repo never reads cx.
|
|
|
|
use async_trait::async_trait;
|
|
use chrono::{DateTime, Utc};
|
|
use picloud_shared::{AdminUserId, AppId, QueueMessageId, TriggerId};
|
|
use sqlx::PgPool;
|
|
use uuid::Uuid;
|
|
|
|
use crate::dead_letter_repo::{DeadLetterRepoError, NewDeadLetter};
|
|
|
|
/// F-P-006: maximum rows scanned by `depth()` / `depth_pending()`.
|
|
/// Scripts calling `queue::depth(name)` per invocation paid an unbounded
|
|
/// COUNT(*); capping the scan at 10k lets the SDK report "10000" as a
|
|
/// floor for very-deep queues without ever touching more than that many
|
|
/// index rows.
|
|
const QUEUE_DEPTH_SCAN_CAP: i64 = 10_000;
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum QueueRepoError {
|
|
#[error("database error: {0}")]
|
|
Db(#[from] sqlx::Error),
|
|
|
|
/// JSON serialization failure on the payload (extremely rare; only
|
|
/// surfaces if a non-serializable Value somehow reached the repo).
|
|
#[error("payload serialization: {0}")]
|
|
Payload(String),
|
|
|
|
/// Surfaces from the dead_letter() helper if the dead-letter insert
|
|
/// fails.
|
|
#[error("dead-letter write: {0}")]
|
|
DeadLetter(String),
|
|
}
|
|
|
|
/// Insert payload — what `QueueService::enqueue` hands the repo.
|
|
#[derive(Debug, Clone)]
|
|
pub struct NewQueueMessage {
|
|
pub app_id: AppId,
|
|
pub queue_name: String,
|
|
pub payload: serde_json::Value,
|
|
/// `Some` → defer until `NOW() + delay`; `None` → immediate.
|
|
pub deliver_after: Option<DateTime<Utc>>,
|
|
pub max_attempts: u32,
|
|
pub enqueued_by_principal: Option<AdminUserId>,
|
|
}
|
|
|
|
/// One claimed message ready for handler dispatch.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ClaimedMessage {
|
|
pub id: QueueMessageId,
|
|
pub app_id: AppId,
|
|
pub queue_name: String,
|
|
pub payload: serde_json::Value,
|
|
pub enqueued_at: DateTime<Utc>,
|
|
pub attempt: u32,
|
|
pub max_attempts: u32,
|
|
pub claim_token: Uuid,
|
|
}
|
|
|
|
/// Aggregate counts for the dashboard's read-only queue overview.
|
|
#[derive(Debug, Clone, Copy, Default)]
|
|
pub struct QueueStats {
|
|
/// `COUNT(*)` for `(app_id, queue_name)`.
|
|
pub total: u64,
|
|
/// Claimable now — `claim_token IS NULL AND (deliver_after IS NULL OR <= NOW())`.
|
|
pub pending: u64,
|
|
/// Currently leased to a dispatcher — `claim_token IS NOT NULL`.
|
|
pub claimed: u64,
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait QueueRepo: Send + Sync {
|
|
/// Insert one row. Returns the new message id.
|
|
async fn enqueue(&self, msg: NewQueueMessage) -> Result<QueueMessageId, QueueRepoError>;
|
|
|
|
/// Atomic claim of a single ready message from `(app_id, queue_name)`.
|
|
/// Returns `Ok(None)` if nothing is claimable.
|
|
async fn claim(
|
|
&self,
|
|
app_id: AppId,
|
|
queue_name: &str,
|
|
) -> Result<Option<ClaimedMessage>, QueueRepoError>;
|
|
|
|
/// Handler succeeded: delete the row IFF the claim_token matches
|
|
/// (so a stale dispatcher whose lease expired doesn't accidentally
|
|
/// delete a re-claimed message). Returns `true` if 1 row deleted.
|
|
async fn ack(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
) -> Result<bool, QueueRepoError>;
|
|
|
|
/// Handler threw: clear the claim, defer redelivery by `retry_delay`.
|
|
/// IFF the claim_token matches. Returns `true` on success.
|
|
async fn nack(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
retry_delay: chrono::Duration,
|
|
) -> Result<bool, QueueRepoError>;
|
|
|
|
/// TRANSIENT release: the handler never ran (gate saturated, or the script
|
|
/// was disabled at fire time), so re-queue AND undo `claim`'s pre-increment
|
|
/// of `attempt` — a non-execution must not burn the retry budget (else a
|
|
/// message could be dead-lettered under sustained overload having executed
|
|
/// zero times). Floors `attempt` at 0. IFF the claim_token matches.
|
|
async fn release(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
retry_delay: chrono::Duration,
|
|
) -> Result<bool, QueueRepoError>;
|
|
|
|
/// Periodic safety net: clear claims whose `claimed_at` exceeds the
|
|
/// per-queue `visibility_timeout_secs` (joined from
|
|
/// `queue_trigger_details`). Returns the number of rows reclaimed.
|
|
async fn reclaim_visibility_timeouts(&self) -> Result<u64, QueueRepoError>;
|
|
|
|
/// `queue::depth(name)` — total rows in the queue.
|
|
async fn depth(&self, app_id: AppId, queue_name: &str) -> Result<u64, QueueRepoError>;
|
|
|
|
/// `queue::depth_pending(name)` — currently claimable rows.
|
|
async fn depth_pending(&self, app_id: AppId, queue_name: &str) -> Result<u64, QueueRepoError>;
|
|
|
|
/// Dashboard read-only view: every distinct queue name in `app_id`
|
|
/// with its aggregate counts.
|
|
async fn list_for_app(
|
|
&self,
|
|
app_id: AppId,
|
|
) -> Result<Vec<(String, QueueStats)>, QueueRepoError>;
|
|
|
|
/// Dispatcher dead-letter helper: move a claimed message to
|
|
/// `dead_letters` and delete it from `queue_messages`, in one
|
|
/// transaction. The dispatcher's existing `fan_out_dead_letter`
|
|
/// fires registered `dead_letter` triggers off the new row.
|
|
#[allow(clippy::too_many_arguments)]
|
|
async fn dead_letter(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
app_id: AppId,
|
|
queue_name: &str,
|
|
trigger_id: Option<TriggerId>,
|
|
script_id: Option<picloud_shared::ScriptId>,
|
|
attempt: u32,
|
|
first_attempt_at: DateTime<Utc>,
|
|
last_error: &str,
|
|
) -> Result<picloud_shared::DeadLetterId, QueueRepoError>;
|
|
}
|
|
|
|
pub struct PostgresQueueRepo {
|
|
pool: PgPool,
|
|
}
|
|
|
|
impl PostgresQueueRepo {
|
|
#[must_use]
|
|
pub fn new(pool: PgPool) -> Self {
|
|
Self { pool }
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl QueueRepo for PostgresQueueRepo {
|
|
async fn enqueue(&self, msg: NewQueueMessage) -> Result<QueueMessageId, QueueRepoError> {
|
|
let (id,): (Uuid,) = sqlx::query_as(
|
|
"INSERT INTO queue_messages ( \
|
|
app_id, queue_name, payload, deliver_after, \
|
|
max_attempts, enqueued_by_principal \
|
|
) VALUES ($1, $2, $3, $4, $5, $6) \
|
|
RETURNING id",
|
|
)
|
|
.bind(msg.app_id.into_inner())
|
|
.bind(&msg.queue_name)
|
|
.bind(&msg.payload)
|
|
.bind(msg.deliver_after)
|
|
.bind(i32::try_from(msg.max_attempts).unwrap_or(3))
|
|
.bind(msg.enqueued_by_principal.map(AdminUserId::into_inner))
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
Ok(id.into())
|
|
}
|
|
|
|
async fn claim(
|
|
&self,
|
|
app_id: AppId,
|
|
queue_name: &str,
|
|
) -> Result<Option<ClaimedMessage>, QueueRepoError> {
|
|
let token = Uuid::new_v4();
|
|
let row: Option<ClaimedRow> = sqlx::query_as(
|
|
"UPDATE queue_messages \
|
|
SET claim_token = $3, claimed_at = NOW(), attempt = attempt + 1 \
|
|
WHERE id = ( \
|
|
SELECT id FROM queue_messages \
|
|
WHERE app_id = $1 AND queue_name = $2 \
|
|
AND claim_token IS NULL \
|
|
AND (deliver_after IS NULL OR deliver_after <= NOW()) \
|
|
ORDER BY enqueued_at \
|
|
FOR UPDATE SKIP LOCKED \
|
|
LIMIT 1 \
|
|
) \
|
|
RETURNING id, app_id, queue_name, payload, enqueued_at, attempt, max_attempts",
|
|
)
|
|
.bind(app_id.into_inner())
|
|
.bind(queue_name)
|
|
.bind(token)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
Ok(row.map(|r| ClaimedMessage {
|
|
id: r.id.into(),
|
|
app_id: r.app_id.into(),
|
|
queue_name: r.queue_name,
|
|
payload: r.payload,
|
|
enqueued_at: r.enqueued_at,
|
|
attempt: u32::try_from(r.attempt).unwrap_or(1),
|
|
max_attempts: u32::try_from(r.max_attempts).unwrap_or(3),
|
|
claim_token: token,
|
|
}))
|
|
}
|
|
|
|
async fn ack(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
) -> Result<bool, QueueRepoError> {
|
|
let res = sqlx::query("DELETE FROM queue_messages WHERE id = $1 AND claim_token = $2")
|
|
.bind(message_id.into_inner())
|
|
.bind(claim_token)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
Ok(res.rows_affected() == 1)
|
|
}
|
|
|
|
async fn nack(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
retry_delay: chrono::Duration,
|
|
) -> Result<bool, QueueRepoError> {
|
|
let next = Utc::now() + retry_delay;
|
|
let res = sqlx::query(
|
|
"UPDATE queue_messages \
|
|
SET claim_token = NULL, claimed_at = NULL, deliver_after = $3 \
|
|
WHERE id = $1 AND claim_token = $2",
|
|
)
|
|
.bind(message_id.into_inner())
|
|
.bind(claim_token)
|
|
.bind(next)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
Ok(res.rows_affected() == 1)
|
|
}
|
|
|
|
async fn release(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
retry_delay: chrono::Duration,
|
|
) -> Result<bool, QueueRepoError> {
|
|
let next = Utc::now() + retry_delay;
|
|
let res = sqlx::query(
|
|
"UPDATE queue_messages \
|
|
SET claim_token = NULL, claimed_at = NULL, deliver_after = $3, \
|
|
attempt = GREATEST(attempt - 1, 0) \
|
|
WHERE id = $1 AND claim_token = $2",
|
|
)
|
|
.bind(message_id.into_inner())
|
|
.bind(claim_token)
|
|
.bind(next)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
Ok(res.rows_affected() == 1)
|
|
}
|
|
|
|
async fn reclaim_visibility_timeouts(&self) -> Result<u64, QueueRepoError> {
|
|
let res = sqlx::query(
|
|
"UPDATE queue_messages m \
|
|
SET claim_token = NULL, claimed_at = NULL \
|
|
FROM triggers t \
|
|
JOIN queue_trigger_details d ON d.trigger_id = t.id \
|
|
WHERE m.app_id = t.app_id \
|
|
AND m.queue_name = d.queue_name \
|
|
AND m.claim_token IS NOT NULL \
|
|
AND m.claimed_at < NOW() - (d.visibility_timeout_secs || ' seconds')::INTERVAL \
|
|
AND t.kind = 'queue' AND t.enabled = TRUE",
|
|
)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
Ok(res.rows_affected())
|
|
}
|
|
|
|
async fn depth(&self, app_id: AppId, queue_name: &str) -> Result<u64, QueueRepoError> {
|
|
// F-P-006: cap the scan at QUEUE_DEPTH_SCAN_CAP rows. Scripts
|
|
// calling queue::depth(name) per invocation were doing
|
|
// unbounded COUNT(*) over the queue_messages partial index;
|
|
// on a backed-up queue with millions of rows that meant a
|
|
// full-partition scan every call. The capped subquery means
|
|
// the worst-case cost is bounded; callers needing exact
|
|
// depth on huge queues should use the admin dashboard which
|
|
// tolerates the slower COUNT for its read frequency.
|
|
let (n,): (i64,) = sqlx::query_as(
|
|
"SELECT COUNT(*) FROM ( \
|
|
SELECT 1 FROM queue_messages \
|
|
WHERE app_id = $1 AND queue_name = $2 \
|
|
LIMIT $3 \
|
|
) sub",
|
|
)
|
|
.bind(app_id.into_inner())
|
|
.bind(queue_name)
|
|
.bind(QUEUE_DEPTH_SCAN_CAP)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
|
|
async fn depth_pending(&self, app_id: AppId, queue_name: &str) -> Result<u64, QueueRepoError> {
|
|
let (n,): (i64,) = sqlx::query_as(
|
|
"SELECT COUNT(*) FROM ( \
|
|
SELECT 1 FROM queue_messages \
|
|
WHERE app_id = $1 AND queue_name = $2 \
|
|
AND claim_token IS NULL \
|
|
AND (deliver_after IS NULL OR deliver_after <= NOW()) \
|
|
LIMIT $3 \
|
|
) sub",
|
|
)
|
|
.bind(app_id.into_inner())
|
|
.bind(queue_name)
|
|
.bind(QUEUE_DEPTH_SCAN_CAP)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
Ok(u64::try_from(n).unwrap_or(0))
|
|
}
|
|
|
|
async fn list_for_app(
|
|
&self,
|
|
app_id: AppId,
|
|
) -> Result<Vec<(String, QueueStats)>, QueueRepoError> {
|
|
let rows: Vec<QueueAggRow> = sqlx::query_as(
|
|
"SELECT queue_name, \
|
|
COUNT(*) AS total, \
|
|
COUNT(*) FILTER ( \
|
|
WHERE claim_token IS NULL \
|
|
AND (deliver_after IS NULL OR deliver_after <= NOW()) \
|
|
) AS pending, \
|
|
COUNT(*) FILTER (WHERE claim_token IS NOT NULL) AS claimed \
|
|
FROM queue_messages \
|
|
WHERE app_id = $1 \
|
|
GROUP BY queue_name \
|
|
ORDER BY queue_name",
|
|
)
|
|
.bind(app_id.into_inner())
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
Ok(rows
|
|
.into_iter()
|
|
.map(|r| {
|
|
(
|
|
r.queue_name,
|
|
QueueStats {
|
|
total: u64::try_from(r.total).unwrap_or(0),
|
|
pending: u64::try_from(r.pending).unwrap_or(0),
|
|
claimed: u64::try_from(r.claimed).unwrap_or(0),
|
|
},
|
|
)
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
async fn dead_letter(
|
|
&self,
|
|
message_id: QueueMessageId,
|
|
claim_token: Uuid,
|
|
app_id: AppId,
|
|
queue_name: &str,
|
|
trigger_id: Option<TriggerId>,
|
|
script_id: Option<picloud_shared::ScriptId>,
|
|
attempt: u32,
|
|
first_attempt_at: DateTime<Utc>,
|
|
last_error: &str,
|
|
) -> Result<picloud_shared::DeadLetterId, QueueRepoError> {
|
|
let mut tx = self.pool.begin().await?;
|
|
// Pull the row inside the transaction; we need the original
|
|
// payload for the dead_letters row. Filtered by claim_token so
|
|
// a lost lease can't dead-letter a re-claimed message.
|
|
let row: QueueRowMinimal = sqlx::query_as(
|
|
"SELECT id, payload, enqueued_at FROM queue_messages \
|
|
WHERE id = $1 AND claim_token = $2",
|
|
)
|
|
.bind(message_id.into_inner())
|
|
.bind(claim_token)
|
|
.fetch_one(&mut *tx)
|
|
.await?;
|
|
|
|
let payload = serde_json::json!({
|
|
"source": "queue",
|
|
"queue_name": queue_name,
|
|
"message": row.payload,
|
|
"enqueued_at": row.enqueued_at,
|
|
"attempt": attempt,
|
|
"message_id": row.id.to_string(),
|
|
});
|
|
|
|
let dl_id = picloud_shared::DeadLetterId::new();
|
|
sqlx::query(
|
|
"INSERT INTO dead_letters ( \
|
|
id, app_id, original_event_id, source, op, trigger_id, script_id, \
|
|
payload, attempt_count, first_attempt_at, last_attempt_at, last_error \
|
|
) VALUES ($1, $2, $3, 'queue', 'receive', $4, $5, $6, $7, $8, NOW(), $9)",
|
|
)
|
|
.bind(dl_id.into_inner())
|
|
.bind(app_id.into_inner())
|
|
.bind(row.id)
|
|
.bind(trigger_id.map(TriggerId::into_inner))
|
|
.bind(script_id.map(picloud_shared::ScriptId::into_inner))
|
|
.bind(&payload)
|
|
.bind(i32::try_from(attempt).unwrap_or(0))
|
|
.bind(first_attempt_at)
|
|
.bind(last_error)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
sqlx::query("DELETE FROM queue_messages WHERE id = $1 AND claim_token = $2")
|
|
.bind(message_id.into_inner())
|
|
.bind(claim_token)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
tx.commit().await?;
|
|
Ok(dl_id)
|
|
}
|
|
}
|
|
|
|
/// Bridge between this repo and the existing `DeadLetterRepo`. Not used
|
|
/// by `dead_letter()` directly (that path inlines its own INSERT for one-
|
|
/// transaction atomicity), but exposed for callers that already have a
|
|
/// `DeadLetterRepo` and want to keep one writer.
|
|
#[must_use]
|
|
pub fn build_new_dead_letter_from_claimed(
|
|
msg: &ClaimedMessage,
|
|
trigger_id: Option<TriggerId>,
|
|
script_id: Option<picloud_shared::ScriptId>,
|
|
last_error: String,
|
|
first_attempt_at: DateTime<Utc>,
|
|
) -> NewDeadLetter {
|
|
let payload = serde_json::json!({
|
|
"source": "queue",
|
|
"queue_name": msg.queue_name,
|
|
"message": msg.payload,
|
|
"enqueued_at": msg.enqueued_at,
|
|
"attempt": msg.attempt,
|
|
"message_id": msg.id.to_string(),
|
|
});
|
|
NewDeadLetter {
|
|
app_id: msg.app_id,
|
|
original_event_id: msg.id.into_inner(),
|
|
source: "queue".into(),
|
|
op: "receive".into(),
|
|
trigger_id,
|
|
script_id,
|
|
payload,
|
|
attempt_count: msg.attempt,
|
|
first_attempt_at,
|
|
last_attempt_at: Utc::now(),
|
|
last_error,
|
|
}
|
|
}
|
|
|
|
/// Surface to bridge DeadLetterRepoError → QueueRepoError when callers
|
|
/// route the dead-letter write through the standard repo instead of the
|
|
/// inline `dead_letter()` path.
|
|
impl From<DeadLetterRepoError> for QueueRepoError {
|
|
fn from(e: DeadLetterRepoError) -> Self {
|
|
Self::DeadLetter(e.to_string())
|
|
}
|
|
}
|
|
|
|
#[derive(sqlx::FromRow)]
|
|
struct ClaimedRow {
|
|
id: Uuid,
|
|
app_id: Uuid,
|
|
queue_name: String,
|
|
payload: serde_json::Value,
|
|
enqueued_at: DateTime<Utc>,
|
|
attempt: i32,
|
|
max_attempts: i32,
|
|
}
|
|
|
|
#[derive(sqlx::FromRow)]
|
|
struct QueueAggRow {
|
|
queue_name: String,
|
|
total: i64,
|
|
pending: i64,
|
|
claimed: i64,
|
|
}
|
|
|
|
#[derive(sqlx::FromRow)]
|
|
struct QueueRowMinimal {
|
|
id: Uuid,
|
|
payload: serde_json::Value,
|
|
enqueued_at: DateTime<Utc>,
|
|
}
|