athena_rs 3.26.2

//! ## Deferred write-through cache
//!
//! When enabled via `deferred_writes.enabled: true` in `config.yaml`, gateway
//! mutations (INSERT / UPDATE / DELETE) are **not** sent to the database
//! immediately. Instead they are:
//!
//! 1. Appended to an optional WAL file on disk for crash recovery.
//! 2. Pushed into an in-memory [`WriteBuffer`].
//! 3. Acknowledged to the caller immediately (optimistic success).
//!
//! A background flush executor ([`spawn_flush_loop`](crate::deferred_write::spawn_flush_loop)) drains the buffer into
//! Postgres on a configurable time / size schedule and then performs the
//! normal scoped cache invalidation so that subsequent reads see fresh data.
//!
//! ### Trade-offs
//!
//! | Mode | Latency | Consistency |
//! |------|---------|-------------|
//! | Default (immediate) | Higher (DB round-trip) | Strong |
//! | Deferred | Lower (cache only) | Eventual (flush window) |
//! | skip_invalidation only | Same as default | Slightly stale |
//!
//! Enable only the `skip_cache_invalidation` flag to keep the cache hot
//! without bypassing the database at all.

use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::Arc;

use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;
use tracing::{info, warn};
use uuid::Uuid;

use crate::api::cache::invalidation::gateway_cache_entry_matches_table_invalidation;
use crate::drivers::postgresql::sqlx_driver::{
    PostgresClientRegistry, delete_rows, insert_row, update_rows,
};
use crate::parser::query_builder::condition::Condition;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Runtime configuration for the deferred write subsystem, derived from the
/// `deferred_writes:` section in `config.yaml`.
#[derive(Debug, Clone)]
pub struct DeferredWriteConfig {
    /// Master toggle. `false` = normal immediate write path.
    pub enabled: bool,
    /// Milliseconds between automatic flush runs. `0` = timer disabled.
    pub batch_window_ms: u64,
    /// Flush when buffer reaches this many entries. `0` = size trigger disabled.
    pub batch_max_size: usize,
    /// Write each entry to the WAL file before acknowledging.
    pub wal_enabled: bool,
    /// Directory where WAL files live.
    pub wal_dir: String,
    /// Skip cache invalidation on mutations even when `enabled` is `false`.
    pub skip_cache_invalidation: bool,
}

impl Default for DeferredWriteConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            batch_window_ms: 1000,
            batch_max_size: 100,
            wal_enabled: true,
            wal_dir: "./data/wal".to_string(),
            skip_cache_invalidation: false,
        }
    }
}

// ---------------------------------------------------------------------------
// Entry types
// ---------------------------------------------------------------------------

/// An equality condition stored alongside a deferred UPDATE or DELETE entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredCondition {
    pub column: String,
    pub value: Value,
}

impl StoredCondition {
    /// Reconstruct as a library [`Condition`] for use in driver calls.
    pub fn to_condition(&self) -> Condition {
        Condition::eq(self.column.clone(), self.value.clone())
    }
}

/// The specific mutation payload for a deferred write.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum DeferredOperation {
    Insert {
        payload: Value,
    },
    Update {
        conditions: Vec<StoredCondition>,
        set_payload: Value,
    },
    Delete {
        conditions: Vec<StoredCondition>,
    },
}

/// A single pending write record, stored in the [`WriteBuffer`] and the WAL.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeferredEntry {
    /// Stable ID used for deduplication and WAL tracking.
    pub id: String,
    /// Logical Postgres client that owns the target table.
    pub client_name: String,
    /// Target table name.
    pub table_name: String,
    /// Mutation to replay against the database during flush.
    pub operation: DeferredOperation,
    /// Wall-clock time when the entry was buffered (Unix ms).
    pub created_at_unix_ms: i64,
    /// Number of flush attempts so far (starts at 0, incremented on each failure).
    #[serde(default)]
    pub flush_attempts: u32,
}

impl DeferredEntry {
    pub fn new_insert(client_name: String, table_name: String, payload: Value) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            client_name,
            table_name,
            operation: DeferredOperation::Insert { payload },
            created_at_unix_ms: chrono::Utc::now().timestamp_millis(),
            flush_attempts: 0,
        }
    }

    pub fn new_update(
        client_name: String,
        table_name: String,
        conditions: Vec<StoredCondition>,
        set_payload: Value,
    ) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            client_name,
            table_name,
            operation: DeferredOperation::Update {
                conditions,
                set_payload,
            },
            created_at_unix_ms: chrono::Utc::now().timestamp_millis(),
            flush_attempts: 0,
        }
    }

    pub fn new_delete(
        client_name: String,
        table_name: String,
        conditions: Vec<StoredCondition>,
    ) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            client_name,
            table_name,
            operation: DeferredOperation::Delete { conditions },
            created_at_unix_ms: chrono::Utc::now().timestamp_millis(),
            flush_attempts: 0,
        }
    }
}

// ---------------------------------------------------------------------------
// Write buffer
// ---------------------------------------------------------------------------

/// Thread-safe in-memory queue of pending deferred writes.
pub struct WriteBuffer {
    entries: Mutex<VecDeque<DeferredEntry>>,
    /// Drain the buffer as soon as it reaches this many entries (`0` = disabled).
    pub batch_max_size: usize,
}

impl WriteBuffer {
    pub fn new(batch_max_size: usize) -> Self {
        Self {
            entries: Mutex::new(VecDeque::new()),
            batch_max_size,
        }
    }

    /// Append an entry.
    pub async fn push(&self, entry: DeferredEntry) {
        self.entries.lock().await.push_back(entry);
    }

    /// Drain and return all pending entries, leaving the buffer empty.
    pub async fn drain_all(&self) -> Vec<DeferredEntry> {
        let mut guard: tokio::sync::MutexGuard<'_, VecDeque<DeferredEntry>> =
            self.entries.lock().await;
        guard.drain(..).collect()
    }

    /// Re-prepend entries that failed a flush attempt (preserves FIFO order).
    pub async fn requeue_failed(&self, mut failed: Vec<DeferredEntry>) {
        if failed.is_empty() {
            return;
        }
        let mut guard = self.entries.lock().await;
        // Re-insert at front so they retry before newly arriving entries.
        failed.reverse();
        for entry in failed {
            guard.push_front(entry);
        }
    }

    /// Number of pending entries.
    pub async fn len(&self) -> usize {
        self.entries.lock().await.len()
    }

    /// `true` when the buffer has reached or exceeded the configured threshold.
    pub async fn is_threshold_exceeded(&self) -> bool {
        self.batch_max_size > 0 && self.len().await >= self.batch_max_size
    }
}

// ---------------------------------------------------------------------------
// Write-Ahead Log
// ---------------------------------------------------------------------------

/// Append-only JSONL Write-Ahead Log for crash recovery.
///
/// Each buffered entry is written as one JSON line.  On a successful flush the
/// file is truncated.  On restart, any remaining lines are recovered and
/// replayed by [`recover_from_wal`].
pub struct WalManager {
    wal_path: PathBuf,
    lock: Mutex<()>,
}

impl WalManager {
    /// Create a new `WalManager`, creating the WAL directory if needed.
    pub fn new(wal_dir: &str) -> std::io::Result<Self> {
        let dir: PathBuf = PathBuf::from(wal_dir);
        std::fs::create_dir_all(&dir)?;
        Ok(Self {
            wal_path: dir.join("pending.wal"),
            lock: Mutex::new(()),
        })
    }

    /// Append a single entry to the WAL file (fsync on every write).
    pub async fn append(&self, entry: &DeferredEntry) -> std::io::Result<()> {
        let _lock: tokio::sync::MutexGuard<'_, ()> = self.lock.lock().await;
        let line: String = serde_json::to_string(entry)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        let mut file = tokio::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.wal_path)
            .await?;
        file.write_all(line.as_bytes()).await?;
        file.write_all(b"\n").await?;
        file.flush().await
    }

    /// Read all entries that were appended since the last [`WalManager::clear`].
    pub async fn read_pending(&self) -> std::io::Result<Vec<DeferredEntry>> {
        let _lock = self.lock.lock().await;
        if !self.wal_path.exists() {
            return Ok(Vec::new());
        }
        let content: String = tokio::fs::read_to_string(&self.wal_path).await?;
        let entries: Vec<DeferredEntry> = content
            .lines()
            .filter(|l| !l.trim().is_empty())
            .filter_map(|line| serde_json::from_str::<DeferredEntry>(line).ok())
            .collect();
        Ok(entries)
    }

    /// Truncate the WAL after a successful flush.
    pub async fn clear(&self) -> std::io::Result<()> {
        let _lock = self.lock.lock().await;
        if self.wal_path.exists() {
            tokio::fs::write(&self.wal_path, b"").await?;
        }
        Ok(())
    }

    /// Rewrite the WAL to contain only the provided entries (used after a
    /// partial flush to preserve only the ones that still need to be retried).
    pub async fn rewrite(&self, entries: &[DeferredEntry]) -> std::io::Result<()> {
        let _lock = self.lock.lock().await;
        if entries.is_empty() {
            if self.wal_path.exists() {
                tokio::fs::write(&self.wal_path, b"").await?;
            }
            return Ok(());
        }
        let mut content = String::new();
        for entry in entries {
            if let Ok(line) = serde_json::to_string(entry) {
                content.push_str(&line);
                content.push('\n');
            }
        }
        tokio::fs::write(&self.wal_path, content.as_bytes()).await
    }

    pub fn wal_path(&self) -> &PathBuf {
        &self.wal_path
    }
}

// ---------------------------------------------------------------------------
// Flush
// ---------------------------------------------------------------------------

/// Summary produced by a single flush run.
#[derive(Debug, Default)]
pub struct FlushSummary {
    pub total: usize,
    pub succeeded: usize,
    pub failed: usize,
    /// Tables that were touched and whose cache should now be invalidated.
    pub invalidated_tables: Vec<(String, String)>, // (client_name, table_name)
}

/// Maximum flush attempts per deferred entry before it is permanently discarded.
///
/// After this many consecutive failures the entry is logged at ERROR level and
/// dropped.  The value is intentionally large enough to survive transient
/// outages without losing writes to a recoverable pool.
const MAX_FLUSH_ATTEMPTS: u32 = 10;

/// Execute all pending entries in `buffer` against their respective Postgres
/// pools and, on success, invalidate the matching entries in `cache`.
///
/// Failed entries are re-queued into `buffer` (up to `MAX_FLUSH_ATTEMPTS`)
/// and the WAL is rewritten to contain only the still-pending entries so that
/// a subsequent crash can recover them.
///
/// Returns a [`FlushSummary`] describing what happened.
pub async fn flush_pending(
    buffer: &WriteBuffer,
    wal: Option<&WalManager>,
    pg_registry: &PostgresClientRegistry,
    cache: &moka::future::Cache<String, serde_json::Value>,
) -> FlushSummary {
    let entries: Vec<DeferredEntry> = buffer.drain_all().await;
    if entries.is_empty() {
        return FlushSummary::default();
    }

    let total: usize = entries.len();
    let mut succeeded: usize = 0usize;
    let mut failed: usize = 0usize;
    let mut dirty: std::collections::HashSet<(String, String)> = std::collections::HashSet::new();
    let mut to_requeue: Vec<DeferredEntry> = Vec::new();

    for entry in entries {
        let pool: sqlx::Pool<sqlx::Postgres> = match pg_registry.get_pool(&entry.client_name) {
            Some(p) => p,
            None => {
                warn!(
                    entry_id = %entry.id,
                    client = %entry.client_name,
                    table = %entry.table_name,
                    "deferred_write flush: no pool for client, re-queuing entry"
                );
                let mut retry = entry;
                retry.flush_attempts += 1;
                if retry.flush_attempts < MAX_FLUSH_ATTEMPTS {
                    to_requeue.push(retry);
                } else {
                    tracing::error!(
                        entry_id = %retry.id,
                        client = %retry.client_name,
                        table = %retry.table_name,
                        attempts = retry.flush_attempts,
                        "deferred_write flush: entry exceeded max attempts, discarding permanently"
                    );
                }
                failed += 1;
                continue;
            }
        };

        let result: Result<(), String> = match &entry.operation {
            DeferredOperation::Insert { payload } => insert_row(&pool, &entry.table_name, payload)
                .await
                .map(|_| ())
                .map_err(|e| format!("{e:?}")),

            DeferredOperation::Update {
                conditions,
                set_payload,
            } => {
                let conds: Vec<Condition> = conditions
                    .iter()
                    .map(StoredCondition::to_condition)
                    .collect();
                update_rows(&pool, &entry.table_name, &conds, set_payload)
                    .await
                    .map(|_| ())
                    .map_err(|e| e.to_string())
            }

            DeferredOperation::Delete { conditions } => {
                let conds: Vec<Condition> = conditions
                    .iter()
                    .map(StoredCondition::to_condition)
                    .collect();
                delete_rows(&pool, &entry.table_name, &conds)
                    .await
                    .map(|_| ())
                    .map_err(|e| e.to_string())
            }
        };

        match result {
            Ok(()) => {
                succeeded += 1;
                dirty.insert((entry.client_name.clone(), entry.table_name.clone()));
            }
            Err(e) => {
                let mut retry = entry;
                retry.flush_attempts += 1;
                if retry.flush_attempts < MAX_FLUSH_ATTEMPTS {
                    warn!(
                        entry_id = %retry.id,
                        client = %retry.client_name,
                        table = %retry.table_name,
                        attempts = retry.flush_attempts,
                        error = %e,
                        "deferred_write flush: entry failed, will be re-queued for retry"
                    );
                    to_requeue.push(retry);
                } else {
                    tracing::error!(
                        entry_id = %retry.id,
                        client = %retry.client_name,
                        table = %retry.table_name,
                        attempts = retry.flush_attempts,
                        error = %e,
                        "deferred_write flush: entry exceeded max attempts, discarding permanently"
                    );
                }
                failed += 1;
            }
        }
    }

    // Invalidate cache for every (client, table) that had at least one
    // successful flush, mirroring the normal post-mutation behaviour.
    let table_name_clone_for_closure = dirty.clone();
    let _ = cache
        .invalidate_entries_if(move |key, _value| {
            table_name_clone_for_closure
                .iter()
                .any(|(_, table)| gateway_cache_entry_matches_table_invalidation(key, table))
        })
        .ok();
    cache.run_pending_tasks().await;

    let invalidated_tables: Vec<(String, String)> = dirty.into_iter().collect();

    // Rewrite WAL to contain only entries that still need to be retried.
    // This preserves crash-recovery durability for failed entries while
    // removing successfully flushed ones from the WAL.
    if let Some(wal_mgr) = wal {
        if to_requeue.is_empty() {
            if let Err(e) = wal_mgr.clear().await {
                warn!(error = %e, "deferred_write: failed to truncate WAL after flush");
            }
        } else {
            if let Err(e) = wal_mgr.rewrite(&to_requeue).await {
                warn!(error = %e, "deferred_write: failed to rewrite WAL with pending entries");
            }
        }
    }

    // Re-push failed entries (under max attempts) back to the front of the buffer.
    if !to_requeue.is_empty() {
        buffer.requeue_failed(to_requeue).await;
    }

    if succeeded > 0 || failed > 0 {
        info!(
            succeeded,
            failed,
            total,
            tables = ?invalidated_tables.iter().map(|(_, t)| t).collect::<Vec<_>>(),
            "deferred_write flush complete"
        );
    }

    FlushSummary {
        total,
        succeeded,
        failed,
        invalidated_tables,
    }
}

/// Recover any entries written to the WAL since the last successful flush and
/// push them back into `buffer` so the next flush run picks them up.
pub async fn recover_from_wal(wal: &WalManager, buffer: &WriteBuffer) {
    match wal.read_pending().await {
        Ok(entries) if entries.is_empty() => {}
        Ok(entries) => {
            let count: usize = entries.len();
            for entry in entries {
                buffer.push(entry).await;
            }
            info!(
                recovered = count,
                wal_path = %wal.wal_path().display(),
                "deferred_write: recovered entries from WAL"
            );
        }
        Err(e) => {
            warn!(
                error = %e,
                "deferred_write: failed to read WAL for recovery, entries may be lost"
            );
        }
    }
}

/// Spawn a background task that periodically flushes the write buffer.
///
/// The loop also triggers an early flush whenever [`WriteBuffer::is_threshold_exceeded`]
/// is true, so large write storms are drained without waiting for the full window.
pub fn spawn_flush_loop(
    buffer: Arc<WriteBuffer>,
    wal: Option<Arc<WalManager>>,
    pg_registry: Arc<PostgresClientRegistry>,
    cache: Arc<moka::future::Cache<String, serde_json::Value>>,
    batch_window_ms: u64,
) {
    if batch_window_ms == 0 {
        return;
    }
    tokio::spawn(async move {
        let interval: std::time::Duration = std::time::Duration::from_millis(batch_window_ms);
        let mut ticker: tokio::time::Interval = tokio::time::interval(interval);
        ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        loop {
            ticker.tick().await;
            flush_pending(&buffer, wal.as_deref(), &pg_registry, &cache).await;
        }
    });
}

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

    #[tokio::test]
    async fn write_buffer_push_len_and_drain_round_trip() {
        let buffer: WriteBuffer = WriteBuffer::new(10);
        assert_eq!(buffer.len().await, 0);

        buffer
            .push(DeferredEntry::new_insert(
                "athena_logging".to_string(),
                "bench_table".to_string(),
                serde_json::json!({ "name": "row-1" }),
            ))
            .await;

        assert_eq!(buffer.len().await, 1);
        let drained = buffer.drain_all().await;
        assert_eq!(drained.len(), 1);
        assert_eq!(buffer.len().await, 0);
    }

    #[tokio::test]
    async fn write_buffer_threshold_respects_batch_size() {
        let buffer: WriteBuffer = WriteBuffer::new(2);
        assert!(!buffer.is_threshold_exceeded().await);

        buffer
            .push(DeferredEntry::new_insert(
                "athena_logging".to_string(),
                "bench_table".to_string(),
                serde_json::json!({ "name": "row-1" }),
            ))
            .await;
        assert!(!buffer.is_threshold_exceeded().await);

        buffer
            .push(DeferredEntry::new_insert(
                "athena_logging".to_string(),
                "bench_table".to_string(),
                serde_json::json!({ "name": "row-2" }),
            ))
            .await;
        assert!(buffer.is_threshold_exceeded().await);
    }

    #[tokio::test]
    async fn wal_append_read_pending_and_clear() {
        let wal_dir: PathBuf =
            std::env::temp_dir().join(format!("athena_dw_test_{}", Uuid::new_v4()));
        let wal: WalManager =
            WalManager::new(wal_dir.to_string_lossy().as_ref()).expect("wal create");

        let entry: DeferredEntry = DeferredEntry::new_insert(
            "athena_logging".to_string(),
            "bench_table".to_string(),
            serde_json::json!({ "name": "row-1" }),
        );

        wal.append(&entry).await.expect("wal append");
        let pending: Vec<DeferredEntry> = wal.read_pending().await.expect("wal read");
        assert_eq!(pending.len(), 1);
        assert_eq!(pending[0].client_name, "athena_logging");
        assert_eq!(pending[0].table_name, "bench_table");

        wal.clear().await.expect("wal clear");
        let after_clear: Vec<DeferredEntry> =
            wal.read_pending().await.expect("wal read after clear");
        assert!(after_clear.is_empty());

        let _ = std::fs::remove_file(wal.wal_path());
        let _ = std::fs::remove_dir_all(wal_dir);
    }

    #[tokio::test]
    async fn recover_from_wal_pushes_entries_back_into_buffer() {
        let wal_dir: PathBuf =
            std::env::temp_dir().join(format!("athena_dw_test_{}", Uuid::new_v4()));
        let wal: WalManager =
            WalManager::new(wal_dir.to_string_lossy().as_ref()).expect("wal create");
        let buffer: WriteBuffer = WriteBuffer::new(10);

        wal.append(&DeferredEntry::new_insert(
            "athena_logging".to_string(),
            "bench_table".to_string(),
            serde_json::json!({ "name": "row-1" }),
        ))
        .await
        .expect("append 1");
        wal.append(&DeferredEntry::new_insert(
            "athena_logging".to_string(),
            "bench_table".to_string(),
            serde_json::json!({ "name": "row-2" }),
        ))
        .await
        .expect("append 2");

        recover_from_wal(&wal, &buffer).await;
        assert_eq!(buffer.len().await, 2);

        let _ = std::fs::remove_file(wal.wal_path());
        let _ = std::fs::remove_dir_all(wal_dir);
    }
}