PiCloud/crates/manager-core/tests/outbox_reclaim.rs

//! A dispatcher that dies mid-dispatch must not strand its claimed outbox rows.
//!
//! The dispatcher claims a row, executes it, then deletes it (success) or
//! reschedules it (failure) — both of which clear the claim. If the PROCESS DIES
//! in between, neither runs. And `claim_due` only ever selects
//! `claimed_at IS NULL`, so before `reclaim_stale_claims` that row was stranded
//! FOREVER: its trigger never fired, and no retry could notice. The outbox is
//! the universal trigger path (kv/docs/files/cron/pubsub/email/invoke_async/
//! dead-letter), so a crash or restart mid-dispatch silently lost all of it.
//!
//! Every other claim-based store — `queue_messages`, `group_queue_messages`,
//! `workflow_steps` — already had this reclaimer. The outbox was the one gap.
//!
//! Skips cleanly when `DATABASE_URL` is unset.

use picloud_manager_core::outbox_repo::{
    NewOutboxRow, OutboxRepo, OutboxSourceKind, PostgresOutboxRepo,
};
use picloud_shared::AppId;
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use uuid::Uuid;

async fn pool_or_skip() -> Option<PgPool> {
    let Ok(url) = std::env::var("DATABASE_URL") else {
        eprintln!("outbox_reclaim: DATABASE_URL unset — skipping");
        return None;
    };
    let pool = PgPoolOptions::new()
        .max_connections(3)
        .connect(&url)
        .await
        .expect("connect");
    sqlx::migrate!("./migrations")
        .run(&pool)
        .await
        .expect("migrate");
    Some(pool)
}

async fn mk_app(pool: &PgPool) -> Uuid {
    let uniq = Uuid::new_v4().simple().to_string();
    let (group,): (Uuid,) =
        sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id")
            .bind(format!("or-grp-{uniq}"))
            .fetch_one(pool)
            .await
            .expect("group");
    let (app,): (Uuid,) =
        sqlx::query_as("INSERT INTO apps (slug, name, group_id) VALUES ($1, $1, $2) RETURNING id")
            .bind(format!("or-app-{uniq}"))
            .bind(group)
            .fetch_one(pool)
            .await
            .expect("app");
    app
}

/// Exactly `claim_due`'s predicate: a row is dispatchable iff its claim is clear.
async fn claimable(pool: &PgPool, id: Uuid) -> bool {
    let (n,): (i64,) =
        sqlx::query_as("SELECT COUNT(*) FROM outbox WHERE id = $1 AND claimed_at IS NULL")
            .bind(id)
            .fetch_one(pool)
            .await
            .expect("claimable");
    n == 1
}

async fn attempt_count(pool: &PgPool, id: Uuid) -> i32 {
    let (n,): (i32,) = sqlx::query_as("SELECT attempt_count FROM outbox WHERE id = $1")
        .bind(id)
        .fetch_one(pool)
        .await
        .expect("attempt_count");
    n
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_stranded_claim_is_reclaimed_without_burning_the_retry_budget() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let repo = PostgresOutboxRepo::new(pool.clone());
    let app = mk_app(&pool).await;

    let id = repo
        .insert(NewOutboxRow {
            app_id: AppId::from(app),
            source_kind: OutboxSourceKind::Kv,
            trigger_id: None,
            script_id: None,
            reply_to: None,
            payload: serde_json::json!({}),
            origin_principal: None,
            trigger_depth: 0,
            root_execution_id: None,
        })
        .await
        .expect("insert");

    // A dispatcher claims it. (Asserted against `claim_due`'s own predicate
    // rather than by calling it: the dev DB is shared, and a real `claim_due`
    // would claim other tests' rows out from under them.)
    sqlx::query("UPDATE outbox SET claimed_at = NOW(), claimed_by = 'dispatcher-a' WHERE id = $1")
        .bind(id)
        .execute(&pool)
        .await
        .expect("claim");
    assert!(
        !claimable(&pool, id).await,
        "a claimed row is invisible to every other dispatcher — that is the trap"
    );

    // …and then the process dies. Nothing clears the claim, so before the
    // reclaimer existed this row would sit here forever.
    sqlx::query("UPDATE outbox SET claimed_at = NOW() - INTERVAL '20 minutes' WHERE id = $1")
        .bind(id)
        .execute(&pool)
        .await
        .expect("backdate");

    let n = repo.reclaim_stale_claims(600).await.expect("reclaim");
    assert!(n >= 1, "the stale claim must be reclaimed");
    assert!(
        claimable(&pool, id).await,
        "the reclaimed row is dispatchable again — the trigger fires after all"
    );
    assert_eq!(
        attempt_count(&pool, id).await,
        0,
        "the handler never ran, so a reclaim must NOT consume the retry budget — \
         otherwise repeated restarts would dead-letter an event that executed zero times"
    );

    // A FRESH claim must not be reclaimed out from under a LIVE dispatcher.
    sqlx::query("UPDATE outbox SET claimed_at = NOW(), claimed_by = 'dispatcher-b' WHERE id = $1")
        .bind(id)
        .execute(&pool)
        .await
        .expect("re-claim");
    assert_eq!(
        repo.reclaim_stale_claims(600).await.expect("reclaim again"),
        0,
        "a claim within the timeout belongs to a live dispatcher and must be left alone"
    );

    sqlx::query("DELETE FROM apps WHERE id = $1")
        .bind(app)
        .execute(&pool)
        .await
        .expect("cleanup");
}