zc2 0.0.29

//! Write-Ahead Log for crash recovery.
//!
//! Append-only JSON Lines file that tracks the lifecycle of each request:
//! Reserved → Executed → Committed/Failed
//!
//! On broker restart, uncommitted entries are replayed to ensure
//! no credits are lost if the broker crashes after worker execution.
//!
//! Optimizations:
//! - In-memory DashMap index for O(1) update_status lookups
//! - Lock-free `SegQueue` durability buffer; the hot path (`append`) only pushes
//!   — it never holds a lock or `fsync`s. A dedicated flush worker
//!   (`run_flush_worker`, spawned by the broker) drains the queue and `fsync`s
//!   on a 100ms age bound (or sooner when the buffer crosses the size
//!   threshold), so durability stays time-bounded with zero fsync on the
//!   request path (group commit).

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
use std::time::{Duration, Instant};

use crossbeam_queue::SegQueue;

/// Max entries before flushing the write buffer to disk. Large batches keep the
/// `fsync` rate bounded and off the per-request critical path under high load
/// (group commit); the `WRITE_BUFFER_MAX_AGE_MS` bound below still caps how long
/// any entry waits, so durability is time-bounded regardless of batch size.
const WRITE_BUFFER_MAX_ENTRIES: usize = 1024;
/// Max time (ms) before flushing the write buffer to disk.
const WRITE_BUFFER_MAX_AGE_MS: u128 = 100;

/// WAL entry status
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum WalStatus {
    Reserved,
    Executed,
    /// A peer-earn credit has been applied to the in-memory ledger and queued
    /// for dashboard delivery, but has not yet been confirmed flushed. Treated
    /// as non-terminal (like `Reserved`/`Executed`): `read_uncommitted` and
    /// `compact` both keep it around until it's marked `Committed`.
    Earned,
    Committed,
    Failed,
}

/// A single WAL entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalEntry {
    pub request_id: String,
    pub user_id: String,
    pub reservation_id: String,
    pub estimated_cost: f64,
    pub actual_cost: Option<f64>,
    pub worker_id: String,
    pub duration_ms: Option<f64>,
    pub timestamp: DateTime<Utc>,
    pub status: WalStatus,
}

/// Write-Ahead Log
///
/// Hot path (`append`) is fully non-blocking: it inserts into the sharded
/// `DashMap` index and pushes a serialized line onto a lock-free `SegQueue`. It
/// never takes the file lock and never `fsync`s. The `file` mutex + `fsync` live
/// entirely in `flush_buffer`, which is run by the dedicated `run_flush_worker`
/// thread (and by rare maintenance ops: `read_uncommitted`, `compact`, shutdown).
/// So no request thread ever blocks on `fsync`.
pub struct Wal {
    path: PathBuf,
    file: Mutex<File>,
    /// In-memory index: request_id → latest WalEntry (O(1) lookup).
    entries_index: DashMap<String, WalEntry>,
    /// Lock-free durability buffer: serialized lines awaiting fsync.
    pending: SegQueue<String>,
    /// Approximate count of `pending` (drives the size-based flush trigger).
    pending_count: AtomicUsize,
    /// Flush-worker wake predicate: set true when `append` crosses the size
    /// threshold, so the worker flushes early instead of waiting the full age
    /// bound. Paired with `flush_cv`.
    flush_signal: Mutex<bool>,
    /// Condvar the flush worker parks on (woken by `flush_signal` or its own
    /// `WRITE_BUFFER_MAX_AGE_MS` timeout).
    flush_cv: Condvar,
    /// Monotonic clock base for the age-based flush trigger.
    start: Instant,
    /// Nanos (since `start`) of the last flush, for the age-based trigger.
    last_flush_nanos: AtomicU64,
    /// Whether the index has been populated from disk.
    index_populated: AtomicBool,
}

/// Exposed for tests only — reads entries directly from disk (bypasses index).
#[cfg(test)]
pub fn read_all_from_disk(path: &std::path::Path) -> std::io::Result<Vec<WalEntry>> {
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    let mut entries = Vec::new();
    for line in reader.lines() {
        let line = line?;
        if !line.trim().is_empty() {
            if let Ok(e) = serde_json::from_str::<WalEntry>(&line) {
                entries.push(e);
            }
        }
    }
    Ok(entries)
}

impl Wal {
    /// Open or create a WAL file at the given path
    pub fn open<P: AsRef<Path>>(path: P) -> std::io::Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }

        let file = OpenOptions::new().create(true).append(true).open(&path)?;

        let wal = Self {
            path,
            file: Mutex::new(file),
            entries_index: DashMap::new(),
            pending: SegQueue::new(),
            pending_count: AtomicUsize::new(0),
            flush_signal: Mutex::new(false),
            flush_cv: Condvar::new(),
            start: Instant::now(),
            last_flush_nanos: AtomicU64::new(0),
            index_populated: AtomicBool::new(false),
        };

        // Populate index from existing WAL file
        wal.populate_index();

        Ok(wal)
    }

    /// Populate the in-memory index from the WAL file on disk.
    fn populate_index(&self) {
        if let Ok(entries) = self.read_all() {
            for entry in entries {
                self.entries_index.insert(entry.request_id.clone(), entry);
            }
        }
        self.index_populated.store(true, Ordering::Release);
    }

    /// Append a new entry to the WAL
    pub fn append(&self, entry: &WalEntry) -> std::io::Result<()> {
        let line = serde_json::to_string(entry).map_err(std::io::Error::other)?;

        // Update in-memory index (DashMap is sharded — no global lock).
        self.entries_index
            .insert(entry.request_id.clone(), entry.clone());

        // Push onto the lock-free durability buffer. The hot path stops here:
        // no file lock, no fsync — the dedicated flush worker drains the queue.
        self.pending.push(line);
        let n = self.pending_count.fetch_add(1, Ordering::Relaxed) + 1;

        // When the buffer crosses the size threshold, wake the worker early
        // (best-effort, non-blocking): a missed wake just defers to the worker's
        // age-bound timeout, so durability is still capped at the age bound.
        if n >= WRITE_BUFFER_MAX_ENTRIES {
            if let Ok(mut signaled) = self.flush_signal.try_lock() {
                *signaled = true;
                self.flush_cv.notify_one();
            }
        }

        Ok(())
    }

    /// Dedicated WAL flush worker. The broker spawns one of these at startup; it
    /// owns every `fsync` so no request thread ever blocks on durable I/O.
    ///
    /// Parks on `flush_cv` with a `WRITE_BUFFER_MAX_AGE_MS` timeout: it wakes on
    /// the timeout (age-bounded group commit) or early when `append` crosses the
    /// size threshold, then drains the buffer with a single `fsync`. Runs until
    /// `running` is cleared, then does a final drain so nothing is lost on a
    /// graceful shutdown. Blocks the calling thread.
    pub fn run_flush_worker(&self, running: &AtomicBool) {
        let age = Duration::from_millis(WRITE_BUFFER_MAX_AGE_MS as u64);
        while running.load(Ordering::Relaxed) {
            {
                let guard = self.flush_signal.lock().unwrap_or_else(|e| e.into_inner());
                let (mut guard, _timeout) = self
                    .flush_cv
                    .wait_timeout(guard, age)
                    .unwrap_or_else(|e| e.into_inner());
                *guard = false; // consume the signal
            }
            if !running.load(Ordering::Relaxed) {
                break;
            }
            if self.pending_count.load(Ordering::Relaxed) > 0 {
                let _ = self.flush_buffer();
            }
        }
        // Final drain on shutdown.
        let _ = self.flush_buffer();
    }

    /// Update the status of an existing entry by appending a new line
    /// with the same request_id but updated status/fields.
    /// Uses the in-memory index for O(1) lookup instead of reading from disk.
    pub fn update_status(
        &self,
        request_id: &str,
        status: WalStatus,
        actual_cost: Option<f64>,
        duration_ms: Option<f64>,
    ) -> std::io::Result<()> {
        // Look up from in-memory index (O(1))
        let original = self.entries_index.get(request_id).map(|e| e.clone());

        if let Some(orig) = original {
            let updated = WalEntry {
                request_id: orig.request_id.clone(),
                user_id: orig.user_id.clone(),
                reservation_id: orig.reservation_id.clone(),
                estimated_cost: orig.estimated_cost,
                actual_cost: actual_cost.or(orig.actual_cost),
                worker_id: orig.worker_id.clone(),
                duration_ms: duration_ms.or(orig.duration_ms),
                timestamp: Utc::now(),
                status,
            };
            self.append(&updated)?;
        }
        Ok(())
    }

    /// Flush the durability buffer to disk with a single fsync.
    ///
    /// The file lock is taken first and held across the drain+write, so only
    /// one flusher runs at a time and the `SegQueue`'s FIFO order is preserved
    /// on disk (a request's Reserved→Executed→Committed lines stay ordered).
    /// Concurrent `append`s during a flush are lock-free and simply land in the
    /// next batch.
    pub fn flush_buffer(&self) -> std::io::Result<()> {
        let mut file = self
            .file
            .lock()
            .map_err(|_| std::io::Error::other("WAL lock poisoned"))?;

        let mut drained = 0usize;
        let mut wrote_any = false;
        while let Some(line) = self.pending.pop() {
            writeln!(file, "{}", line)?;
            drained += 1;
            wrote_any = true;
        }
        if !wrote_any {
            return Ok(());
        }
        file.flush()?;
        file.sync_data()?;

        self.pending_count.fetch_sub(
            drained.min(self.pending_count.load(Ordering::Relaxed)),
            Ordering::Relaxed,
        );
        self.last_flush_nanos
            .store(self.start.elapsed().as_nanos() as u64, Ordering::Relaxed);

        Ok(())
    }

    /// Read all entries from the WAL file
    fn read_all(&self) -> std::io::Result<Vec<WalEntry>> {
        let file = File::open(&self.path)?;
        let reader = BufReader::new(file);
        let mut entries = Vec::new();

        for line in reader.lines() {
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }
            if let Ok(entry) = serde_json::from_str::<WalEntry>(&line) {
                entries.push(entry);
            }
        }
        Ok(entries)
    }

    /// Read uncommitted entries (the latest status for each request_id
    /// that is not Committed or Failed).
    /// Uses in-memory index when available.
    pub fn read_uncommitted(&self) -> std::io::Result<Vec<WalEntry>> {
        // Flush any buffered writes first
        let _ = self.flush_buffer();

        if self.index_populated.load(Ordering::Acquire) {
            // Fast path: read from DashMap index
            Ok(self
                .entries_index
                .iter()
                .filter(|e| e.status != WalStatus::Committed && e.status != WalStatus::Failed)
                .map(|e| e.value().clone())
                .collect())
        } else {
            // Fallback: read from disk
            let entries = self.read_all()?;
            let mut latest: std::collections::HashMap<String, WalEntry> =
                std::collections::HashMap::new();
            for entry in entries {
                latest.insert(entry.request_id.clone(), entry);
            }
            Ok(latest
                .into_values()
                .filter(|e| e.status != WalStatus::Committed && e.status != WalStatus::Failed)
                .collect())
        }
    }

    /// Path to the WAL file (for tests).
    #[cfg(test)]
    pub fn path(&self) -> &std::path::Path {
        &self.path
    }

    /// Compact the WAL: rewrite keeping only uncommitted entries.
    /// Also clears completed entries from the in-memory index.
    pub fn compact(&self) -> std::io::Result<()> {
        // Flush buffer before compacting
        self.flush_buffer()?;

        let uncommitted = self.read_uncommitted()?;

        // Write to temp file, then rename
        let tmp_path = self.path.with_extension("tmp");
        {
            let mut tmp = File::create(&tmp_path)?;
            for entry in &uncommitted {
                let line = serde_json::to_string(entry).map_err(std::io::Error::other)?;
                writeln!(tmp, "{}", line)?;
            }
            tmp.sync_all()?;
        }

        // Swap files
        fs::rename(&tmp_path, &self.path)?;

        // Reopen the file handle for future appends
        let new_file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.path)?;

        let mut file = self
            .file
            .lock()
            .map_err(|_| std::io::Error::other("WAL lock poisoned"))?;
        *file = new_file;

        // Clear committed/failed entries from index
        let to_remove: Vec<String> = self
            .entries_index
            .iter()
            .filter(|e| e.status == WalStatus::Committed || e.status == WalStatus::Failed)
            .map(|e| e.key().clone())
            .collect();
        for key in to_remove {
            self.entries_index.remove(&key);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    fn tmp_wal(name: &str) -> std::path::PathBuf {
        std::path::PathBuf::from(format!("/tmp/zc_wal_test_{}.jsonl", name))
    }

    fn cleanup(p: &std::path::Path) {
        let _ = fs::remove_file(p);
        let _ = fs::remove_file(p.with_extension("tmp"));
    }

    fn entry(request_id: &str, status: WalStatus) -> WalEntry {
        WalEntry {
            request_id: request_id.to_string(),
            user_id: "alice".to_string(),
            reservation_id: format!("res-{}", request_id),
            estimated_cost: 0.05,
            actual_cost: None,
            worker_id: "worker-1".to_string(),
            duration_ms: None,
            timestamp: Utc::now(),
            status,
        }
    }

    #[test]
    fn test_open_creates_file() {
        let path = tmp_wal("open");
        cleanup(&path);
        let _wal = Wal::open(&path).unwrap();
        assert!(path.exists());
        cleanup(&path);
    }

    #[test]
    fn test_append_then_flush_persists_to_disk() {
        let path = tmp_wal("append");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-1", WalStatus::Reserved)).unwrap();
        wal.flush_buffer().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        assert_eq!(disk.len(), 1);
        assert_eq!(disk[0].request_id, "req-1");
        assert_eq!(disk[0].status, WalStatus::Reserved);

        cleanup(&path);
    }

    #[test]
    fn test_update_status_appends_new_line() {
        let path = tmp_wal("update");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-2", WalStatus::Reserved)).unwrap();
        wal.update_status("req-2", WalStatus::Committed, Some(0.03), Some(100.0))
            .unwrap();
        wal.flush_buffer().unwrap();

        // Two lines on disk: reserved + committed
        let disk = read_all_from_disk(&path).unwrap();
        assert_eq!(disk.len(), 2);
        assert_eq!(disk[1].status, WalStatus::Committed);
        assert_eq!(disk[1].actual_cost, Some(0.03));

        cleanup(&path);
    }

    #[test]
    fn test_update_status_unknown_request_is_noop() {
        let path = tmp_wal("update_noop");
        cleanup(&path);
        let wal = Wal::open(&path).unwrap();
        // Should not panic or error
        wal.update_status("nonexistent", WalStatus::Committed, None, None)
            .unwrap();
        wal.flush_buffer().unwrap();
        let disk = read_all_from_disk(&path).unwrap();
        assert!(disk.is_empty());
        cleanup(&path);
    }

    #[test]
    fn test_read_uncommitted_excludes_committed_and_failed() {
        let path = tmp_wal("uncommitted");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-a", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-b", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-c", WalStatus::Reserved)).unwrap();
        wal.update_status("req-b", WalStatus::Committed, Some(0.01), Some(10.0))
            .unwrap();
        wal.update_status("req-c", WalStatus::Failed, None, None)
            .unwrap();
        wal.flush_buffer().unwrap();

        let uncommitted = wal.read_uncommitted().unwrap();
        let ids: Vec<&str> = uncommitted.iter().map(|e| e.request_id.as_str()).collect();
        assert!(ids.contains(&"req-a"), "req-a should be uncommitted");
        assert!(
            !ids.contains(&"req-b"),
            "req-b (committed) should be excluded"
        );
        assert!(!ids.contains(&"req-c"), "req-c (failed) should be excluded");

        cleanup(&path);
    }

    #[test]
    fn test_compact_removes_committed_entries_from_disk() {
        let path = tmp_wal("compact");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-keep", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-done", WalStatus::Reserved)).unwrap();
        wal.update_status("req-done", WalStatus::Committed, Some(0.01), Some(5.0))
            .unwrap();
        wal.flush_buffer().unwrap();

        wal.compact().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        let ids: Vec<&str> = disk.iter().map(|e| e.request_id.as_str()).collect();
        assert!(
            ids.contains(&"req-keep"),
            "uncommitted entry should survive compaction"
        );
        assert!(
            !ids.contains(&"req-done"),
            "committed entry should be removed"
        );

        cleanup(&path);
    }

    #[test]
    fn test_compact_also_clears_failed_entries() {
        let path = tmp_wal("compact_failed");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-fail", WalStatus::Reserved)).unwrap();
        wal.update_status("req-fail", WalStatus::Failed, None, None)
            .unwrap();
        wal.flush_buffer().unwrap();
        wal.compact().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        assert!(disk.is_empty());

        cleanup(&path);
    }

    #[test]
    fn test_reload_populates_index_from_disk() {
        let path = tmp_wal("reload");
        cleanup(&path);

        // Write and close
        {
            let wal = Wal::open(&path).unwrap();
            wal.append(&entry("req-persist", WalStatus::Reserved))
                .unwrap();
            wal.flush_buffer().unwrap();
        }

        // Reopen: index must be repopulated from disk
        {
            let wal = Wal::open(&path).unwrap();
            let uncommitted = wal.read_uncommitted().unwrap();
            assert!(
                uncommitted.iter().any(|e| e.request_id == "req-persist"),
                "Reloaded WAL must surface persisted reserved entry"
            );
        }

        cleanup(&path);
    }

    #[test]
    fn test_flush_worker_persists_without_manual_flush() {
        use std::sync::Arc;
        use std::time::Duration;

        let path = tmp_wal("flush_worker");
        cleanup(&path);

        let wal = Arc::new(Wal::open(&path).unwrap());
        let running = Arc::new(AtomicBool::new(true));

        // Spawn the dedicated flush worker (what the broker does at startup).
        let w = wal.clone();
        let r = running.clone();
        let handle = std::thread::spawn(move || w.run_flush_worker(&r));

        // append() must NOT fsync inline — the worker persists it within the age bound.
        wal.append(&entry("req-bg", WalStatus::Reserved)).unwrap();

        // Wait comfortably past the age bound for the worker to flush.
        std::thread::sleep(Duration::from_millis(WRITE_BUFFER_MAX_AGE_MS as u64 * 4));

        let disk = read_all_from_disk(&path).unwrap();
        assert_eq!(
            disk.len(),
            1,
            "worker should have persisted the appended entry"
        );
        assert_eq!(disk[0].request_id, "req-bg");

        running.store(false, Ordering::Release);
        handle.join().unwrap();
        cleanup(&path);
    }

    #[test]
    fn test_flush_worker_drains_on_shutdown() {
        use std::sync::Arc;

        let path = tmp_wal("flush_worker_shutdown");
        cleanup(&path);

        let wal = Arc::new(Wal::open(&path).unwrap());
        let running = Arc::new(AtomicBool::new(true));
        let w = wal.clone();
        let r = running.clone();
        let handle = std::thread::spawn(move || w.run_flush_worker(&r));

        // Append then immediately request shutdown; the worker must drain on exit.
        wal.append(&entry("req-drain", WalStatus::Reserved))
            .unwrap();
        running.store(false, Ordering::Release);
        handle.join().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        assert!(
            disk.iter().any(|e| e.request_id == "req-drain"),
            "worker must flush pending entries on shutdown"
        );
        cleanup(&path);
    }

    #[test]
    fn test_earned_status_is_uncommitted_until_marked_committed() {
        let path = tmp_wal("earned_uncommitted");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        let mut e = entry("earn-req-1", WalStatus::Earned);
        e.actual_cost = Some(0.02);
        wal.append(&e).unwrap();
        wal.flush_buffer().unwrap();

        // Earned is non-terminal: it must show up as uncommitted (replay-eligible),
        // same as Reserved/Executed.
        let uncommitted = wal.read_uncommitted().unwrap();
        assert!(
            uncommitted.iter().any(|e| e.request_id == "earn-req-1"),
            "Earned entries must be replay-eligible until marked Committed"
        );

        // compact() must NOT drop it while still Earned.
        wal.compact().unwrap();
        let disk = read_all_from_disk(&path).unwrap();
        assert!(
            disk.iter().any(|e| e.request_id == "earn-req-1"),
            "Earned entries must survive compaction until Committed"
        );

        // Once marked Committed, it's excluded and compaction drops it.
        wal.update_status("earn-req-1", WalStatus::Committed, Some(0.02), None)
            .unwrap();
        wal.flush_buffer().unwrap();
        let uncommitted = wal.read_uncommitted().unwrap();
        assert!(!uncommitted.iter().any(|e| e.request_id == "earn-req-1"));

        cleanup(&path);
    }

    #[test]
    fn test_full_lifecycle_reserved_executed_committed() {
        let path = tmp_wal("lifecycle");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        let e = entry("req-life", WalStatus::Reserved);
        wal.append(&e).unwrap();

        // Transition to Executed
        wal.update_status("req-life", WalStatus::Executed, None, Some(250.0))
            .unwrap();

        // Transition to Committed
        wal.update_status("req-life", WalStatus::Committed, Some(0.04), Some(250.0))
            .unwrap();
        wal.flush_buffer().unwrap();

        // No uncommitted entries should remain
        let uncommitted = wal.read_uncommitted().unwrap();
        assert!(!uncommitted.iter().any(|e| e.request_id == "req-life"));

        cleanup(&path);
    }
}