pgsync_postgres/

job.rs

//! Job orchestration: bootstrap (snapshot + handoff), run, status, cleanup, and the WAL
//! safety monitor.

use crate::client::{self, PgClient};
use crate::errors::from_pg;
use crate::replication::conn::ReplicationConnection;
use crate::replication::slot::{create_slot_with_snapshot, identify_system};
use crate::replication::stream::PgStreamFactory;
use crate::schema;
use crate::sink::{self, PgSink, SinkLimits};
use crate::snapshot;
use crate::source::{self, PreflightMode, PreflightReport, SourcePreflight};
use crate::url::ConnectionInfo;
use pgsync_core::checkpoint::Checkpoint;
use pgsync_core::config::{Config, SchemaMode, TableName, WalRetentionAction};
use pgsync_core::engine::run_replication;
use pgsync_core::metrics::Metrics;
use pgsync_core::state::JobState;
use pgsync_core::traits::SinkFactory;
use pgsync_core::{Error, ErrorKind, Lsn, Result};
use postgres_types::PgLsn;
use serde::Serialize;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};

/// Everything a command needs.
pub struct JobContext {
    pub config: Config,
    pub source: ConnectionInfo,
    pub target: ConnectionInfo,
    pub metrics: Arc<Metrics>,
    pub cancel: CancellationToken,
    pub progress: pgsync_core::progress::ProgressHandle,
    /// Source `version()` once known (preflight or stream open), for the heartbeat row.
    pub source_version: Arc<Mutex<Option<String>>>,
}

impl JobContext {
    pub fn new(config: Config, cancel: CancellationToken) -> Result<Self> {
        let source = ConnectionInfo::from_env(&config.source.url_env, &config.safety)?;
        let target = ConnectionInfo::from_env(&config.target.url_env, &config.safety)?;
        Ok(JobContext {
            config,
            source,
            target,
            metrics: Arc::new(Metrics::new()),
            cancel,
            progress: pgsync_core::progress::none(),
            source_version: Arc::new(Mutex::new(None)),
        })
    }

    fn timeout(&self) -> Duration {
        self.config.resources.connection_timeout
    }

    async fn source_client(&self, name: &str) -> Result<PgClient> {
        client::connect(&self.source, &self.config.safety, self.timeout(), name).await
    }

    async fn target_client(&self, name: &str) -> Result<PgClient> {
        client::connect(&self.target, &self.config.safety, self.timeout(), name).await
    }

    fn limits(&self) -> SinkLimits {
        SinkLimits {
            max_batch_rows: self.config.resources.max_batch_rows,
            max_batch_bytes: self.config.resources.max_batch_bytes as usize,
        }
    }

    pub fn stream_factory(&self) -> PgStreamFactory {
        PgStreamFactory {
            info: self.source.clone(),
            safety: self.config.safety.clone(),
            timeout: self.timeout(),
            slot: self.config.source.slot.clone(),
            publication: self.config.source.publication.clone(),
            max_message_bytes: self.config.resources.effective_max_message_bytes() as usize,
        }
    }
}

pub struct PgSinkFactory {
    pub target: ConnectionInfo,
    pub safety: pgsync_core::config::SafetyConfig,
    pub timeout: Duration,
    pub max_batch_rows: usize,
    pub max_batch_bytes: usize,
    pub tables: Vec<TableName>,
}

impl SinkFactory for PgSinkFactory {
    type Sink = PgSink;
    async fn connect(&mut self) -> Result<PgSink> {
        let c = client::connect(&self.target, &self.safety, self.timeout, "pgsync-apply").await?;
        Ok(PgSink::new(
            c,
            SinkLimits {
                max_batch_rows: self.max_batch_rows,
                max_batch_bytes: self.max_batch_bytes,
            },
        )
        .await?
        .restrict_to(&self.tables))
    }
}

// ---- check -------------------------------------------------------------------------

#[derive(Serialize)]
pub struct CheckOutput {
    pub source: PreflightReport,
    pub target: PreflightReport,
    pub ok: bool,
}

pub async fn check(
    ctx: &JobContext,
    mode: PreflightMode,
) -> Result<(CheckOutput, Option<SourcePreflight>)> {
    let mut target_report = PreflightReport::default();
    let src = match ctx.source_client("pgsync-check").await {
        Ok(c) => Some(c),
        Err(e) => {
            let mut r = PreflightReport::default();
            r.fail("source_connection", e.to_string());
            return Ok((
                CheckOutput {
                    source: r,
                    target: target_report,
                    ok: false,
                },
                None,
            ));
        }
    };
    let src = src.unwrap();
    let mut pre = source::preflight(
        &src,
        &ctx.config.source.publication,
        &ctx.config.source.slot,
        &ctx.config.replication.tables,
        mode,
        ctx.config.resources.feedback_interval,
    )
    .await?;
    *ctx.source_version.lock().unwrap() = Some(pre.server.version.clone());
    // Replication protocol connectivity (pg_hba, REPLICATION attribute in practice).
    match ReplicationConnection::connect(&ctx.source, &ctx.config.safety, ctx.timeout()).await {
        Ok(mut conn) => match identify_system(&mut conn).await {
            Ok(sys) => {
                pre.report.pass(
                    "replication_connection",
                    format!("system {} timeline {}", sys.system_id, sys.timeline),
                );
                conn.close().await;
            }
            Err(e) => pre.report.fail("replication_connection", e.to_string()),
        },
        Err(e) => pre.report.fail(
            "replication_connection",
            format!(
                "{e} (check pg_hba.conf allows replication connections for this role/database)"
            ),
        ),
    }
    match session_tls(&src).await {
        Ok(true) => pre
            .report
            .pass("source_tls", format!("encrypted ({:?})", ctx.source.tls)),
        Ok(false) if ctx.source.is_loopback => {
            pre.report.warn("source_tls", "plain connection (loopback)")
        }
        Ok(false) => pre.report.fail("source_tls", "connection is not encrypted"),
        Err(e) => pre
            .report
            .warn("source_tls", format!("cannot determine: {e}")),
    }

    match ctx.target_client("pgsync-check").await {
        Ok(tgt) => {
            let row = tgt
                .query_one("SELECT version(), current_setting('server_version_num')::int, (SELECT rolsuper FROM pg_roles WHERE rolname = current_user)", &[])
                .await
                .map_err(|e| from_pg("target version", e))?;
            let version: String = row.get(0);
            let vnum: i32 = row.get(1);
            let superuser: bool = row.get(2);
            if vnum < 140000 {
                target_report.fail("target_version", version);
            } else {
                target_report.pass("target_version", version);
            }
            if superuser {
                target_report.pass(
                    "target_role",
                    "superuser (session_replication_role = replica available)",
                );
            } else {
                target_report.warn("target_role", "not a superuser: target triggers and FK checks stay active during apply; tables must be ordered parent-first");
            }
            match tgt
                .batch_execute("CREATE SCHEMA IF NOT EXISTS pgsync")
                .await
            {
                Ok(()) => target_report.pass("target_pgsync_schema", "writable"),
                Err(e) => target_report.fail(
                    "target_pgsync_schema",
                    from_pg("create schema", e).to_string(),
                ),
            }
            if ctx.config.schema.mode == SchemaMode::PgDump {
                let bin = ctx
                    .config
                    .schema
                    .pg_dump_path
                    .clone()
                    .unwrap_or_else(|| "pg_dump".into());
                match tokio::process::Command::new(&bin)
                    .arg("--version")
                    .output()
                    .await
                {
                    Ok(o) if o.status.success() => target_report.pass(
                        "pg_dump",
                        String::from_utf8_lossy(&o.stdout).trim().to_owned(),
                    ),
                    _ => target_report.fail(
                        "pg_dump",
                        format!(
                            "{bin} not runnable (set schema.pg_dump_path or schema.mode: existing)"
                        ),
                    ),
                }
            }
            for t in &ctx.config.replication.tables {
                match schema::target_columns(&tgt, &t.schema, &t.name).await? {
                    Some(cols) => {
                        if ctx.config.schema.mode == SchemaMode::PgDump
                            && mode == PreflightMode::Bootstrap
                        {
                            target_report.fail(&format!("target table {t}"), "already exists; schema.mode is pg_dump (use --restart-snapshot to drop and recreate, or schema.mode: existing)");
                        } else if let Some(info) = pre.tables.iter().find(|i| &i.table == t) {
                            match schema::verify_columns(t, &info.columns, &cols) {
                                Ok(w) if w.is_empty() => target_report
                                    .pass(&format!("target table {t}"), "columns match"),
                                Ok(w) => {
                                    target_report.warn(&format!("target table {t}"), w.join("; "))
                                }
                                Err(e) => target_report
                                    .fail(&format!("target table {t}"), e.message().to_owned()),
                            }
                        }
                    }
                    None => {
                        if ctx.config.schema.mode == SchemaMode::Existing {
                            target_report.fail(
                                &format!("target table {t}"),
                                "does not exist and schema.mode is existing",
                            );
                        } else if mode == PreflightMode::Run {
                            target_report
                                .fail(&format!("target table {t}"), "does not exist on the target");
                        } else {
                            target_report.pass(
                                &format!("target table {t}"),
                                "will be created by pg_dump schema bootstrap",
                            );
                        }
                    }
                }
            }
            match session_tls(&tgt).await {
                Ok(true) => {
                    target_report.pass("target_tls", format!("encrypted ({:?})", ctx.target.tls))
                }
                Ok(false) if ctx.target.is_loopback => {
                    target_report.warn("target_tls", "plain connection (loopback)")
                }
                Ok(false) => target_report.fail("target_tls", "connection is not encrypted"),
                Err(e) => target_report.warn("target_tls", format!("cannot determine: {e}")),
            }
        }
        Err(e) => target_report.fail("target_connection", e.to_string()),
    }
    let ok = !pre.report.failed() && !target_report.failed();
    Ok((
        CheckOutput {
            source: pre.report.clone(),
            target: target_report,
            ok,
        },
        Some(pre),
    ))
}

async fn session_tls(c: &PgClient) -> Result<bool> {
    let row = c
        .query_one(
            "SELECT coalesce((SELECT ssl FROM pg_stat_ssl WHERE pid = pg_backend_pid()), false)",
            &[],
        )
        .await
        .map_err(|e| from_pg("pg_stat_ssl", e))?;
    Ok(row.get(0))
}

fn preflight_error(out: &CheckOutput) -> Error {
    let mut lines = Vec::new();
    for c in out.source.failures() {
        lines.push(format!("source {}: {}", c.name, c.detail));
    }
    for c in out.target.failures() {
        lines.push(format!("target {}: {}", c.name, c.detail));
    }
    Error::preflight(format!("preflight failed:\n  {}", lines.join("\n  ")))
}

// ---- bootstrap ---------------------------------------------------------------------

pub async fn bootstrap(ctx: &JobContext, restart_snapshot: bool) -> Result<()> {
    let cfg = &ctx.config;
    let job_id = cfg.job.id.clone();
    ctx.metrics.set_state(JobState::Initializing);

    // Existing job?
    let tgt = ctx.target_client("pgsync-bootstrap").await?;
    tgt.batch_execute(sink::SCHEMA_SQL)
        .await
        .map_err(|e| from_pg("create pgsync schema", e))?;
    let mut probe = PgSink::new(ctx.target_client("pgsync-bootstrap").await?, ctx.limits()).await?;
    let existing =
        pgsync_core::traits::ReplicationSink::load_checkpoint(&mut probe, &job_id).await?;
    drop(probe);
    if let Some(cp) = &existing {
        if cp.state == JobState::Initializing {
            // A previous bootstrap died between recording its intent and recording the
            // snapshot. Nothing was copied; the only thing to undo is a slot it may have
            // created. Target tables are left alone.
            info!("previous bootstrap did not reach the snapshot; discarding its metadata");
            drop_owned_slot(ctx, cp, "left by an interrupted bootstrap").await?;
            sink::delete_job(&tgt, &job_id).await?;
        } else if !restart_snapshot {
            return Err(Error::config(format!(
                "job {job_id:?} already exists on the target in state {}. Use `pgsync run` to resume, \
                 `pgsync bootstrap --restart-snapshot` to discard the incomplete copy and start over, \
                 or `pgsync cleanup` to remove it.",
                cp.state
            )));
        } else {
            info!(state = %cp.state, "restarting snapshot: discarding previous copy");
            restart_cleanup(ctx, &tgt, Some(cp)).await?;
        }
    } else if restart_snapshot {
        info!("restarting snapshot: no job metadata; resetting target tables only");
        restart_cleanup(ctx, &tgt, None).await?;
    }

    // Preflight (never modifies the source).
    let (out, pre) = check(ctx, PreflightMode::Bootstrap).await?;
    if !out.ok {
        return Err(preflight_error(&out));
    }
    let pre = pre.expect("preflight present when ok");
    let src_system_id = pre.server.system_identifier.clone();
    let warns = out
        .source
        .checks
        .iter()
        .chain(out.target.checks.iter())
        .filter(|c| c.status == source::CheckStatus::Warn)
        .count();
    ctx.progress.step(&format!(
        "connected to source {} · system {}",
        short_version(&pre.server.version),
        src_system_id
    ));
    ctx.progress.step(&format!(
        "publication {} · {} tables",
        cfg.source.publication,
        pre.tables.len()
    ));
    ctx.progress.step(&format!(
        "preflight: {} checks passed{}",
        out.source.checks.len() + out.target.checks.len(),
        if warns > 0 {
            format!(" ({warns} warnings)")
        } else {
            String::new()
        }
    ));

    // Record the intent to create the slot *before* creating it, so that a crash in between
    // leaves a job row that marks the slot as owned (and therefore droppable by `cleanup`
    // or the next `bootstrap`) instead of an orphaned slot that retains WAL.
    let mut cp = Checkpoint::new_job(
        &job_id,
        &src_system_id,
        &pre.server.dbname,
        &cfg.source.slot,
        &cfg.source.publication,
    );
    cp.state = JobState::Initializing;
    cp.slot_owned = true;
    sink::upsert_checkpoint(&tgt, &cp).await?;

    // Slot + exported snapshot. The replication connection must stay idle until all copy
    // sessions have imported the snapshot.
    let mut repl = ReplicationConnection::connect(&ctx.source, &cfg.safety, ctx.timeout()).await?;
    let sys = identify_system(&mut repl).await?;
    if sys.system_id != src_system_id {
        return Err(Error::internal(
            "system identifier differs between SQL and replication connections",
        ));
    }
    let created = create_slot_with_snapshot(&mut repl, &cfg.source.slot).await?;
    info!(slot = %created.slot_name, consistent_point = %created.consistent_point, snapshot = %created.snapshot_name, "replication slot created with exported snapshot");
    ctx.progress.step(&format!(
        "replication slot {} created (pgoutput) · snapshot at {}",
        created.slot_name, created.consistent_point
    ));

    cp.state = JobState::Snapshotting;
    cp.consistent_point = Some(created.consistent_point);
    cp.commit_lsn = created.consistent_point;
    cp.committed_lsn = created.consistent_point;
    cp.acked_lsn = created.consistent_point;
    sink::upsert_checkpoint(&tgt, &cp).await?;
    ctx.metrics.set_state(JobState::Snapshotting);

    let hb_cancel = ctx.cancel.child_token();
    let reporter = spawn_reporter(ctx, hb_cancel.clone());
    let snapshot_result = run_snapshot(ctx, &tgt, &pre, &created.snapshot_name).await;
    hb_cancel.cancel();
    let _ = reporter.await;
    // The exported snapshot is no longer needed once the copies have committed (or failed).
    repl.close().await;
    match snapshot_result {
        Ok(()) => {}
        Err(e) => {
            // The copy is all-or-nothing: a slot without a usable snapshot would only retain
            // WAL on the source until someone notices. Drop the slot this invocation created
            // (never one that predates it) and leave the job in SNAPSHOTTING so that
            // `bootstrap --restart-snapshot` is the only way forward.
            let dropped = match drop_owned_slot(ctx, &cp, "its snapshot failed").await {
                Ok(()) => "the slot was dropped (no WAL is retained)".to_owned(),
                Err(de) => {
                    error!(error = %de, slot = %cp.slot, "could not drop the slot; it retains WAL until `bootstrap --restart-snapshot` or `cleanup --yes`");
                    format!(
                        "WARNING: slot {:?} could not be dropped and retains WAL: {de}",
                        cp.slot
                    )
                }
            };
            error!(error = %e, "snapshot failed; job stays in SNAPSHOTTING and needs `bootstrap --restart-snapshot`");
            let _ = sink::set_job_state(
                &tgt,
                &job_id,
                JobState::Snapshotting,
                Some(&format!("{e}; {dropped}")),
            )
            .await;
            return Err(e);
        }
    }

    cp.state = JobState::CatchingUp;
    sink::upsert_checkpoint(&tgt, &cp).await?;
    info!(consistent_point = %created.consistent_point, "snapshot complete; switching to CDC");
    drop(tgt);

    run(ctx).await.map(|_| ())
}

/// Drop the job's slot on the source if the job owns it and it exists. Refuses while a
/// walsender holds it (another process is streaming from it).
async fn drop_owned_slot(ctx: &JobContext, cp: &Checkpoint, why: &str) -> Result<()> {
    if !cp.slot_owned {
        return Ok(());
    }
    let src = ctx.source_client("pgsync-bootstrap").await?;
    let Some(info) = source::slot_info(&src, &cp.slot).await? else {
        return Ok(());
    };
    if info.active {
        return Err(Error::preflight(format!(
            "slot {:?} is active (pid {:?}); stop the consumer first",
            cp.slot, info.active_pid
        )));
    }
    src.execute("SELECT pg_drop_replication_slot($1)", &[&cp.slot])
        .await
        .map_err(|e| from_pg("drop slot", e))?;
    info!(slot = %cp.slot, "dropped replication slot owned by this job ({why})");
    Ok(())
}

async fn restart_cleanup(ctx: &JobContext, tgt: &PgClient, cp: Option<&Checkpoint>) -> Result<()> {
    if let Some(cp) = cp {
        drop_owned_slot(ctx, cp, "restarting the snapshot").await?;
    }
    for t in &ctx.config.replication.tables {
        let sql = match ctx.config.schema.mode {
            SchemaMode::PgDump => format!("DROP TABLE IF EXISTS {} CASCADE", t.quoted()),
            SchemaMode::Existing => format!("TRUNCATE TABLE {} CASCADE", t.quoted()),
        };
        tgt.batch_execute(&sql)
            .await
            .map_err(|e| from_pg(&format!("reset target table {t}"), e))?;
    }
    if let Some(cp) = cp {
        sink::delete_job(tgt, &cp.job_id).await?;
    }
    Ok(())
}

async fn run_snapshot(
    ctx: &JobContext,
    tgt: &PgClient,
    pre: &SourcePreflight,
    snapshot_name: &str,
) -> Result<()> {
    let cfg = &ctx.config;
    let started = Instant::now();

    if cfg.schema.mode == SchemaMode::PgDump {
        let sql = schema::dump_schema(
            cfg.schema.pg_dump_path.as_deref(),
            &ctx.source,
            &cfg.safety,
            &cfg.replication.tables,
        )
        .await?;
        schema::apply_schema(tgt, &cfg.replication.tables, &sql).await?;
        info!("target schema created from pg_dump");
        ctx.progress.step(&format!(
            "schema: {} tables created on target (pg_dump)",
            cfg.replication.tables.len()
        ));
    }
    for info in &pre.tables {
        let cols = schema::target_columns(tgt, &info.table.schema, &info.table.name)
            .await?
            .ok_or_else(|| {
                Error::preflight(format!(
                    "{} missing on target after schema bootstrap",
                    info.table
                ))
            })?;
        for w in schema::verify_columns(&info.table, &info.columns, &cols)? {
            warn!("{w}");
        }
    }

    // Copy tables with bounded parallelism. Each worker owns one source connection (with
    // the exported snapshot) and one target connection (one transaction per table).
    let sem = Arc::new(Semaphore::new(cfg.resources.max_parallel_copies));
    let mut set: JoinSet<Result<(TableName, snapshot::CopyStats)>> = JoinSet::new();
    let deadline = tokio::time::Instant::now() + cfg.resources.snapshot_timeout;
    let buffer = (cfg.resources.buffer_size_kb * 1024) as usize;
    let job_id = cfg.job.id.clone();
    for info in &pre.tables {
        let table = info.table.clone();
        let columns = info.column_names();
        let sem = sem.clone();
        let source = ctx.source.clone();
        let target = ctx.target.clone();
        let safety = cfg.safety.clone();
        let timeout = ctx.timeout();
        let metrics = ctx.metrics.clone();
        let snapshot_name = snapshot_name.to_owned();
        let cancel = ctx.cancel.clone();
        let job_id = job_id.clone();
        let progress = ctx.progress.clone();
        set.spawn(async move {
            let _permit = sem.acquire().await.map_err(|_| Error::cancelled())?;
            if cancel.is_cancelled() {
                return Err(Error::cancelled());
            }
            let src = client::connect(&source, &safety, timeout, "pgsync-copy").await?;
            snapshot::import_snapshot(&src, &snapshot_name).await?;
            let dst = client::connect(&target, &safety, timeout, "pgsync-copy").await?;
            sink::configure_session(&dst).await?;
            if let Err(e) = dst
                .batch_execute("SET session_replication_role = replica")
                .await
            {
                warn!(error = %e, %table, "copying with triggers and FK checks active");
            }
            dst.batch_execute("BEGIN")
                .await
                .map_err(|e| from_pg("BEGIN copy", e))?;
            progress.copy_started(&table.to_string());
            let t0 = Instant::now();
            let stats = tokio::select! {
                r = snapshot::copy_table(&src, &dst, &table, &columns, buffer, &metrics) => r?,
                _ = cancel.cancelled() => return Err(Error::cancelled()),
            };
            PgSink::record_table_copied(
                &dst,
                &job_id,
                &table.schema,
                &table.name,
                stats.rows,
                stats.bytes,
            )
            .await?;
            dst.batch_execute("COMMIT")
                .await
                .map_err(|e| from_pg("COMMIT copy", e))?;
            src.batch_execute("COMMIT")
                .await
                .map_err(|e| from_pg("end snapshot txn", e))?;
            info!(%table, rows = stats.rows, bytes = stats.bytes, "copied");
            progress.copy_done(&table.to_string(), stats.rows, stats.bytes, t0.elapsed());
            Ok((table, stats))
        });
    }
    let mut total_rows = 0u64;
    let mut total_bytes = 0u64;
    while let Some(res) = tokio::select! {
        r = set.join_next() => r,
        _ = tokio::time::sleep_until(deadline) => {
            set.abort_all();
            return Err(Error::resource(format!("snapshot exceeded snapshot_timeout ({:?})", cfg.resources.snapshot_timeout)));
        }
    } {
        match res {
            Ok(Ok((_, stats))) => {
                total_rows += stats.rows;
                total_bytes += stats.bytes;
            }
            Ok(Err(e)) => {
                set.abort_all();
                return Err(e);
            }
            Err(e) => {
                set.abort_all();
                return Err(Error::internal(format!("copy task failed: {e}")));
            }
        }
    }
    let secs = started.elapsed().as_secs_f64();
    info!(
        tables = pre.tables.len(),
        rows = total_rows,
        bytes = total_bytes,
        seconds = format!("{secs:.1}"),
        mb_per_s = format!("{:.1}", total_bytes as f64 / 1e6 / secs.max(0.001)),
        "initial copy finished"
    );
    ctx.progress.step(&format!(
        "initial copy complete: {} rows in {secs:.1}s",
        group_digits(total_rows)
    ));
    Ok(())
}

fn short_version(v: &str) -> String {
    // "PostgreSQL 18.6 (Debian ...) on ..." -> "PostgreSQL 18.6"
    v.split(" (")
        .next()
        .unwrap_or(v)
        .split(" on ")
        .next()
        .unwrap_or(v)
        .to_owned()
}

pub fn group_digits(n: u64) -> String {
    let s = n.to_string();
    let mut out = String::new();
    for (i, c) in s.chars().enumerate() {
        if i > 0 && (s.len() - i) % 3 == 0 {
            out.push(',');
        }
        out.push(c);
    }
    out
}

// ---- run -----------------------------------------------------------------------------

pub async fn run(ctx: &JobContext) -> Result<Checkpoint> {
    let cfg = &ctx.config;
    let monitor_cancel = ctx.cancel.child_token();
    let monitor_reason: Arc<Mutex<Option<Error>>> = Arc::new(Mutex::new(None));
    let monitor = tokio::spawn(monitor_task(
        ctx.source.clone(),
        cfg.safety.clone(),
        ctx.timeout(),
        cfg.source.slot.clone(),
        cfg.safety.wal_retention_alarm_bytes,
        cfg.resources.max_replication_lag_bytes,
        cfg.safety.wal_retention_action,
        cfg.safety.monitor_interval,
        ctx.metrics.clone(),
        monitor_cancel.clone(),
        ctx.cancel.clone(),
        monitor_reason.clone(),
    ));
    let reporter = spawn_reporter(ctx, monitor_cancel.clone());
    if ctx.source_version.lock().unwrap().is_none()
        && let Ok(src) = ctx.source_client("pgsync-run").await
        && let Ok(info) = source::server_info(&src).await
    {
        *ctx.source_version.lock().unwrap() = Some(info.version);
    }

    let result = run_replication(
        &cfg.job.id,
        ctx.stream_factory(),
        PgSinkFactory {
            target: ctx.target.clone(),
            safety: cfg.safety.clone(),
            timeout: ctx.timeout(),
            max_batch_rows: cfg.resources.max_batch_rows,
            max_batch_bytes: cfg.resources.max_batch_bytes as usize,
            tables: cfg.replication.tables.clone(),
        },
        &cfg.resources,
        ctx.metrics.clone(),
        ctx.cancel.clone(),
        ctx.progress.clone(),
    )
    .await;
    monitor_cancel.cancel();
    let _ = monitor.await;
    let _ = reporter.await;

    let monitor_reason = monitor_reason.lock().unwrap().take();
    if let Some(reason) = monitor_reason {
        // The monitor stopped the job: surface its reason and record the state.
        if let Ok(tgt) = ctx.target_client("pgsync-run").await {
            let state = if reason.kind() == ErrorKind::ResyncRequired {
                JobState::ResyncRequired
            } else {
                JobState::Paused
            };
            let _ = sink::set_job_state(&tgt, &cfg.job.id, state, Some(&reason.to_string())).await;
        }
        return Err(reason);
    }
    match result {
        Ok(cp) => {
            info!(committed = %cp.committed_lsn, acked = %cp.acked_lsn, "replication stopped cleanly");
            Ok(cp)
        }
        Err(e) if e.kind() == ErrorKind::Cancelled => {
            // Cancelled before the loop ran (e.g. during connect). Reload the checkpoint.
            let mut s = PgSink::new(ctx.target_client("pgsync-run").await?, ctx.limits()).await?;
            pgsync_core::traits::ReplicationSink::load_checkpoint(&mut s, &cfg.job.id)
                .await?
                .ok_or_else(|| Error::config("job has no checkpoint"))
        }
        Err(e) => Err(e),
    }
}

#[allow(clippy::too_many_arguments)]
async fn monitor_task(
    source: ConnectionInfo,
    safety: pgsync_core::config::SafetyConfig,
    timeout: Duration,
    slot: String,
    retention_alarm: u64,
    lag_alarm: u64,
    action: WalRetentionAction,
    interval: Duration,
    metrics: Arc<Metrics>,
    cancel: CancellationToken,
    job_cancel: CancellationToken,
    reason: Arc<Mutex<Option<Error>>>,
) {
    let mut client: Option<PgClient> = None;
    loop {
        tokio::select! {
            _ = cancel.cancelled() => return,
            _ = tokio::time::sleep(interval) => {}
        }
        if client.is_none() {
            match client::connect(&source, &safety, timeout, "pgsync-monitor").await {
                Ok(c) => client = Some(c),
                Err(e) => {
                    warn!(error = %e, "monitor: cannot connect to source");
                    continue;
                }
            }
        }
        let c = client.as_ref().unwrap();
        let (slot_info, current) = match (
            source::slot_info(c, &slot).await,
            source::current_wal_lsn(c).await,
        ) {
            (Ok(s), Ok(l)) => (s, l),
            (Err(e), _) | (_, Err(e)) => {
                warn!(error = %e, "monitor: source query failed");
                client = None;
                continue;
            }
        };
        metrics.set_lsn(&metrics.server_wal_end, current);
        let Some(s) = slot_info else {
            let e = Error::resync_required(format!("replication slot {slot:?} disappeared"));
            error!("{e}");
            *reason.lock().unwrap() = Some(e);
            job_cancel.cancel();
            return;
        };
        if s.wal_status.as_deref() == Some("lost") || s.invalidation_reason.is_some() {
            let e = Error::resync_required(format!(
                "replication slot {slot:?} invalidated (wal_status={:?}, reason={:?})",
                s.wal_status, s.invalidation_reason
            ));
            error!("{e}");
            *reason.lock().unwrap() = Some(e);
            job_cancel.cancel();
            return;
        }
        let retained = s.restart_lsn.map(|l| current.bytes_since(l)).unwrap_or(0);
        let lag = s
            .confirmed_flush_lsn
            .map(|l| current.bytes_since(l))
            .unwrap_or(0);
        let snap = metrics.snapshot();
        info!(
            slot = %slot,
            active = s.active,
            wal_status = ?s.wal_status,
            retained_wal_bytes = retained,
            slot_lag_bytes = lag,
            apply_lag_bytes = snap.apply_lag_bytes,
            safe_wal_size = ?s.safe_wal_size,
            "wal monitor"
        );
        if lag > lag_alarm {
            warn!(
                slot_lag_bytes = lag,
                threshold = lag_alarm,
                "ALARM: replication lag exceeds max_replication_lag_bytes"
            );
        }
        if retained > retention_alarm {
            warn!(
                retained_wal_bytes = retained,
                threshold = retention_alarm,
                "ALARM: WAL retained by the slot exceeds safety.wal_retention_alarm_bytes"
            );
            if action == WalRetentionAction::Stop {
                let e = Error::resource(format!(
                    "WAL retention {retained} bytes exceeds safety.wal_retention_alarm_bytes ({retention_alarm}); stopping as configured. \
                     Stopping does not release WAL: resume with `pgsync run`, or `pgsync cleanup` to drop the slot."
                ));
                error!("{e}");
                *reason.lock().unwrap() = Some(e);
                job_cancel.cancel();
                return;
            }
        }
    }
}

/// Periodic status line in the log plus the `pgsync.job_runtime` heartbeat row on the
/// target (its own connection, best effort, reconnects on failure). On cancellation a final
/// row with `stopped_at` is written.
async fn report_task(
    target: ConnectionInfo,
    safety: pgsync_core::config::SafetyConfig,
    timeout: Duration,
    job_id: String,
    source_version: Arc<Mutex<Option<String>>>,
    interval: Duration,
    metrics: Arc<Metrics>,
    cancel: CancellationToken,
) {
    let started_at_unix = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0);
    let mut client: Option<PgClient> = None;
    let mut last = (Instant::now(), 0u64, 0u64);
    let mut last_log = Instant::now();
    let mut stopping = false;
    loop {
        if !stopping {
            tokio::select! {
                _ = cancel.cancelled() => stopping = true,
                _ = tokio::time::sleep(interval) => {}
            }
        }
        let s = metrics.snapshot();
        let elapsed = last.0.elapsed().as_secs_f64().max(0.001);
        let txns_per_s = (s.txns_applied.saturating_sub(last.1)) as f64 / elapsed;
        let rows_per_s = (s.rows_applied.saturating_sub(last.2)) as f64 / elapsed;
        last = (Instant::now(), s.txns_applied, s.rows_applied);
        if client.is_none() {
            client = client::connect(&target, &safety, timeout, "pgsync-heartbeat")
                .await
                .ok();
        }
        if let Some(c) = &client {
            let sv = source_version.lock().unwrap().clone();
            let row = sink::RuntimeRow {
                job_id: &job_id,
                started_at_unix,
                stopped: stopping,
                source_version: sv.as_deref(),
                snapshot: &s,
                txns_per_s,
                rows_per_s,
            };
            if let Err(e) = sink::upsert_runtime(c, &row).await {
                warn!(error = %e, "heartbeat: cannot write pgsync.job_runtime");
                client = None;
            }
        }
        if stopping {
            return;
        }
        if last_log.elapsed() < Duration::from_secs(30) {
            continue;
        }
        last_log = Instant::now();
        info!(
            state = %s.state,
            txns_applied = s.txns_applied,
            rows_applied = s.rows_applied,
            events = s.events_received,
            bytes_received = s.bytes_received,
            committed_lsn = %s.committed_lsn,
            acked_lsn = %s.acked_lsn,
            apply_lag_bytes = s.apply_lag_bytes,
            inflight_items = s.inflight_items,
            rss_mb = s.rss_bytes.map(|b| b / (1024 * 1024)),
            reconnects = s.reconnects,
            connection = ?s.connection,
            "status"
        );
    }
}

/// Start the heartbeat/report task for a `bootstrap` or `run`.
fn spawn_reporter(ctx: &JobContext, cancel: CancellationToken) -> tokio::task::JoinHandle<()> {
    tokio::spawn(report_task(
        ctx.target.clone(),
        ctx.config.safety.clone(),
        ctx.timeout(),
        ctx.config.job.id.clone(),
        ctx.source_version.clone(),
        ctx.config.safety.heartbeat_interval,
        ctx.metrics.clone(),
        cancel,
    ))
}

// ---- status / cleanup ----------------------------------------------------------------

#[derive(Serialize)]
pub struct StatusOutput {
    pub job: Option<Checkpoint>,
    /// Last heartbeat of the process running this job, if any was ever written.
    pub runtime: Option<RuntimeStatus>,
    pub source: Option<SourceStatus>,
    pub slot: Option<crate::replication::slot::SlotInfo>,
    pub source_wal_lsn: Option<Lsn>,
    pub slot_lag_bytes: Option<u64>,
    pub retained_wal_bytes: Option<u64>,
    /// Seconds between now and the commit time of the last applied source transaction.
    pub seconds_since_last_commit: Option<f64>,
    pub snapshot: SnapshotStatus,
    pub tables: Vec<TableSyncRow>,
    pub resync_required: bool,
    /// Operator action, when one is needed.
    pub action: Option<String>,
}

#[derive(Serialize)]
pub struct SourceStatus {
    pub version: Option<String>,
    pub system_id: String,
    pub dbname: String,
}

#[derive(Serialize)]
pub struct SnapshotStatus {
    pub tables_total: usize,
    pub tables_copied: usize,
    pub rows_copied: i64,
    pub bytes_copied: i64,
}

#[derive(Serialize)]
pub struct RuntimeStatus {
    pub pid: i32,
    pub hostname: String,
    pub pgsync_version: String,
    pub started_at: String,
    pub updated_at: String,
    pub stopped_at: Option<String>,
    /// Seconds since the heartbeat row was last written.
    pub age_seconds: f64,
    /// True when the row is older than three heartbeat intervals and the process did not
    /// record a clean stop: the process is probably gone.
    pub stale: bool,
    pub state: String,
    pub connection: String,
    pub received_lsn: Lsn,
    pub server_wal_end: Lsn,
    pub committed_lsn: Lsn,
    pub acked_lsn: Lsn,
    pub apply_lag_bytes: i64,
    pub txns_applied: i64,
    pub rows_applied: i64,
    pub txns_per_s: f64,
    pub rows_per_s: f64,
    pub reconnects: i64,
    pub rss_bytes: Option<i64>,
    pub last_apply_at: Option<String>,
}

#[derive(Serialize)]
pub struct TableSyncRow {
    pub table: String,
    pub state: String,
    pub rows_copied: i64,
    pub bytes_copied: i64,
}

pub async fn status(ctx: &JobContext) -> Result<StatusOutput> {
    let tgt = ctx.target_client("pgsync-status").await?;
    let mut s = PgSink::new(tgt, ctx.limits()).await?;
    let job =
        pgsync_core::traits::ReplicationSink::load_checkpoint(&mut s, &ctx.config.job.id).await?;
    let rows = s
        .client()
        .query(
            "SELECT schema_name || '.' || table_name, state, rows_copied, bytes_copied FROM pgsync.table_sync WHERE job_id = $1 ORDER BY 1",
            &[&ctx.config.job.id],
        )
        .await
        .map_err(|e| from_pg("table_sync", e))?;
    let tables = rows
        .iter()
        .map(|r| TableSyncRow {
            table: r.get(0),
            state: r.get(1),
            rows_copied: r.get(2),
            bytes_copied: r.get(3),
        })
        .collect();
    let runtime = s
        .client()
        .query_opt(
            "SELECT pid, hostname, pgsync_version, started_at::text, updated_at::text, stopped_at::text, \
             extract(epoch from now() - updated_at)::float8, state, connection, received_lsn, server_wal_end, \
             committed_lsn, acked_lsn, apply_lag_bytes, txns_applied, rows_applied, txns_per_s, rows_per_s, \
             reconnects, rss_bytes, last_apply_at::text, source_version \
             FROM pgsync.job_runtime WHERE job_id = $1",
            &[&ctx.config.job.id],
        )
        .await
        .map_err(|e| from_pg("job_runtime", e))?;
    let stale_after = ctx.config.safety.heartbeat_interval.as_secs_f64() * 3.0;
    let mut source_version: Option<String> = None;
    let runtime = runtime.map(|r| {
        let age: f64 = r.get(6);
        let stopped: Option<String> = r.get(5);
        source_version = r.get(21);
        let lsn = |i: usize| Lsn(u64::from(r.get::<_, PgLsn>(i)));
        RuntimeStatus {
            pid: r.get(0),
            hostname: r.get(1),
            pgsync_version: r.get(2),
            started_at: r.get(3),
            updated_at: r.get(4),
            stale: stopped.is_none() && age > stale_after,
            stopped_at: stopped,
            age_seconds: age,
            state: r.get(7),
            connection: r.get(8),
            received_lsn: lsn(9),
            server_wal_end: lsn(10),
            committed_lsn: lsn(11),
            acked_lsn: lsn(12),
            apply_lag_bytes: r.get(13),
            txns_applied: r.get(14),
            rows_applied: r.get(15),
            txns_per_s: r.get(16),
            rows_per_s: r.get(17),
            reconnects: r.get(18),
            rss_bytes: r.get(19),
            last_apply_at: r.get(20),
        }
    });
    let (slot, source_wal_lsn, live_version) = match ctx.source_client("pgsync-status").await {
        Ok(src) => {
            let v = src
                .query_one("SELECT version()", &[])
                .await
                .ok()
                .map(|r| r.get::<_, String>(0));
            (
                source::slot_info(&src, &ctx.config.source.slot).await?,
                Some(source::current_wal_lsn(&src).await?),
                v,
            )
        }
        Err(e) => {
            warn!(error = %e, "source unreachable");
            (None, None, None)
        }
    };
    let source = job.as_ref().map(|j| SourceStatus {
        version: live_version.or(source_version),
        system_id: j.source_system_id.clone(),
        dbname: j.source_dbname.clone(),
    });
    let seconds_since_last_commit = job.as_ref().and_then(|j| j.last_commit_ts).map(|ts| {
        let unix_micros = pgsync_core::event::pg_micros_to_unix_micros(ts);
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_micros() as i64)
            .unwrap_or(0);
        (now - unix_micros) as f64 / 1e6
    });
    let slot_lag_bytes = match (&slot, source_wal_lsn) {
        (Some(s), Some(cur)) => s.confirmed_flush_lsn.map(|l| cur.bytes_since(l)),
        _ => None,
    };
    let retained_wal_bytes = match (&slot, source_wal_lsn) {
        (Some(s), Some(cur)) => s.restart_lsn.map(|l| cur.bytes_since(l)),
        _ => None,
    };
    let tables: Vec<TableSyncRow> = tables;
    let snapshot = SnapshotStatus {
        tables_total: ctx.config.replication.tables.len(),
        tables_copied: tables.iter().filter(|t| t.state == "COPIED").count(),
        rows_copied: tables.iter().map(|t| t.rows_copied).sum(),
        bytes_copied: tables.iter().map(|t| t.bytes_copied).sum(),
    };
    let resync_required = job
        .as_ref()
        .map(|j| j.state == JobState::ResyncRequired)
        .unwrap_or(false)
        || slot
            .as_ref()
            .map(|s| s.wal_status.as_deref() == Some("lost") || s.invalidation_reason.is_some())
            .unwrap_or(false)
        || (job.is_some() && slot.is_none() && source_wal_lsn.is_some());
    let action = match (&job, resync_required) {
        (None, _) => Some("run `pgsync bootstrap`".to_owned()),
        (Some(_), true) => Some("run `pgsync cleanup --yes` then `pgsync bootstrap`".to_owned()),
        (Some(j), false) => match j.state {
            JobState::Snapshotting | JobState::Initializing => {
                Some("snapshot incomplete: run `pgsync bootstrap --restart-snapshot`".to_owned())
            }
            JobState::Error => Some("fix the cause in last_error, then `pgsync run`".to_owned()),
            JobState::Paused => {
                Some("resume with `pgsync run` (WAL is still retained by the slot)".to_owned())
            }
            _ => None,
        },
    };
    Ok(StatusOutput {
        job,
        runtime,
        source,
        slot,
        source_wal_lsn,
        slot_lag_bytes,
        retained_wal_bytes,
        seconds_since_last_commit,
        snapshot,
        tables,
        resync_required,
        action,
    })
}

/// Drop the slot PGSync owns and delete the job metadata. Target data tables are kept.
pub async fn cleanup(ctx: &JobContext) -> Result<Vec<String>> {
    let mut actions = Vec::new();
    let tgt = ctx.target_client("pgsync-cleanup").await?;
    let mut s = PgSink::new(tgt, ctx.limits()).await?;
    let job =
        pgsync_core::traits::ReplicationSink::load_checkpoint(&mut s, &ctx.config.job.id).await?;
    let slot_owned =
        job.as_ref().map(|j| j.slot_owned).unwrap_or(false) || ctx.config.safety.auto_drop_slot;
    let src = ctx.source_client("pgsync-cleanup").await?;
    if let Some(info) = source::slot_info(&src, &ctx.config.source.slot).await? {
        if slot_owned {
            if info.active {
                return Err(Error::preflight(format!(
                    "slot {:?} is active (pid {:?}); stop the consumer first",
                    ctx.config.source.slot, info.active_pid
                )));
            }
            src.execute(
                "SELECT pg_drop_replication_slot($1)",
                &[&ctx.config.source.slot],
            )
            .await
            .map_err(|e| from_pg("drop slot", e))?;
            actions.push(format!(
                "dropped replication slot {:?} on the source",
                ctx.config.source.slot
            ));
        } else {
            actions.push(format!(
                "left replication slot {:?} in place (not owned by this job; set safety.auto_drop_slot to override)",
                ctx.config.source.slot
            ));
        }
    } else {
        actions.push(format!(
            "replication slot {:?} does not exist",
            ctx.config.source.slot
        ));
    }
    if job.is_some() {
        sink::delete_job(s.client(), &ctx.config.job.id).await?;
        actions.push(format!(
            "deleted job {:?} metadata from the target (data tables kept)",
            ctx.config.job.id
        ));
    } else {
        actions.push(format!(
            "job {:?} has no metadata on the target",
            ctx.config.job.id
        ));
    }
    Ok(actions)
}