Files
PiCloud/crates/manager-core/src/queue_repo.rs
MechaCat02 31e6fc964d fix(queue): don't count never-executed transient releases against max_attempts
`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>
2026-07-13 22:05:32 +02:00

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>,
}