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).
//!
//! Deferred (documented): shared-queue dead-lettering — an exhausted message is
//! nacked with backoff by the dispatcher (never silently dropped), but there is
//! no group dead-letter store yet.

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>;

    /// Drop a message that exhausted its attempts (no group dead-letter store
    /// yet — see the module deferral note). Deletes iff `claim_token` matches.
    async fn drop_exhausted(
        &self,
        message_id: QueueMessageId,
        claim_token: Uuid,
    ) -> Result<bool, 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 drop_exhausted(
        &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 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,
}