PiCloud/crates/picloud/tests/common/mod.rs

//! A hermetic Postgres database per test.
//!
//! ## Why this exists
//!
//! Every test in `dispatcher_e2e` / `queue_e2e` / `invoke_e2e` / `retry_e2e` /
//! `email_inbound` calls `build_app`, which spawns a **real dispatcher**. When
//! those tests shared one database, they shared one `outbox` — and
//! `OutboxRepo::claim_due` is deliberately NOT app-scoped (in production one
//! dispatcher serves the whole instance, so claiming any due row is correct).
//!
//! So test A's dispatcher would claim test B's outbox row, test A would finish,
//! its server would be dropped mid-dispatch, and the claim would be **stranded**.
//! Test B then polled for its handler's side effect until it timed out. The
//! result was a flake that moved around between runs — `dispatcher_e2e` failed
//! roughly one run in three, on a different test each time, and never when run
//! with `--test-threads=1`.
//!
//! Note what the harness was modelling: N independent *instances* sharing one
//! database. Production never does that — one instance owns one database. The
//! fix is therefore to make the tests match production rather than to weaken the
//! dispatcher (app-scoping `claim_due` would be wrong) or to shorten
//! `PICLOUD_OUTBOX_CLAIM_TIMEOUT_SEC` (600s is right for a data plane where a
//! script may legitimately run 300s; a 10s test window is not a reason to make
//! production abandon live claims).
//!
//! ## How
//!
//! Migrations are slow to replay 30+ times, so the first caller builds a
//! migrated **template** database once and every test then clones it with
//! `CREATE DATABASE ... TEMPLATE`, which is a file copy. Database names are
//! derived from (suite, test name), so they are stable: a test drops and
//! recreates its own database on entry. That means a crashed test leaves no
//! garbage to clean up — the next run reclaims the name — and there is no
//! cleanup step to forget.
//!
//! Skips cleanly (returns `None`) when `DATABASE_URL` is unset, matching the
//! existing `pool_or_skip` contract.

#![allow(dead_code)]

use sqlx::postgres::PgPoolOptions;
use sqlx::{Connection, Executor, PgConnection, PgPool};

/// The migrated database every test database is cloned from.
const TEMPLATE_DB: &str = "picloud_test_template";

/// Serializes template creation and `CREATE DATABASE ... TEMPLATE` across every
/// concurrently-running test *process*. Advisory locks are cluster-wide, so this
/// works across test binaries, not just across threads. It is needed because
/// Postgres refuses to clone a template that has an open connection.
const TEMPLATE_LOCK: i64 = 0x0070_6963_6C6F_7564;

/// Split `postgres://…/picloud?sslmode=…` into (`everything before the db name`,
/// `db name`, `query suffix`) so we can address a sibling database on the same
/// server.
fn split_url(url: &str) -> Option<(String, String)> {
    let (root, tail) = url.rsplit_once('/')?;
    let (_db, query) = match tail.split_once('?') {
        Some((db, q)) => (db, format!("?{q}")),
        None => (tail, String::new()),
    };
    Some((root.to_string(), query))
}

/// Postgres identifiers cap at 63 bytes, so a long test name is truncated and
/// disambiguated by a hash of the full (suite, test) pair.
fn db_name(suite: &str, test: &str) -> String {
    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
    for b in suite
        .bytes()
        .chain(b"::".iter().copied())
        .chain(test.bytes())
    {
        hash ^= u64::from(b);
        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
    }
    let safe: String = test
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() {
                c.to_ascii_lowercase()
            } else {
                '_'
            }
        })
        .take(38)
        .collect();
    format!("pt_{safe}_{:08x}", hash & 0xffff_ffff)
}

/// The running test's name. libtest names each test's thread after the test, and
/// `#[tokio::test]` drives the future on that same thread (`block_on` polls on the
/// caller), so this is stable for both the current-thread and multi-thread
/// flavors. Falls back to a constant if a future harness ever stops naming them —
/// tests would then share a database again, which is exactly today's behavior, so
/// the fallback cannot be worse than the status quo.
fn test_name() -> String {
    std::thread::current()
        .name()
        .unwrap_or("unnamed")
        .rsplit("::")
        .next()
        .unwrap_or("unnamed")
        .to_string()
}

/// Connect to a database private to this test, or `None` when `DATABASE_URL` is
/// unset. `suite` disambiguates same-named tests in different binaries.
pub async fn test_pool(suite: &str) -> Option<PgPool> {
    let Ok(url) = std::env::var("DATABASE_URL") else {
        eprintln!("{suite}: DATABASE_URL unset — skipping");
        return None;
    };
    let (root, query) = split_url(&url).expect("DATABASE_URL must name a database");
    let name = db_name(suite, &test_name());

    let mut admin = PgConnection::connect(&url)
        .await
        .expect("connect to DATABASE_URL");

    // Hold the lock across BOTH the template build and the clone: Postgres will
    // not clone a template that has an open connection, and the migrating
    // connection below is exactly such a connection.
    sqlx::query("SELECT pg_advisory_lock($1)")
        .bind(TEMPLATE_LOCK)
        .execute(&mut admin)
        .await
        .expect("take the template lock");

    ensure_template(&mut admin, &root, &query).await;

    // A stable name means a crashed test leaves nothing behind: the next run
    // simply reclaims it. FORCE evicts any connection a previous run leaked.
    admin
        .execute(format!(r#"DROP DATABASE IF EXISTS "{name}" WITH (FORCE)"#).as_str())
        .await
        .expect("drop the previous test database");
    admin
        .execute(format!(r#"CREATE DATABASE "{name}" TEMPLATE "{TEMPLATE_DB}""#).as_str())
        .await
        .expect("clone the template");

    sqlx::query("SELECT pg_advisory_unlock($1)")
        .bind(TEMPLATE_LOCK)
        .execute(&mut admin)
        .await
        .expect("release the template lock");
    admin.close().await.ok();

    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect(&format!("{root}/{name}{query}"))
        .await
        .expect("connect to the test database");
    Some(pool)
}

/// Create the template if absent, then migrate it. Migrating on every call (not
/// just on creation) is what picks up a newly-added migration; it is a handful of
/// queries once the template is current. The caller must hold [`TEMPLATE_LOCK`],
/// and this must leave no connection open to the template.
async fn ensure_template(admin: &mut PgConnection, root: &str, query: &str) {
    let exists: Option<(i32,)> = sqlx::query_as("SELECT 1 FROM pg_database WHERE datname = $1")
        .bind(TEMPLATE_DB)
        .fetch_optional(&mut *admin)
        .await
        .expect("look up the template");
    if exists.is_none() {
        admin
            .execute(format!(r#"CREATE DATABASE "{TEMPLATE_DB}""#).as_str())
            .await
            .expect("create the template");
    }

    let mut tmpl = PgConnection::connect(&format!("{root}/{TEMPLATE_DB}{query}"))
        .await
        .expect("connect to the template");
    sqlx::migrate!("../manager-core/migrations")
        .run(&mut tmpl)
        .await
        .expect("migrate the template");
    // Must be closed before the clone — see the lock comment above.
    tmpl.close().await.expect("close the template connection");
}