knut_bifrost/

state.rs

//! Per-table sync progress — "what is already done".
//!
//! Two backends behind one [`SyncState`]:
//!
//! * **File** (default) — one JSON file, written atomically (tmp + rename)
//!   **after** a sink round succeeds, so a crash before the write replays the
//!   round. Single-process.
//! * **Valkey** — the ledger lives in a Valkey hash (Valkey 9+, the Redis-API/
//!   RESP-conformant store FalkorDB itself runs on; reached via the `redis`
//!   crate). The whole ledger is loaded into memory on open and written through
//!   on each advance, so **multiple bifröst instances share one ledger** and a
//!   restart resumes where any instance left off.
//!
//! Both serve `get` from an in-memory cache (sync); `advance` is async (the
//! Valkey backend does a network write-through).

use std::collections::HashMap;
use std::path::PathBuf;

use redis::aio::MultiplexedConnection;
use serde::{Deserialize, Serialize};

use crate::error::Result;

#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct TableState {
    /// Last fully processed Iceberg snapshot id.
    pub snapshot_id: Option<i64>,
    /// Last watermark value, for watermark-column sources.
    pub watermark: Option<String>,
    /// Files already ingested, for file/asset sources (the AIP-82 "which
    /// files are done" ledger). Empty for snapshot/watermark sources.
    #[serde(default)]
    pub done_files: Vec<String>,
    /// ISO-8601 of the last successful round.
    pub synced_at: Option<String>,
    /// **Which CDC backend served this table, and why** — recorded on every
    /// changelog round so an operator reads the decision out of `state_json`
    /// instead of inferring it from timings. Carries the Spark version, whether
    /// the native 4.2 path was expected, whether it was actually usable, and —
    /// when it was not — that it **fell back** and the engine's own reason.
    ///
    /// `#[serde(default)]` + `skip_serializing_if`: a ledger written before this
    /// field existed still loads, and a non-CDC table does not grow the file.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cdc: Option<crate::cdc::CdcCapability>,
}

/// The sync-progress ledger: table name → [`TableState`], over a file or a
/// Valkey hash. Construct with [`SyncState::open`] (file) or
/// [`SyncState::open_valkey`].
#[derive(Debug)]
pub enum SyncState {
    File {
        path: PathBuf,
        tables: HashMap<String, TableState>,
        /// Set when a staged advance has changed the in-memory ledger but the
        /// file has not been rewritten yet. [`SyncState::flush`] persists and
        /// clears it; a clean ledger skips the write entirely (H8).
        dirty: bool,
    },
    Valkey {
        conn: MultiplexedConnection,
        /// Valkey hash key holding `table → JSON(TableState)`.
        key: String,
        tables: HashMap<String, TableState>,
    },
}

impl SyncState {
    /// Open (or initialise) a **file**-backed ledger.
    pub fn open(path: impl Into<PathBuf>) -> Result<Self> {
        let path = path.into();
        let tables = match std::fs::read_to_string(&path) {
            Ok(s) => serde_json::from_str(&s)?,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => HashMap::new(),
            Err(e) => return Err(e.into()),
        };
        Ok(SyncState::File {
            path,
            tables,
            dirty: false,
        })
    }

    /// Open a **Valkey**-backed ledger at hash `key` on `url` (e.g.
    /// `redis://127.0.0.1:6379`). Loads the existing ledger into memory so a
    /// fresh instance resumes from whatever any prior instance committed.
    pub async fn open_valkey(url: &str, key: impl Into<String>) -> Result<Self> {
        let key = key.into();
        let client = redis::Client::open(url)?;
        let mut conn = client.get_multiplexed_async_connection().await?;
        let raw: HashMap<String, String> = redis::cmd("HGETALL")
            .arg(&key)
            .query_async(&mut conn)
            .await?;
        let mut tables = HashMap::with_capacity(raw.len());
        for (table, json) in raw {
            tables.insert(table, serde_json::from_str(&json)?);
        }
        Ok(SyncState::Valkey { conn, key, tables })
    }

    fn cache(&self) -> &HashMap<String, TableState> {
        match self {
            SyncState::File { tables, .. } | SyncState::Valkey { tables, .. } => tables,
        }
    }

    /// The recorded state for `table` (default if never synced). Served from
    /// the in-memory cache — no IO. Clones the state; the round loop uses
    /// [`Self::get_ref`] instead to avoid cloning `done_files` every poll (M3).
    pub fn get(&self, table: &str) -> TableState {
        self.cache().get(table).cloned().unwrap_or_default()
    }

    /// Borrow the recorded state for `table` from the in-memory cache, or
    /// `None` if the table was never synced. No IO, no clone — the hot path
    /// (`Listener::sync_table`, every poll round) reads through this so an idle
    /// round never deep-clones a possibly-large `done_files` ledger (M3). The
    /// missing case is left to the caller, which substitutes a borrowed default
    /// rather than allocating one here.
    pub fn get_ref(&self, table: &str) -> Option<&TableState> {
        self.cache().get(table)
    }

    /// Record a completed round and persist immediately (write-through for
    /// Valkey, atomic file replace for File). Kept for callers that advance a
    /// single table and want it durable at once; the listener's round loop uses
    /// [`Self::advance_staged`] + [`Self::flush`] instead (H8).
    pub async fn advance(&mut self, table: &str, mut state: TableState) -> Result<()> {
        state.synced_at = Some(now_iso8601());
        match self {
            SyncState::File {
                path,
                tables,
                dirty,
            } => {
                tables.insert(table.to_string(), state);
                persist_file(path, tables)?;
                *dirty = false;
                Ok(())
            }
            SyncState::Valkey { conn, key, tables } => {
                write_through_valkey(conn, key, tables, table, state).await
            }
        }
    }

    /// Like [`Self::advance`] but, for the **file** backend, only updates the
    /// in-memory ledger and marks it dirty — the round persists once at the end
    /// via [`Self::flush`] (H8: one file write per round, not per table).
    ///
    /// The **Valkey** backend still writes through per advance: a shared ledger
    /// must be durable per table for multi-instance correctness, so its
    /// behaviour is identical to [`Self::advance`].
    ///
    /// As with `advance`, the caller invokes this only *after* the sink for
    /// `table` succeeded, so a table that errors leaves no staged state and the
    /// crash-replay guarantee holds (the un-flushed in-memory change is lost on
    /// crash exactly like the old un-written one was).
    pub async fn advance_staged(&mut self, table: &str, mut state: TableState) -> Result<()> {
        state.synced_at = Some(now_iso8601());
        match self {
            SyncState::File { tables, dirty, .. } => {
                tables.insert(table.to_string(), state);
                *dirty = true;
                Ok(())
            }
            SyncState::Valkey { conn, key, tables } => {
                write_through_valkey(conn, key, tables, table, state).await
            }
        }
    }

    /// Persist any staged file-backend changes, then clear the dirty flag. A
    /// no-op when nothing changed (skips the write — H8) and for Valkey (which
    /// already wrote through per advance).
    pub fn flush(&mut self) -> Result<()> {
        match self {
            SyncState::File {
                path,
                tables,
                dirty,
            } => {
                if *dirty {
                    persist_file(path, tables)?;
                    *dirty = false;
                }
                Ok(())
            }
            SyncState::Valkey { .. } => Ok(()),
        }
    }
}

/// Valkey write-through: HSET the one table then update the in-memory cache.
async fn write_through_valkey(
    conn: &mut MultiplexedConnection,
    key: &str,
    tables: &mut HashMap<String, TableState>,
    table: &str,
    state: TableState,
) -> Result<()> {
    let json = serde_json::to_string(&state)?;
    redis::cmd("HSET")
        .arg(key)
        .arg(table)
        .arg(json)
        .query_async::<()>(conn)
        .await?;
    tables.insert(table.to_string(), state);
    Ok(())
}

/// Atomic file write: tmp + rename. Compact JSON (it's a machine file).
fn persist_file(path: &PathBuf, tables: &HashMap<String, TableState>) -> Result<()> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, serde_json::to_vec(tables)?)?;
    std::fs::rename(&tmp, path)?;
    Ok(())
}

/// Current time as an RFC3339 / ISO-8601 UTC string, e.g.
/// `2026-06-30T12:00:00+00:00`. Stored in [`TableState::synced_at`], whose doc
/// promises ISO-8601; the previous `@<unix_secs>` encoding was not parsable as
/// a timestamp by any consumer.
fn now_iso8601() -> String {
    let secs = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    chrono::DateTime::<chrono::Utc>::from_timestamp(secs as i64, 0)
        .unwrap_or_else(|| chrono::DateTime::<chrono::Utc>::from_timestamp(0, 0).unwrap())
        .to_rfc3339()
}

#[cfg(test)]
mod now_iso8601_tests {
    use super::now_iso8601;

    #[test]
    fn synced_at_is_parsable_iso8601() {
        // FAIL-ON-BUG: the old `@<unix_secs>` form ("@1751299200") is not
        // ISO-8601 and `parse::<DateTime>()` rejects it. PASS-ON-FIX: a real
        // RFC3339/ISO-8601 string round-trips through chrono's parser.
        let s = now_iso8601();
        assert!(
            !s.starts_with('@'),
            "synced_at must not use the @epoch encoding, got {s:?}"
        );
        let parsed = chrono::DateTime::parse_from_rfc3339(&s);
        assert!(parsed.is_ok(), "synced_at {s:?} must parse as RFC3339");
    }
}