server/

fleet.rs

//! Bringing every tenant database to the current schema.
//!
//! `boot()` migrates only the global database, and provisioning migrates a
//! tenant once, at creation. Nothing else ever touched an existing tenant — so
//! before this module the first tenant-facing migration after launch reached
//! the global database and newly registered users, and no existing user, ever.
//! CI stayed green (fresh databases), the deploy stayed green (global
//! migrates), and each pre-existing user broke at runtime, one feature at a
//! time.
//!
//! This runs as its own process, never inside a serving one. `boot()` is
//! awaited before either binary serves and a boot error exits, so a fan-out
//! there would let one unreachable tenant crash-loop web and sshd for every
//! user.
//!
//! Reconciliation is against each tenant's own `_sqlx_migrations`, via the
//! ordinary migrator. There is deliberately no central version table: applied
//! state is already recorded per tenant, a mirror can drift, and under the
//! shared DDL-only rule its `CREATE TABLE` would land in every tenant database
//! anyway.

use crate::db::{DBError, migrate_on};
use crate::provision::dsn_for;
use sqlx::types::Uuid;
use sqlx::types::chrono::{DateTime, Utc};
use sqlx::{Connection, PgConnection};
use std::time::Duration;

#[cfg(test)]
mod tests;

// Tenants are migrated ONE AT A TIME, deliberately.
//
// Every concurrent worker costs a tenant connection on top of the admin pool,
// and ADR 0032 reserves five connections for this whole process; four workers
// spent about thirteen. Sequential also keeps the migration on the very
// connection that holds the claim, which a spawned worker cannot do — sqlx's
// migrator future is not `Send` while it borrows a connection.
//
// The cost is wall-clock on a large fleet. Revisit by giving the migrator its
// own capped pool and a `Send`-safe worker, not by adding threads here.

/// Ceiling on one tenant's migration. Without it a single tenant holding a lock
/// stalls the whole run.
const TENANT_DEADLINE: Duration = Duration::from_secs(120);

/// Bound on reaching a tenant at all, so connect cannot eat the outer budget.
const CONNECT_DEADLINE: Duration = Duration::from_secs(10);

/// Backstop around connect, claim and migration together. Looser than
/// [`TENANT_DEADLINE`] so the inner deadline normally fires first and gets to
/// close the connection.
const VISIT_DEADLINE: Duration = Duration::from_secs(200);

/// Connection attempts before a tenant is quarantined.
const CONNECT_ATTEMPTS: u32 = 2;

/// Pause between those attempts.
const CONNECT_RETRY_DELAY: Duration = Duration::from_secs(3);

/// Advisory-lock key, so two migrator processes cannot work one tenant at once.
///
/// Taken inside the *tenant's own* database rather than a shared table: there
/// is no central state to drift, and a lock dies with the connection holding
/// it, so a killed worker frees its claim without a lease to expire.
const CLAIM_LOCK: i64 = 0x6e6f_6d69_0001;

/// How often to re-attempt a claim another worker currently holds.
const CLAIM_POLL: Duration = Duration::from_secs(2);

/// How many times the directory is re-read looking for tenants that appeared
/// mid-run. Bounded so a steady stream of registrations cannot keep the Job
/// alive forever; anything still missed is caught by the next release's run.
const INVENTORY_ROUNDS: usize = 3;

/// How long to wait out a peer's claim before giving up on this tenant.
/// Bounded well inside [`TENANT_DEADLINE`] so the migration itself still has
/// room once the claim is won.
const CLAIM_DEADLINE: Duration = Duration::from_secs(45);

/// What became of one tenant.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Outcome {
    /// Verified at the current schema by *this* run.
    Migrated,
    /// Another worker held the claim for the whole deadline, so this run never
    /// saw the tenant's schema. Not a success: the holder may since have timed
    /// out, panicked or been killed, and nothing would say so.
    Unresolved,
    /// Failed, and recorded in the global directory.
    Quarantined(String),
}

#[derive(Debug, Default)]
pub struct FleetReport {
    pub migrated: usize,
    /// Tenants this run could not verify because another worker held the claim.
    pub unresolved: Vec<Uuid>,
    pub quarantined: Vec<(Uuid, String)>,
    /// Tenants whose verdict could not be written to the directory. The
    /// dangerous case: a tenant that failed but whose quarantine did not stick
    /// is still being served.
    pub unpublished: Vec<(Uuid, String)>,
}

impl FleetReport {
    /// Whether every tenant was accounted for.
    ///
    /// A quarantined tenant does not stop the *fleet* — that is the whole point
    /// of quarantine — but it does mean this run left work undone, and the exit
    /// code is what a deployment gate reads. Unresolved and unpublished tenants
    /// count too: in both cases the run finished without knowing, or without
    /// recording, the state of a customer's database.
    #[must_use]
    pub const fn is_healthy(&self) -> bool {
        self.quarantined.is_empty() && self.unresolved.is_empty() && self.unpublished.is_empty()
    }

    #[must_use]
    pub fn summary(&self) -> String {
        format!(
            "migrated {}, quarantined {}, unresolved {}, unpublished {}",
            self.migrated,
            self.quarantined.len(),
            self.unresolved.len(),
            self.unpublished.len()
        )
    }
}

/// Migrates one tenant, holding a claim for the duration.
///
/// Waits for a claim another worker holds rather than declaring victory over a
/// tenant it never looked at. Once the claim is won the migrator runs
/// regardless of what the previous holder did: it is idempotent, so running it
/// after a peer's success costs a no-op, and running it after a peer *died*
/// mid-way is exactly the repair that is needed.
async fn migrate_tenant(database_name: &str) -> Result<Outcome, DBError> {
    let dsn = dsn_for(database_name)?;
    let mut conn = connect(&dsn).await?;

    if !claim(&mut conn, CLAIM_DEADLINE).await? {
        let _ = conn.close().await;
        return Ok(Outcome::Unresolved);
    }

    // Migrated on the claim connection itself. A second connection per worker
    // would put the run at roughly thirteen against the five ADR 0032 reserves
    // for it — a budget breach during exactly the operation that must not fail.
    //
    // The deadline is applied HERE rather than around the whole call: cancelling
    // this future from outside would drop the connection without closing it.
    let migrated = tokio::time::timeout(TENANT_DEADLINE, migrate_on(&mut conn)).await;

    // Closing the connection drops the lock; do it whether or not we succeeded,
    // and whether or not we ran out of time.
    let _ = conn.close().await;

    match migrated {
        Ok(result) => result.map(|()| Outcome::Migrated),
        Err(_) => Ok(Outcome::Quarantined(format!(
            "migration exceeded {}s",
            TENANT_DEADLINE.as_secs()
        ))),
    }
}

/// Connects to a tenant, retrying a refused or slow connection once.
///
/// The only retry in a visit, and deliberately before any claim is taken: a
/// migration that failed may still hold its advisory lock, so retrying one
/// would find its own claim and learn nothing.
async fn connect(dsn: &str) -> Result<PgConnection, DBError> {
    let mut attempt = 1;
    loop {
        match tokio::time::timeout(CONNECT_DEADLINE, PgConnection::connect(dsn)).await {
            Ok(Ok(conn)) => return Ok(conn),
            _ if attempt < CONNECT_ATTEMPTS => {
                tokio::time::sleep(CONNECT_RETRY_DELAY).await;
                attempt += 1;
            }
            Ok(Err(e)) => return Err(DBError::Sqlx(e)),
            Err(_) => return Err(DBError::Config(crate::config::ConfigError::DB)),
        }
    }
}

/// Takes the tenant's claim, waiting out a peer that already holds it.
///
/// Polls rather than using the blocking `pg_advisory_lock`, so the surrounding
/// timeout can actually interrupt it — a blocked lock wait inside Postgres
/// would ignore a cancelled future and hold the connection open.
async fn claim(conn: &mut PgConnection, deadline: Duration) -> Result<bool, DBError> {
    let give_up_at = tokio::time::Instant::now() + deadline;

    loop {
        let taken: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
            .bind(CLAIM_LOCK)
            .fetch_one(&mut *conn)
            .await
            .map_err(DBError::Sqlx)?;

        if taken {
            return Ok(true);
        }

        if tokio::time::Instant::now() >= give_up_at {
            return Ok(false);
        }

        tokio::time::sleep(CLAIM_POLL).await;
    }
}

/// Runs one tenant to a verdict, converting any failure into a quarantine
/// reason rather than propagating it — one sick tenant must not end the run.
///
/// A timeout quarantines rather than passing: the migration was cancelled
/// mid-flight, so its transaction rolled back and the tenant is on whatever
/// schema it had before. Unknown is not the same as fine.
async fn visit(database_name: &str) -> Outcome {
    // The backstop around connect, claim and migration together; looser than
    // `TENANT_DEADLINE` so the inner deadline normally fires first.
    match tokio::time::timeout(VISIT_DEADLINE, migrate_tenant(database_name)).await {
        Ok(Ok(outcome)) => outcome,
        Ok(Err(e)) => Outcome::Quarantined(e.to_string()),
        Err(_) => Outcome::Quarantined(format!("tenant exceeded {}s", VISIT_DEADLINE.as_secs())),
    }
}

/// Migrates every tenant in the directory.
///
/// Skips `admin_user_id`, whose row points at the global database that the
/// ordinary boot migration already handled. Quarantined tenants are retried:
/// quarantine records a failure, it is not a tombstone.
///
/// Never seeds. `server/seed.sql` must not run against a per-user database —
/// an implementer reusing boot-shaped logic here would insert the admin user
/// and global config into every tenant.
///
/// # Errors
/// Only for failures that prevent the run itself: an unreachable *global*
/// directory, or a missing `DATABASE_URL`. A tenant that cannot be migrated is
/// quarantined and reported, not raised.
/// How a database's applied migrations differ from this binary's.
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct SchemaDrift {
    /// Embedded here, not applied there: the directory is behind this binary.
    pub(crate) missing: Vec<i64>,
    /// Applied there, unknown here: the directory is AHEAD, so this binary is
    /// the stale one — a rolled-back image, or one whose migration files were
    /// edited in place after being applied.
    pub(crate) unknown: Vec<i64>,
}

impl SchemaDrift {
    fn is_aligned(&self) -> bool {
        self.missing.is_empty() && self.unknown.is_empty()
    }
}

/// Compares two migration sets by version, in BOTH directions.
///
/// Versions rather than counts, because the set has a deliberate gap
/// (1, 2, 3, 5, 6) and "five applied" is not "the right five applied".
///
/// Both directions, because a one-way check is only half a gate. A binary older
/// than the directory passes `embedded ⊆ applied` and then fails on every single
/// tenant — sqlx refuses a version it does not know — turning a deployment
/// mistake into a fleet-wide quarantine whose verdicts outrank the correct
/// run's, because they were written later.
fn schema_drift(embedded: &[i64], applied: &[i64]) -> SchemaDrift {
    SchemaDrift {
        missing: embedded
            .iter()
            .filter(|version| !applied.contains(version))
            .copied()
            .collect(),
        unknown: applied
            .iter()
            .filter(|version| !embedded.contains(version))
            .copied()
            .collect(),
    }
}

/// Refuses to run unless the GLOBAL database already carries this binary's
/// migration set.
///
/// The Job and the app are applied by different Flux Kustomizations, and
/// `dependsOn` waits for the app to be *Ready* — which it can be on a previous
/// revision. So the Job can start before the release it belongs to has migrated
/// the global directory, read a `users` table shaped by the older schema, and
/// carry that misunderstanding into every tenant.
///
/// Checking the precondition directly is better than trusting the ordering:
/// it holds however the Job came to be scheduled.
async fn assert_global_schema_current() -> Result<(), DBError> {
    let mut conn = crate::db::get_connection().await?;

    let applied: Vec<i64> =
        sqlx::query_scalar!(r#"SELECT version as "version!" FROM _sqlx_migrations WHERE success"#)
            .fetch_all(&mut *conn)
            .await
            .map_err(DBError::Sqlx)?;

    let embedded: Vec<i64> = sqlx::migrate!("../migrations")
        .iter()
        .map(|m| m.version)
        .collect();

    let drift = schema_drift(&embedded, &applied);
    if drift.is_aligned() {
        return Ok(());
    }

    Err(DBError::GlobalSchemaStale(format!(
        "this binary and the global directory disagree: missing {:?}, unknown {:?}. \
         Missing means the directory is behind this release; unknown means this \
         binary is behind the directory. Either way, migrating tenants from here \
         would quarantine them on a disagreement that is not theirs",
        drift.missing, drift.unknown
    )))
}

pub async fn migrate_fleet(admin_user_id: Uuid) -> Result<FleetReport, DBError> {
    // Before anything: never fan out from a directory older than this binary.
    assert_global_schema_current().await?;

    // The database's clock, not this host's. The fence compares against
    // `quarantined_at`, written by `now()` inside Postgres; comparing that to a
    // migrator host's clock makes correctness a function of clock skew between
    // two machines.
    let started = database_now().await?;
    let mut report = FleetReport::default();
    let mut seen: Vec<Uuid> = Vec::new();

    // Re-inventoried until it stops growing. A single snapshot would miss a
    // tenant registered while the run was in progress — by an old pod still
    // serving, whose provisioning carries *its* migration set — and the run
    // would still exit 0 over a tenant nobody brought current.
    //
    // This NARROWS the window; it does not close it. A tenant registered after
    // the last read is still missed, and the run still exits 0. Closing it needs
    // a barrier between provisioning and this run — a shared lock, or a schema
    // generation that provisioning stamps and this reconciles against. Until
    // then the guarantee rests where ADR 0032 puts it: expand/contract, so the
    // serving code works on both schemas across a release.
    for round in 1..=INVENTORY_ROUNDS {
        let pending: Vec<(Uuid, String)> = inventory(admin_user_id)
            .await?
            .into_iter()
            .filter(|(id, _)| !seen.contains(id))
            .collect();

        if pending.is_empty() {
            break;
        }

        log::info!(
            "fleet migration round {round}: {} tenant(s) to visit",
            pending.len()
        );

        for (id, database_name) in pending {
            seen.push(id);
            let outcome = visit(&database_name).await;

            // Published as each tenant lands. Collecting verdicts first would
            // leave a tenant that failed early still being served for as long
            // as the rest of the fleet takes.
            if let Err(e) = publish(id, &outcome, started).await {
                log::error!("could not record the verdict for {database_name}: {e}");
                report.unpublished.push((id, e.to_string()));
            }

            match outcome {
                Outcome::Migrated => report.migrated += 1,
                Outcome::Unresolved => {
                    log::warn!(
                        "{database_name} stayed claimed by another worker; this run did not verify it"
                    );
                    report.unresolved.push(id);
                }
                Outcome::Quarantined(reason) => {
                    log::error!("quarantining {database_name}: {reason}");
                    report.quarantined.push((id, reason));
                }
            }
        }
    }

    log::info!("fleet migration finished: {}", report.summary());
    Ok(report)
}

/// Writes one tenant's verdict to the global directory.
///
/// Errors are returned, never swallowed: a quarantine that fails to stick
/// leaves a tenant on the wrong schema still being served, which is the exact
/// state this module exists to prevent.
async fn publish(
    user_id: Uuid,
    outcome: &Outcome,
    run_started: DateTime<Utc>,
) -> Result<(), DBError> {
    match outcome {
        Outcome::Migrated => clear_quarantine(user_id, run_started).await,
        Outcome::Quarantined(reason) => quarantine(user_id, reason, run_started).await,
        // Nothing was learned about this tenant, so nothing is claimed about it.
        Outcome::Unresolved => Ok(()),
    }
}

/// Reads the wall clock of the database that records quarantines.
///
async fn database_now() -> Result<DateTime<Utc>, DBError> {
    let mut conn = crate::db::get_connection().await?;

    sqlx::query_scalar!(r#"SELECT now() as "now!""#)
        .fetch_one(&mut *conn)
        .await
        .map_err(DBError::Sqlx)
}

async fn inventory(admin_user_id: Uuid) -> Result<Vec<(Uuid, String)>, DBError> {
    let mut conn = crate::db::get_connection().await?;

    let rows = sqlx::query_file!(
        "../server/sql/select/system/tenant_inventory.sql",
        admin_user_id
    )
    .fetch_all(&mut *conn)
    .await
    .map_err(DBError::Sqlx)?;

    Ok(rows
        .into_iter()
        .map(|row| (row.id, row.database_name))
        .collect())
}

async fn quarantine(
    user_id: Uuid,
    reason: &str,
    run_started: DateTime<Utc>,
) -> Result<(), DBError> {
    let mut conn = crate::db::get_connection().await?;

    let recorded = sqlx::query_file!(
        "../server/sql/update/users/quarantine.sql",
        user_id,
        reason,
        run_started
    )
    .execute(&mut *conn)
    .await
    .map_err(DBError::Sqlx)?;

    if recorded.rows_affected() == 0 {
        log::info!("{user_id} already carries a verdict newer than this run");
    }

    Ok(())
}

/// Lifts a quarantine, but only one that predates this run.
///
/// The fence matters when two releases' migrators overlap: without it, an older
/// migrator finishing successfully against its own (older) migration set would
/// clear a quarantine a newer migrator had just written for failing on a newer
/// one — and the newest application would then serve a tenant known not to have
/// the schema it needs.
async fn clear_quarantine(user_id: Uuid, run_started: DateTime<Utc>) -> Result<(), DBError> {
    let mut conn = crate::db::get_connection().await?;

    let cleared = sqlx::query_file!(
        "../server/sql/update/users/clear_quarantine.sql",
        user_id,
        run_started
    )
    .execute(&mut *conn)
    .await
    .map_err(DBError::Sqlx)?;

    // A tenant that was never quarantined still matches the fence, so this is
    // not the ordinary case — it means someone newer has already ruled.
    if cleared.rows_affected() == 0 {
        log::info!("{user_id} was not cleared: its quarantine is newer than this run");
    }

    Ok(())
}