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

//! Audit fix #6: a data write and the trigger fan-out it produces must commit
//! ATOMICALLY.
//!
//! Before this, a service wrote the row (commit #1), then asked the emitter to
//! resolve matching triggers and insert outbox rows (commit #2). If commit #2
//! failed, the row was permanently in the store with its trigger having never
//! fired — unobservable to the caller, unrecoverable by any retry, and only
//! logged. `PostgresKvWriter` runs both on ONE connection in ONE transaction, so
//! an emit failure rolls the write back and surfaces as an error.
//!
//! Fault injection: a Postgres BEFORE-INSERT trigger on `outbox` that raises,
//! scoped to THIS test's freshly-minted `app_id`, so it cannot affect any test
//! running in parallel against the same database. Skips when `DATABASE_URL` is
//! unset.

use std::sync::Arc;

use picloud_manager_core::atomic_write::{
    GroupDocsTarget, GroupDocsWriter, GroupKvTarget, GroupKvWriter, KvWriter,
    PostgresGroupDocsWriter, PostgresGroupKvWriter, PostgresKvWriter,
};
use picloud_manager_core::group_quota::GroupWriteQuota;
use picloud_shared::{
    AppId, ExecutionId, GroupDocsError, GroupId, GroupKvError, RequestId, ScriptId, SdkCallCx,
};
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!("atomic_write: DATABASE_URL unset — skipping");
        return None;
    };
    let pool = PgPoolOptions::new()
        .max_connections(24)
        .connect(&url)
        .await
        .expect("connect");
    sqlx::migrate!("./migrations")
        .run(&pool)
        .await
        .expect("migrate");
    Some(pool)
}

fn cx(app_id: Uuid) -> SdkCallCx {
    SdkCallCx {
        app_id: AppId::from(app_id),
        script_id: ScriptId::new(),
        principal: None,
        execution_id: ExecutionId::new(),
        request_id: RequestId::new(),
        trigger_depth: 0,
        root_execution_id: ExecutionId::new(),
        is_dead_letter_handler: false,
        event: None,
    }
}

/// An app with one script and one enabled KV trigger watching `collection`.
struct Fixture {
    app: Uuid,
}

async fn setup(pool: &PgPool, collection: &str) -> Fixture {
    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!("aw-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!("aw-app-{uniq}"))
            .bind(group)
            .fetch_one(pool)
            .await
            .expect("app");
    let (admin,): (Uuid,) = sqlx::query_as(
        "INSERT INTO admin_users (username, password_hash) VALUES ($1, 'x') RETURNING id",
    )
    .bind(format!("aw-admin-{uniq}"))
    .fetch_one(pool)
    .await
    .expect("admin");
    let (script,): (Uuid,) = sqlx::query_as(
        "INSERT INTO scripts (app_id, name, source) VALUES ($1, 'handler', '') RETURNING id",
    )
    .bind(app)
    .fetch_one(pool)
    .await
    .expect("script");
    let (trigger,): (Uuid,) = sqlx::query_as(
        "INSERT INTO triggers ( \
            app_id, script_id, kind, retry_max_attempts, retry_backoff, \
            retry_base_ms, registered_by_principal \
         ) VALUES ($1, $2, 'kv', 3, 'exponential', 1000, $3) RETURNING id",
    )
    .bind(app)
    .bind(script)
    .bind(admin)
    .fetch_one(pool)
    .await
    .expect("trigger");
    sqlx::query(
        "INSERT INTO kv_trigger_details (trigger_id, collection_glob, ops) \
         VALUES ($1, $2, '{}')",
    )
    .bind(trigger)
    .bind(collection)
    .execute(pool)
    .await
    .expect("kv_trigger_details");
    Fixture { app }
}

/// Make every outbox insert for `app` fail. Scoped by app_id so parallel tests
/// (which use their own fresh apps) are untouched.
async fn break_outbox_for(pool: &PgPool, app: Uuid, tag: &str) {
    sqlx::query(&format!(
        "CREATE FUNCTION break_outbox_{tag}() RETURNS TRIGGER AS $$ \
         BEGIN \
           IF NEW.app_id = '{app}'::uuid THEN \
             RAISE EXCEPTION 'injected outbox failure'; \
           END IF; \
           RETURN NEW; \
         END; $$ LANGUAGE plpgsql"
    ))
    .execute(pool)
    .await
    .expect("create fn");
    sqlx::query(&format!(
        "CREATE TRIGGER break_outbox_{tag} BEFORE INSERT ON outbox \
         FOR EACH ROW EXECUTE FUNCTION break_outbox_{tag}()"
    ))
    .execute(pool)
    .await
    .expect("create trigger");
}

async fn unbreak_outbox(pool: &PgPool, tag: &str) {
    sqlx::query(&format!(
        "DROP TRIGGER IF EXISTS break_outbox_{tag} ON outbox"
    ))
    .execute(pool)
    .await
    .expect("drop trigger");
    sqlx::query(&format!("DROP FUNCTION IF EXISTS break_outbox_{tag}()"))
        .execute(pool)
        .await
        .expect("drop fn");
}

async fn kv_count(pool: &PgPool, app: Uuid) -> i64 {
    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM kv_entries WHERE app_id = $1")
        .bind(app)
        .fetch_one(pool)
        .await
        .expect("count kv");
    n
}

async fn outbox_count(pool: &PgPool, app: Uuid) -> i64 {
    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM outbox WHERE app_id = $1")
        .bind(app)
        .fetch_one(pool)
        .await
        .expect("count outbox");
    n
}

/// `scripts` is RESTRICT on app delete (code is not data), so drop it first.
async fn cleanup(pool: &PgPool, app: Uuid) {
    sqlx::query("DELETE FROM scripts WHERE app_id = $1")
        .bind(app)
        .execute(pool)
        .await
        .expect("cleanup scripts");
    sqlx::query("DELETE FROM apps WHERE id = $1")
        .bind(app)
        .execute(pool)
        .await
        .expect("cleanup app");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_failed_fan_out_rolls_the_kv_write_back() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let f = setup(&pool, "watched").await;
    let writer = PostgresKvWriter::new(pool.clone());
    let cx = cx(f.app);
    // A distinct, valid SQL identifier for this run's trigger/function names.
    let tag = Uuid::new_v4().simple().to_string();

    // --- Control: with the outbox healthy, the write lands AND fans out. -----
    writer
        .set(&cx, "watched", "k", serde_json::json!(1))
        .await
        .expect("healthy set succeeds");
    assert_eq!(kv_count(&pool, f.app).await, 1, "the row is written");
    assert_eq!(
        outbox_count(&pool, f.app).await,
        1,
        "the matching trigger fanned out"
    );

    // --- The fix: a fan-out failure must roll the write back. ---------------
    break_outbox_for(&pool, f.app, &tag).await;
    let err = writer
        .set(&cx, "watched", "k2", serde_json::json!(2))
        .await
        .expect_err("an outbox failure must surface as an error, not be swallowed");
    unbreak_outbox(&pool, &tag).await;

    assert!(
        format!("{err}").contains("event emit"),
        "the error should name the emit failure, got: {err}"
    );
    assert_eq!(
        kv_count(&pool, f.app).await,
        1,
        "k2 must NOT be in the store — the write is rolled back with its fan-out. \
         A committed row whose trigger never fired is exactly the durability hole \
         the transactional outbox closes."
    );
    assert_eq!(
        outbox_count(&pool, f.app).await,
        1,
        "and no outbox row for the rolled-back write"
    );

    // --- An unwatched collection still writes (no trigger → no fan-out). ----
    writer
        .set(&cx, "unwatched", "k", serde_json::json!(3))
        .await
        .expect("a write with no matching trigger needs no outbox row");
    assert_eq!(kv_count(&pool, f.app).await, 2);
    assert_eq!(outbox_count(&pool, f.app).await, 1, "still no new fan-out");

    cleanup(&pool, f.app).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_failed_fan_out_rolls_a_delete_back() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let f = setup(&pool, "watched").await;
    let writer = PostgresKvWriter::new(pool.clone());
    let cx = cx(f.app);
    let tag = Uuid::new_v4().simple().to_string();

    writer
        .set(&cx, "watched", "doomed", serde_json::json!(1))
        .await
        .expect("seed");
    assert_eq!(kv_count(&pool, f.app).await, 1);

    // A delete that can't fan out must leave the row in place — otherwise the
    // key is gone and nothing downstream ever learns it was removed.
    break_outbox_for(&pool, f.app, &tag).await;
    let err = writer
        .delete(&cx, "watched", "doomed")
        .await
        .expect_err("outbox failure surfaces");
    unbreak_outbox(&pool, &tag).await;
    assert!(format!("{err}").contains("event emit"));
    assert_eq!(
        kv_count(&pool, f.app).await,
        1,
        "the delete rolled back — the key is still there"
    );

    cleanup(&pool, f.app).await;
}

// ----------------------------------------------------------------------------
// Audit fix #8: the per-group quota must be a BOUND, not a hint.
// ----------------------------------------------------------------------------

async fn mk_group(pool: &PgPool) -> Uuid {
    let (g,): (Uuid,) =
        sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id")
            .bind(format!("q-grp-{}", Uuid::new_v4().simple()))
            .fetch_one(pool)
            .await
            .expect("group");
    g
}

async fn group_kv_count(pool: &PgPool, group: Uuid) -> i64 {
    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_kv_entries WHERE group_id = $1")
        .bind(group)
        .fetch_one(pool)
        .await
        .expect("count");
    n
}

/// The service read the group's usage on one pooled connection and wrote on
/// another, so N concurrent writers each saw the same pre-write count, each
/// decided it had room, and together blew past the ceiling. The writer now holds
/// a per-group advisory lock across the check AND the write, so they serialize
/// and each sees its predecessor's row.
///
/// Note this is NOT fixed by merely putting the check in the same transaction:
/// under READ COMMITTED each transaction's `COUNT(*)` still sees a snapshot
/// without the others' uncommitted rows. The lock is what does the work.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_writers_cannot_push_a_group_past_its_row_quota() {
    const MAX_ROWS: u64 = 5;

    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let group = mk_group(&pool).await;
    let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone()));
    let quota = GroupWriteQuota {
        max_rows: MAX_ROWS,
        max_total_bytes: u64::MAX,
        max_value_bytes: 256 * 1024,
    };

    // 20 writers race to insert 20 DISTINCT new keys into an empty group whose
    // ceiling is 5. Exactly 5 may win.
    let app = Uuid::new_v4(); // cx.app_id is not used by the group writer's SQL
    let mut set = tokio::task::JoinSet::new();
    for i in 0..20 {
        let w = Arc::clone(&writer);
        set.spawn(async move {
            let cx = cx(app);
            w.set(
                &cx,
                GroupKvTarget {
                    group_id: GroupId::from(group),
                    collection: "shared",
                    key: &format!("k{i}"),
                },
                serde_json::json!(i),
                quota,
            )
            .await
        });
    }
    let mut ok = 0;
    let mut denied = 0;
    while let Some(r) = set.join_next().await {
        match r.expect("task") {
            Ok(()) => ok += 1,
            Err(GroupKvError::QuotaExceeded { .. }) => denied += 1,
            Err(e) => panic!("unexpected error: {e}"),
        }
    }

    assert_eq!(
        group_kv_count(&pool, group).await,
        i64::try_from(MAX_ROWS).unwrap(),
        "the group must hold EXACTLY its ceiling — a check-then-write race would overshoot"
    );
    let max_rows = usize::try_from(MAX_ROWS).unwrap();
    assert_eq!(ok, max_rows, "exactly {MAX_ROWS} writers may win");
    assert_eq!(denied, 20 - max_rows, "the rest are refused");

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

/// The byte ceiling has the same race, and the same fix.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_writers_cannot_push_a_group_past_its_byte_quota() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let group = mk_group(&pool).await;
    let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone()));

    // A ~100-byte value, and a ceiling that fits about 4 of them. `max_value_bytes`
    // is set low enough that the upper-bound fast path can't short-circuit the
    // scan (rows_after * max_value_bytes must exceed the ceiling immediately).
    let value = serde_json::json!("x".repeat(100));
    let quota = GroupWriteQuota {
        max_total_bytes: 450,
        max_rows: u64::MAX,
        max_value_bytes: 200,
    };

    let app = Uuid::new_v4();
    let mut set = tokio::task::JoinSet::new();
    for i in 0..16 {
        let w = Arc::clone(&writer);
        let v = value.clone();
        set.spawn(async move {
            let cx = cx(app);
            w.set(
                &cx,
                GroupKvTarget {
                    group_id: GroupId::from(group),
                    collection: "shared",
                    key: &format!("k{i}"),
                },
                v,
                quota,
            )
            .await
        });
    }
    while let Some(r) = set.join_next().await {
        match r.expect("task") {
            Ok(()) | Err(GroupKvError::TotalBytesQuotaExceeded { .. }) => {}
            Err(e) => panic!("unexpected error: {e}"),
        }
    }

    let (bytes,): (i64,) = sqlx::query_as(
        "SELECT COALESCE(SUM(octet_length(value::text)), 0)::BIGINT \
         FROM group_kv_entries WHERE group_id = $1",
    )
    .bind(group)
    .fetch_one(&pool)
    .await
    .expect("bytes");
    assert!(
        bytes <= 450,
        "stored bytes ({bytes}) must not exceed the 450-byte ceiling — \
         concurrent writers must not each see the same pre-write total"
    );
    assert!(bytes > 0, "some writers should have succeeded");

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

/// Group DOCS carried the identical check-then-write race as group KV, and the
/// identical fix (a per-group advisory lock, on its own `docs`-namespaced key so
/// it doesn't contend with KV writes).
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_writers_cannot_push_a_group_past_its_docs_row_quota() {
    const MAX_DOCS: u64 = 4;

    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let group = mk_group(&pool).await;
    let writer = Arc::new(PostgresGroupDocsWriter::new(pool.clone()));
    let quota = GroupWriteQuota {
        max_rows: MAX_DOCS,
        max_total_bytes: u64::MAX,
        max_value_bytes: 256 * 1024,
    };

    let app = Uuid::new_v4();
    let mut set = tokio::task::JoinSet::new();
    for i in 0..16 {
        let w = Arc::clone(&writer);
        set.spawn(async move {
            let cx = cx(app);
            w.create(
                &cx,
                GroupDocsTarget {
                    group_id: GroupId::from(group),
                    collection: "articles",
                },
                serde_json::json!({ "n": i }),
                quota,
            )
            .await
        });
    }
    let mut ok = 0;
    while let Some(r) = set.join_next().await {
        match r.expect("task") {
            Ok(_) => ok += 1,
            Err(GroupDocsError::QuotaExceeded { .. }) => {}
            Err(e) => panic!("unexpected error: {e}"),
        }
    }

    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_docs WHERE group_id = $1")
        .bind(group)
        .fetch_one(&pool)
        .await
        .expect("count");
    assert_eq!(
        n,
        i64::try_from(MAX_DOCS).unwrap(),
        "the group must hold EXACTLY its doc ceiling"
    );
    assert_eq!(ok, usize::try_from(MAX_DOCS).unwrap());

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