PiCloud/crates/manager-core/src/group_queue_repo.rs

//! `GroupQueueRepo` — CRUD over `group_queue_messages`, the §11.6 D3 group-
//! shared durable queue store.
//!
//! The group counterpart to [`crate::queue_repo::QueueRepo`], keyed by
//! `(group_id, collection)` instead of `(app_id, queue_name)`. Producers enqueue
//! via `GroupQueueServiceImpl` (`queue::shared_collection(name).enqueue(...)`);
//! the dispatcher's queue arm claims from here for a materialized shared-queue
//! consumer (competing consumers — every descendant app runs a consumer copy,
//! all claiming this one store with `FOR UPDATE SKIP LOCKED`, so each message is
//! delivered at-most-once across the subtree).
//!
//! §11.6 D3 (dead-lettering): an exhausted shared-queue message moves to the
//! group dead-letter store (`group_dead_letters`, 0068) via [`GroupQueueRepo::dead_letter`]
//! instead of being dropped — operator-visible at GET .../groups/{id}/dead-letters.
//! Deferred: fan-out to a *shared* dead-letter TRIGGER (needs a new trigger kind).

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use picloud_shared::{AdminUserId, GroupId, QueueMessageId};
use sqlx::PgPool;
use uuid::Uuid;

/// Cap on `depth`/`depth_pending` scans (mirrors `queue_repo`).
const QUEUE_DEPTH_SCAN_CAP: i64 = 10_000;

#[derive(Debug, thiserror::Error)]
pub enum GroupQueueRepoError {
    #[error("database error: {0}")]
    Db(#[from] sqlx::Error),
}

/// Insert payload — what `GroupQueueService::enqueue` hands the repo.
#[derive(Debug, Clone)]
pub struct NewGroupQueueMessage {
    pub group_id: GroupId,
    pub collection: String,
    pub payload: serde_json::Value,
    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 ClaimedGroupMessage {
    pub id: QueueMessageId,
    pub group_id: GroupId,
    pub collection: String,
    pub payload: serde_json::Value,
    pub enqueued_at: DateTime<Utc>,
    pub attempt: u32,
    pub max_attempts: u32,
    pub claim_token: Uuid,
}

#[async_trait]
pub trait GroupQueueRepo: Send + Sync {
    async fn enqueue(
        &self,
        msg: NewGroupQueueMessage,
    ) -> Result<QueueMessageId, GroupQueueRepoError>;

    /// Atomic claim of one ready message from `(group_id, collection)` with
    /// `FOR UPDATE SKIP LOCKED` — the competing-consumer primitive. `Ok(None)`
    /// when nothing is claimable.
    async fn claim(
        &self,
        group_id: GroupId,
        collection: &str,
    ) -> Result<Option<ClaimedGroupMessage>, GroupQueueRepoError>;

    /// Handler succeeded: delete the row iff `claim_token` matches.
    async fn ack(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
    ) -> Result<bool, GroupQueueRepoError>;

    /// Handler threw: clear the claim, defer redelivery by `retry_delay`.
    async fn nack(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
        retry_delay: chrono::Duration,
    ) -> Result<bool, GroupQueueRepoError>;

    /// TRANSIENT release (handler never ran): re-queue AND undo `claim`'s
    /// pre-increment of `attempt`, so a non-execution doesn't burn the retry
    /// budget. Mirrors `queue_repo::release`. Floors `attempt` at 0.
    async fn release(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
        retry_delay: chrono::Duration,
    ) -> Result<bool, GroupQueueRepoError>;

    /// §11.6 D3: a message exhausted its attempts — move it to the group dead-
    /// letter store (`group_dead_letters`) and delete it from the live queue, in
    /// one transaction (mirrors `queue_repo::dead_letter`). Filtered by
    /// `claim_token` so a lost lease can't dead-letter a re-claimed message.
    /// Returns the new dead-letter id.
    #[allow(clippy::too_many_arguments)]
    async fn dead_letter(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
        group_id: GroupId,
        collection: &str,
        trigger_id: Option<picloud_shared::TriggerId>,
        script_id: Option<picloud_shared::ScriptId>,
        attempt: u32,
        first_attempt_at: DateTime<Utc>,
        last_error: &str,
    ) -> Result<picloud_shared::DeadLetterId, GroupQueueRepoError>;

    /// Periodic safety net: clear claims older than a shared-queue consumer's
    /// `visibility_timeout_secs`. Returns the number of rows reclaimed.
    async fn reclaim_visibility_timeouts(&self) -> Result<u64, GroupQueueRepoError>;

    async fn depth(&self, group_id: GroupId, collection: &str) -> Result<u64, GroupQueueRepoError>;

    async fn depth_pending(
        &self,
        group_id: GroupId,
        collection: &str,
    ) -> Result<u64, GroupQueueRepoError>;
}

pub struct PostgresGroupQueueRepo {
    pool: PgPool,
}

impl PostgresGroupQueueRepo {
    #[must_use]
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }
}

#[async_trait]
impl GroupQueueRepo for PostgresGroupQueueRepo {
    async fn enqueue(
        &self,
        msg: NewGroupQueueMessage,
    ) -> Result<QueueMessageId, GroupQueueRepoError> {
        let (id,): (Uuid,) = sqlx::query_as(
            "INSERT INTO group_queue_messages ( \
                group_id, collection, payload, deliver_after, \
                max_attempts, enqueued_by_principal \
             ) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id",
        )
        .bind(msg.group_id.into_inner())
        .bind(&msg.collection)
        .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,
        group_id: GroupId,
        collection: &str,
    ) -> Result<Option<ClaimedGroupMessage>, GroupQueueRepoError> {
        let token = Uuid::new_v4();
        let row: Option<ClaimedRow> = sqlx::query_as(
            "UPDATE group_queue_messages \
             SET claim_token = $3, claimed_at = NOW(), attempt = attempt + 1 \
             WHERE id = ( \
                SELECT id FROM group_queue_messages \
                WHERE group_id = $1 AND collection = $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, group_id, collection, payload, enqueued_at, attempt, max_attempts",
        )
        .bind(group_id.into_inner())
        .bind(collection)
        .bind(token)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.map(|r| ClaimedGroupMessage {
            id: r.id.into(),
            group_id: r.group_id.into(),
            collection: r.collection,
            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, GroupQueueRepoError> {
        let res =
            sqlx::query("DELETE FROM group_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, GroupQueueRepoError> {
        let next = Utc::now() + retry_delay;
        let res = sqlx::query(
            "UPDATE group_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, GroupQueueRepoError> {
        // TRANSIENT release (handler never ran): re-queue AND undo `claim`'s
        // pre-increment of `attempt` so a non-execution doesn't burn the retry
        // budget. Mirrors `queue_repo::release`. Floors at 0.
        let next = Utc::now() + retry_delay;
        let res = sqlx::query(
            "UPDATE group_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)
    }

    #[allow(clippy::too_many_arguments)]
    async fn dead_letter(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
        group_id: GroupId,
        collection: &str,
        trigger_id: Option<picloud_shared::TriggerId>,
        script_id: Option<picloud_shared::ScriptId>,
        attempt: u32,
        first_attempt_at: DateTime<Utc>,
        last_error: &str,
    ) -> Result<picloud_shared::DeadLetterId, GroupQueueRepoError> {
        let mut tx = self.pool.begin().await?;
        // Pull the row inside the tx for its payload; filtered by claim_token so
        // a lost lease can't dead-letter a re-claimed message.
        let row: DeadLetterSourceRow = sqlx::query_as(
            "SELECT id, payload, enqueued_at FROM group_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": collection,
            "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 group_dead_letters ( \
                id, group_id, collection, original_event_id, source, op, \
                trigger_id, script_id, payload, attempt_count, \
                first_attempt_at, last_attempt_at, last_error \
             ) VALUES ($1, $2, $3, $4, 'queue', 'receive', $5, $6, $7, $8, $9, NOW(), $10)",
        )
        .bind(dl_id.into_inner())
        .bind(group_id.into_inner())
        .bind(collection)
        .bind(row.id)
        .bind(trigger_id.map(picloud_shared::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 group_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)
    }

    async fn reclaim_visibility_timeouts(&self) -> Result<u64, GroupQueueRepoError> {
        // A shared-queue consumer is a materialized app-owned trigger
        // (`shared = TRUE`, `group_id` on the SOURCE template) whose
        // queue_trigger_details.queue_name IS the shared collection name. Join
        // the group template to recover the owning group + visibility timeout.
        let res = sqlx::query(
            "UPDATE group_queue_messages m \
             SET claim_token = NULL, claimed_at = NULL \
             FROM triggers copy \
             JOIN triggers tmpl ON tmpl.id = copy.materialized_from \
             JOIN queue_trigger_details d ON d.trigger_id = copy.id \
             WHERE tmpl.group_id = m.group_id \
               AND d.queue_name = m.collection \
               AND tmpl.shared = TRUE \
               AND m.claim_token IS NOT NULL \
               AND m.claimed_at < NOW() - (d.visibility_timeout_secs || ' seconds')::INTERVAL",
        )
        .execute(&self.pool)
        .await?;
        Ok(res.rows_affected())
    }

    async fn depth(&self, group_id: GroupId, collection: &str) -> Result<u64, GroupQueueRepoError> {
        let (n,): (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM ( \
                 SELECT 1 FROM group_queue_messages \
                 WHERE group_id = $1 AND collection = $2 LIMIT $3 \
             ) sub",
        )
        .bind(group_id.into_inner())
        .bind(collection)
        .bind(QUEUE_DEPTH_SCAN_CAP)
        .fetch_one(&self.pool)
        .await?;
        Ok(u64::try_from(n).unwrap_or(0))
    }

    async fn depth_pending(
        &self,
        group_id: GroupId,
        collection: &str,
    ) -> Result<u64, GroupQueueRepoError> {
        let (n,): (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM ( \
                 SELECT 1 FROM group_queue_messages \
                 WHERE group_id = $1 AND collection = $2 \
                   AND claim_token IS NULL \
                   AND (deliver_after IS NULL OR deliver_after <= NOW()) LIMIT $3 \
             ) sub",
        )
        .bind(group_id.into_inner())
        .bind(collection)
        .bind(QUEUE_DEPTH_SCAN_CAP)
        .fetch_one(&self.pool)
        .await?;
        Ok(u64::try_from(n).unwrap_or(0))
    }
}

#[derive(sqlx::FromRow)]
struct ClaimedRow {
    id: Uuid,
    group_id: Uuid,
    collection: String,
    payload: serde_json::Value,
    enqueued_at: DateTime<Utc>,
    attempt: i32,
    max_attempts: i32,
}

/// Minimal projection of a to-be-dead-lettered message: its payload + timing.
#[derive(sqlx::FromRow)]
struct DeadLetterSourceRow {
    id: Uuid,
    payload: serde_json::Value,
    enqueued_at: DateTime<Utc>,
}