storage/

wal.rs

//! Write-Ahead Log (WAL) for Buffer
//!
//! Provides durability and crash recovery through sequential logging:
//! - All mutations logged before application
//! - Supports checkpointing for log compaction
//! - Recovery from incomplete transactions

use crate::snapshot::{SnapshotConfig, SnapshotManager};
use crate::traits::VectorStorage;
use common::{DakeraError, NamespaceId, Result, Vector, VectorId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

/// WAL configuration
#[derive(Debug, Clone)]
pub struct WalConfig {
    /// Directory for WAL files
    pub wal_dir: PathBuf,
    /// Maximum WAL segment size in bytes
    pub max_segment_size: u64,
    /// Sync mode for durability
    pub sync_mode: WalSyncMode,
    /// Maximum entries before forced checkpoint
    pub checkpoint_threshold: u64,
}

impl Default for WalConfig {
    fn default() -> Self {
        Self {
            wal_dir: PathBuf::from("./data/wal"),
            max_segment_size: 64 * 1024 * 1024, // 64MB
            sync_mode: WalSyncMode::EveryWrite,
            checkpoint_threshold: 10000,
        }
    }
}

/// Sync mode for WAL durability
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WalSyncMode {
    /// Sync after every write (safest)
    EveryWrite,
    /// Sync every N writes
    Periodic(u32),
    /// Sync on explicit flush only
    Manual,
}

/// WAL entry types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum WalEntry {
    /// Insert or update vectors
    Upsert {
        namespace: String,
        vectors: Vec<SerializedVector>,
    },
    /// Delete vectors
    Delete { namespace: String, ids: Vec<String> },
    /// Create namespace
    CreateNamespace { namespace: String },
    /// Delete namespace
    DeleteNamespace { namespace: String },
    /// Checkpoint marker
    Checkpoint { lsn: u64 },
    /// Transaction begin
    TxnBegin { txn_id: u64 },
    /// Transaction commit
    TxnCommit { txn_id: u64 },
    /// Transaction rollback
    TxnRollback { txn_id: u64 },
}

/// Serialized vector for WAL storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializedVector {
    pub id: String,
    pub values: Vec<f32>,
    pub metadata: Option<String>,
    /// TTL in seconds, preserved across crash recovery (DAK-7428).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub ttl_seconds: Option<u64>,
    /// Absolute expiry timestamp, preserved across crash recovery (DAK-7428).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<u64>,
}

impl SerializedVector {
    /// Project a `Vector` into its WAL-serializable form.
    pub fn from_vector(v: &Vector) -> Self {
        Self {
            id: v.id.clone(),
            values: v.values.clone(),
            metadata: v.metadata.as_ref().map(|m| m.to_string()),
            ttl_seconds: v.ttl_seconds,
            expires_at: v.expires_at,
        }
    }

    /// Reconstruct a `Vector` from a replayed WAL entry.
    pub fn into_vector(self) -> Vector {
        Vector {
            id: self.id,
            values: self.values,
            metadata: self.metadata.and_then(|s| serde_json::from_str(&s).ok()),
            ttl_seconds: self.ttl_seconds,
            expires_at: self.expires_at,
        }
    }
}

/// WAL segment file
#[derive(Debug)]
struct WalSegment {
    /// Current size
    size: u64,
    /// Starting LSN
    start_lsn: u64,
    /// Ending LSN
    end_lsn: u64,
    /// On-disk path of this segment file. Used to positively identify the
    /// segment currently being appended to so truncation never deletes it
    /// (deleting the active file would unlink the inode the open writer still
    /// targets, silently losing every post-checkpoint write — DAK-7428).
    path: PathBuf,
}

/// Write-Ahead Log manager
pub struct WriteAheadLog {
    config: WalConfig,
    /// Current log sequence number
    lsn: AtomicU64,
    /// Current active segment
    current_segment: RwLock<Option<WalSegment>>,
    /// Active writer
    writer: RwLock<Option<BufWriter<File>>>,
    /// Entries since last checkpoint
    entries_since_checkpoint: AtomicU64,
    /// Last checkpointed LSN
    last_checkpoint_lsn: AtomicU64,
    /// Write count for periodic sync
    write_count: AtomicU64,
    /// Signalled by `append` when `entries_since_checkpoint` first reaches
    /// `checkpoint_threshold`, so a checkpoint driver can bound WAL growth by
    /// entry count under a write burst instead of waiting for the next timed
    /// tick. Without this, `checkpoint_threshold` was tracked but never acted
    /// on — the log could grow without bound between periodic checkpoints
    /// (DAK-7428).
    checkpoint_needed: tokio::sync::Notify,
}

impl WriteAheadLog {
    /// Create a new WAL
    pub fn new(config: WalConfig) -> Result<Self> {
        // Ensure WAL directory exists
        fs::create_dir_all(&config.wal_dir)
            .map_err(|e| DakeraError::Storage(format!("Failed to create WAL dir: {}", e)))?;

        let wal = Self {
            config,
            lsn: AtomicU64::new(0),
            current_segment: RwLock::new(None),
            writer: RwLock::new(None),
            entries_since_checkpoint: AtomicU64::new(0),
            last_checkpoint_lsn: AtomicU64::new(0),
            write_count: AtomicU64::new(0),
            checkpoint_needed: tokio::sync::Notify::new(),
        };

        // Recover LSN from existing segments
        wal.recover_lsn()?;

        Ok(wal)
    }

    /// Recover LSN from existing WAL segments
    fn recover_lsn(&self) -> Result<()> {
        let segments = self.list_segments()?;
        if let Some(last_segment) = segments.last() {
            // Read the last segment to find the highest LSN
            let entries = self.read_segment(last_segment)?;
            if let Some(last_entry) = entries.last() {
                self.lsn.store(last_entry.0 + 1, Ordering::SeqCst);
            }
        }
        Ok(())
    }

    /// List all WAL segment files
    fn list_segments(&self) -> Result<Vec<PathBuf>> {
        let mut segments = Vec::new();

        if let Ok(entries) = fs::read_dir(&self.config.wal_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.extension().map(|e| e == "wal").unwrap_or(false) {
                    segments.push(path);
                }
            }
        }

        segments.sort();
        Ok(segments)
    }

    /// Read entries from a segment file
    fn read_segment(&self, path: &Path) -> Result<Vec<(u64, WalEntry)>> {
        let file = File::open(path)
            .map_err(|e| DakeraError::Storage(format!("Failed to open WAL: {}", e)))?;

        let reader = BufReader::new(file);
        let mut entries = Vec::new();

        for line_result in reader.lines() {
            let line =
                line_result.map_err(|e| DakeraError::Storage(format!("WAL read error: {}", e)))?;

            if line.trim().is_empty() {
                continue;
            }

            // Format: LSN|JSON_ENTRY
            if let Some((lsn_str, entry_json)) = line.split_once('|') {
                let lsn: u64 = lsn_str.parse().unwrap_or(0);
                if let Ok(entry) = serde_json::from_str::<WalEntry>(entry_json) {
                    entries.push((lsn, entry));
                }
            }
        }

        Ok(entries)
    }

    /// Append an entry to the WAL
    pub fn append(&self, entry: WalEntry) -> Result<u64> {
        let lsn = self.lsn.fetch_add(1, Ordering::SeqCst);

        // Ensure we have an active segment
        self.ensure_segment()?;

        // Serialize and write
        let entry_json = serde_json::to_string(&entry)
            .map_err(|e| DakeraError::Storage(format!("WAL serialize error: {}", e)))?;

        let line = format!("{}|{}\n", lsn, entry_json);

        {
            let mut writer_guard = self.writer.write();
            let writer = writer_guard.as_mut().ok_or_else(|| {
                DakeraError::Storage("WAL writer not available after ensure_segment".to_string())
            })?;

            writer
                .write_all(line.as_bytes())
                .map_err(|e| DakeraError::Storage(format!("WAL write error: {}", e)))?;

            // Handle sync based on mode. `flush()` only drains the userspace
            // BufWriter into the OS page cache — an OS crash or power loss would
            // still lose "durably logged" entries. A write-ahead log's whole job
            // is to survive that, so the syncing modes additionally fsync the
            // file data to physical storage via `sync_data()` (fdatasync). This
            // makes `EveryWrite` ("safest") and `Periodic` honour the crash
            // durability their docs promise instead of only page-cache flushing
            // (DAK-7428). `Manual` deliberately defers all persistence to an
            // explicit `flush()`/checkpoint for throughput.
            let write_count = self.write_count.fetch_add(1, Ordering::Relaxed) + 1;
            match self.config.sync_mode {
                WalSyncMode::EveryWrite => {
                    if let Err(e) = writer.flush() {
                        tracing::warn!(error = %e, "WAL flush failed during EveryWrite sync");
                    } else if let Err(e) = writer.get_ref().sync_data() {
                        tracing::warn!(error = %e, "WAL fsync failed during EveryWrite sync");
                    }
                }
                WalSyncMode::Periodic(n) if n > 0 && write_count.is_multiple_of(n as u64) => {
                    if let Err(e) = writer.flush() {
                        tracing::warn!(error = %e, write_count, "WAL flush failed during periodic sync");
                    } else if let Err(e) = writer.get_ref().sync_data() {
                        tracing::warn!(error = %e, write_count, "WAL fsync failed during periodic sync");
                    }
                }
                _ => {}
            }
        }

        // Update segment size
        {
            let mut segment_guard = self.current_segment.write();
            if let Some(ref mut segment) = *segment_guard {
                segment.size += line.len() as u64;
                segment.end_lsn = lsn;
            }
        }

        // Track entries for checkpoint threshold. When the count first crosses
        // the configured threshold, wake the checkpoint driver so WAL growth is
        // bounded by entry count (not just by the periodic timer) under bursty
        // writes (DAK-7428).
        let entries = self
            .entries_since_checkpoint
            .fetch_add(1, Ordering::Relaxed)
            + 1;
        if self.config.checkpoint_threshold > 0 && entries == self.config.checkpoint_threshold {
            self.checkpoint_needed.notify_one();
        }

        Ok(lsn)
    }

    /// Ensure we have an active segment for writing
    fn ensure_segment(&self) -> Result<()> {
        let needs_new_segment = {
            let segment_guard = self.current_segment.read();
            match &*segment_guard {
                None => true,
                Some(seg) => seg.size >= self.config.max_segment_size,
            }
        };

        if needs_new_segment {
            self.rotate_segment()?;
        }

        Ok(())
    }

    /// Rotate to a new segment file
    fn rotate_segment(&self) -> Result<()> {
        // Close current writer
        {
            let mut writer_guard = self.writer.write();
            if let Some(ref mut writer) = *writer_guard {
                if let Err(e) = writer.flush() {
                    tracing::warn!(error = %e, "WAL flush failed during segment rotation");
                }
            }
            *writer_guard = None;
        }

        // Create new segment
        let current_lsn = self.lsn.load(Ordering::SeqCst);
        let segment_id = current_lsn;
        let segment_path = self.config.wal_dir.join(format!("{:020}.wal", segment_id));

        let file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&segment_path)
            .map_err(|e| DakeraError::Storage(format!("Failed to create WAL segment: {}", e)))?;

        let writer = BufWriter::new(file);

        // Update state
        {
            let mut segment_guard = self.current_segment.write();
            *segment_guard = Some(WalSegment {
                size: 0,
                start_lsn: current_lsn,
                end_lsn: current_lsn,
                path: segment_path.clone(),
            });
        }

        tracing::debug!(
            segment_id = segment_id,
            path = %segment_path.display(),
            "Created new WAL segment"
        );

        {
            let mut writer_guard = self.writer.write();
            *writer_guard = Some(writer);
        }

        Ok(())
    }

    /// Write a checkpoint marker
    pub fn checkpoint(&self) -> Result<u64> {
        let lsn = self.lsn.load(Ordering::SeqCst);

        // Write checkpoint entry
        self.append(WalEntry::Checkpoint { lsn })?;

        // Update checkpoint tracking
        self.last_checkpoint_lsn.store(lsn, Ordering::SeqCst);
        self.entries_since_checkpoint.store(0, Ordering::SeqCst);

        // Flush writer
        {
            let mut writer_guard = self.writer.write();
            if let Some(ref mut writer) = *writer_guard {
                if let Err(e) = writer.flush() {
                    tracing::warn!(error = %e, lsn, "WAL flush failed during checkpoint");
                }
            }
        }

        Ok(lsn)
    }

    /// Recover entries since last checkpoint
    pub fn recover(&self) -> Result<Vec<WalEntry>> {
        let segments = self.list_segments()?;
        let checkpoint_lsn = self.last_checkpoint_lsn.load(Ordering::SeqCst);

        let mut entries = Vec::new();

        for segment_path in segments {
            let segment_entries = self.read_segment(&segment_path)?;

            for (lsn, entry) in segment_entries {
                // Skip entries before checkpoint (but only if checkpoint was done)
                if checkpoint_lsn > 0 && lsn <= checkpoint_lsn {
                    // But track checkpoint LSN updates
                    if let WalEntry::Checkpoint { lsn: cp_lsn } = entry {
                        self.last_checkpoint_lsn.store(cp_lsn, Ordering::SeqCst);
                    }
                    continue;
                }

                // Skip transaction control entries for recovery
                match entry {
                    WalEntry::TxnBegin { .. }
                    | WalEntry::TxnCommit { .. }
                    | WalEntry::TxnRollback { .. }
                    | WalEntry::Checkpoint { .. } => continue,
                    _ => entries.push(entry),
                }
            }
        }

        Ok(entries)
    }

    /// Truncate WAL up to (and including) the given LSN
    pub fn truncate(&self, up_to_lsn: u64) -> Result<u64> {
        let segments = self.list_segments()?;
        let mut removed_count = 0u64;

        // Identify the active segment by its on-disk path. We must NEVER delete
        // the file the open writer is appending to: its LSN range still grows
        // with post-checkpoint writes, and unlinking it would drop those writes
        // into a deleted inode (the exact bug that lost the WAL tail — DAK-7428).
        // Path identity is used deliberately instead of comparing start LSNs,
        // which do not line up with a segment's first on-disk entry LSN.
        let active_path = {
            let segment_guard = self.current_segment.read();
            segment_guard.as_ref().map(|s| s.path.clone())
        };

        for segment_path in segments {
            // Never truncate the segment currently being written to.
            if active_path.as_deref() == Some(segment_path.as_path()) {
                continue;
            }

            // Check if entire segment is before truncation point
            let segment_entries = self.read_segment(&segment_path)?;

            if let Some((last_lsn, _)) = segment_entries.last() {
                if *last_lsn <= up_to_lsn {
                    // Safe to remove this segment
                    fs::remove_file(&segment_path).ok();
                    removed_count += segment_entries.len() as u64;
                }
            }
        }

        Ok(removed_count)
    }

    /// Get current LSN
    pub fn current_lsn(&self) -> u64 {
        self.lsn.load(Ordering::SeqCst)
    }

    /// Notify handle that fires when pending entries reach `checkpoint_threshold`.
    /// A checkpoint driver awaits `checkpoint_needed().notified()` to trigger a
    /// size-bounded checkpoint between periodic ticks (DAK-7428).
    pub fn checkpoint_needed(&self) -> &tokio::sync::Notify {
        &self.checkpoint_needed
    }

    /// Get WAL statistics
    pub fn stats(&self) -> WalStats {
        let segment_count = self.list_segments().map(|s| s.len()).unwrap_or(0);

        let (current_segment_size, current_segment_entries) = {
            let segment_guard = self.current_segment.read();
            match &*segment_guard {
                Some(seg) => (seg.size, seg.end_lsn.saturating_sub(seg.start_lsn)),
                None => (0, 0),
            }
        };

        WalStats {
            current_lsn: self.lsn.load(Ordering::SeqCst),
            last_checkpoint_lsn: self.last_checkpoint_lsn.load(Ordering::SeqCst),
            segment_count,
            current_segment_size,
            current_segment_entries,
            entries_since_checkpoint: self.entries_since_checkpoint.load(Ordering::Relaxed),
        }
    }

    /// Flush WAL to disk
    pub fn flush(&self) -> Result<()> {
        let mut writer_guard = self.writer.write();
        if let Some(ref mut writer) = *writer_guard {
            writer
                .flush()
                .map_err(|e| DakeraError::Storage(format!("WAL flush error: {}", e)))?;
        }
        Ok(())
    }
}

/// WAL statistics
#[derive(Debug, Clone)]
pub struct WalStats {
    /// Current log sequence number
    pub current_lsn: u64,
    /// Last checkpointed LSN
    pub last_checkpoint_lsn: u64,
    /// Number of segment files
    pub segment_count: usize,
    /// Current segment size in bytes
    pub current_segment_size: u64,
    /// Entries in current segment
    pub current_segment_entries: u64,
    /// Entries since last checkpoint
    pub entries_since_checkpoint: u64,
}

/// WAL-wrapped storage that logs all mutations
pub struct WalStorage<S> {
    /// Underlying storage
    inner: S,
    /// Write-ahead log
    wal: WriteAheadLog,
    /// Optional snapshot manager enabling log truncation without losing the ability
    /// to reconstruct a volatile inner store (DAK-7428). `None` => full-log replay.
    snapshot_mgr: Option<SnapshotManager>,
    /// Serializes checkpoints against writes: mutations take a read lock (concurrent),
    /// `snapshot_and_checkpoint` takes the write lock so the snapshot sees a quiescent
    /// inner store (no in-flight write can hold an LSN that is not yet applied when we
    /// truncate). This is what makes truncation crash-safe under concurrency.
    write_gate: tokio::sync::RwLock<()>,
}

impl<S> WalStorage<S> {
    /// Create a new WAL-wrapped storage
    pub fn new(inner: S, wal_config: WalConfig) -> Result<Self> {
        let wal = WriteAheadLog::new(wal_config)?;
        Ok(Self {
            inner,
            wal,
            snapshot_mgr: None,
            write_gate: tokio::sync::RwLock::new(()),
        })
    }

    /// Enable snapshot-backed checkpointing. With snapshots, `snapshot_and_checkpoint`
    /// can persist inner state and truncate the log to bound its growth; without them
    /// the log is retained in full so a volatile inner can always be replayed (DAK-7428).
    pub fn with_snapshots(mut self, snapshot_config: SnapshotConfig) -> Result<Self> {
        self.snapshot_mgr = Some(SnapshotManager::new(snapshot_config)?);
        Ok(self)
    }

    /// Get WAL reference
    pub fn wal(&self) -> &WriteAheadLog {
        &self.wal
    }

    /// Get inner storage reference
    pub fn inner(&self) -> &S {
        &self.inner
    }

    /// Checkpoint the WAL
    pub fn checkpoint(&self) -> Result<u64> {
        self.wal.checkpoint()
    }
}

impl<S: VectorStorage> WalStorage<S> {
    /// Replay the WAL into the inner storage at startup so a volatile inner
    /// (e.g. `InMemoryStorage`) is fully reconstructed after a crash (DAK-7428).
    /// Applies every logged mutation not superseded by a checkpoint, in LSN
    /// order; returns the number of entries applied.
    pub async fn recover_into_inner(&self) -> Result<usize> {
        // If checkpointing is enabled, first restore the most recent snapshot so the
        // inner store starts from the last checkpoint; the retained WAL tail is then
        // replayed on top (idempotently) to reach the latest state (DAK-7428).
        if let Some(mgr) = &self.snapshot_mgr {
            if let Some(latest) = mgr
                .list_snapshots()?
                .into_iter()
                .max_by_key(|m| m.created_at)
            {
                let restored = mgr.restore_snapshot(&self.inner, &latest.id).await?;
                tracing::info!(
                    snapshot_id = %latest.id,
                    namespaces = restored.namespaces_restored,
                    vectors = restored.vectors_restored,
                    "WAL recovery restored base snapshot"
                );
            }
        }

        let entries = self.wal.recover()?;
        let mut applied = 0usize;
        for entry in entries {
            match entry {
                WalEntry::CreateNamespace { namespace } => {
                    self.inner.ensure_namespace(&namespace).await?;
                }
                WalEntry::Upsert { namespace, vectors } => {
                    let vecs: Vec<Vector> = vectors
                        .into_iter()
                        .map(SerializedVector::into_vector)
                        .collect();
                    if !vecs.is_empty() {
                        self.inner.upsert(&namespace, vecs).await?;
                    }
                }
                WalEntry::Delete { namespace, ids } => {
                    self.inner.delete(&namespace, &ids).await?;
                }
                WalEntry::DeleteNamespace { namespace } => {
                    self.inner.delete_namespace(&namespace).await?;
                }
                // recover() already filters Txn*/Checkpoint markers; ignore any.
                WalEntry::Checkpoint { .. }
                | WalEntry::TxnBegin { .. }
                | WalEntry::TxnCommit { .. }
                | WalEntry::TxnRollback { .. } => continue,
            }
            applied += 1;
        }
        if applied > 0 {
            tracing::info!(
                entries = applied,
                "WAL recovery replayed mutations into inner storage"
            );
        }
        Ok(applied)
    }

    /// Persist inner state to a durable snapshot, checkpoint the WAL, and truncate the
    /// now-snapshotted prefix — bounding log growth. The write gate is held exclusively
    /// so no mutation is in flight: the snapshot therefore reflects every write up to
    /// `current_lsn`, making the subsequent truncate crash-safe. A crash at any point
    /// still recovers (snapshot + idempotent replay of the retained tail). Returns
    /// `Ok(None)` when snapshots are not enabled.
    pub async fn snapshot_and_checkpoint(&self) -> Result<Option<u64>> {
        let Some(mgr) = &self.snapshot_mgr else {
            return Ok(None);
        };
        // Exclusive: block writers so the inner store is quiescent for the snapshot.
        let _gate = self.write_gate.write().await;
        let lsn = self.wal.current_lsn();
        mgr.create_snapshot(&self.inner, Some(format!("wal-checkpoint-{lsn}")))
            .await?;
        self.wal.checkpoint()?;
        let removed = self.wal.truncate(lsn)?;
        tracing::info!(
            lsn,
            segments_removed = removed,
            "WAL snapshot + checkpoint + truncate complete"
        );
        Ok(Some(lsn))
    }
}

#[async_trait::async_trait]
impl<S: VectorStorage> VectorStorage for WalStorage<S> {
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        // Shared gate: concurrent with other writes, but excluded during a checkpoint
        // snapshot so truncation stays crash-safe (DAK-7428).
        let _gate = self.write_gate.read().await;
        // Write-ahead: durably log the mutation BEFORE applying it, so a crash
        // between the log append and the (volatile) apply is recovered by replay.
        let serialized: Vec<SerializedVector> =
            vectors.iter().map(SerializedVector::from_vector).collect();
        self.wal.append(WalEntry::Upsert {
            namespace: namespace.clone(),
            vectors: serialized,
        })?;
        self.inner.upsert(namespace, vectors).await
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        self.inner.get(namespace, ids).await
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        self.inner.get_all(namespace).await
    }

    async fn get_all_meta(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        self.inner.get_all_meta(namespace).await
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        let _gate = self.write_gate.read().await;
        self.wal.append(WalEntry::Delete {
            namespace: namespace.clone(),
            ids: ids.to_vec(),
        })?;
        self.inner.delete(namespace, ids).await
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        self.inner.namespace_exists(namespace).await
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        let _gate = self.write_gate.read().await;
        self.wal.append(WalEntry::CreateNamespace {
            namespace: namespace.clone(),
        })?;
        self.inner.ensure_namespace(namespace).await
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        self.inner.count(namespace).await
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        self.inner.dimension(namespace).await
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        self.inner.list_namespaces().await
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        let _gate = self.write_gate.read().await;
        self.wal.append(WalEntry::DeleteNamespace {
            namespace: namespace.clone(),
        })?;
        self.inner.delete_namespace(namespace).await
    }

    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
        // TTL expiry is deterministic from vector state, so it is reconstructed by
        // replay + a subsequent sweep — no need to WAL-log derived cleanups.
        self.inner.cleanup_expired(namespace).await
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        self.inner.cleanup_all_expired().await
    }

    async fn get_page(
        &self,
        namespace: &NamespaceId,
        offset: usize,
        limit: usize,
    ) -> Result<Vec<Vector>> {
        self.inner.get_page(namespace, offset, limit).await
    }

    async fn reclaim_derived_caches(&self) {
        self.inner.reclaim_derived_caches().await;
    }
}

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

    fn test_config(dir: &Path) -> WalConfig {
        WalConfig {
            wal_dir: dir.to_path_buf(),
            max_segment_size: 1024,
            sync_mode: WalSyncMode::EveryWrite,
            checkpoint_threshold: 100,
        }
    }

    fn tv(id: &str) -> Vector {
        Vector {
            id: id.to_string(),
            values: vec![0.1, 0.2, 0.3],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    // DAK-7428: snapshot + checkpoint + truncate must not lose data across a crash.
    // Writes before a checkpoint survive via the snapshot; writes after survive via
    // the retained WAL tail; a post-checkpoint delete stays applied.
    #[tokio::test]
    async fn wal_snapshot_checkpoint_recovers_all_data() {
        use crate::InMemoryStorage;
        let wal_dir = TempDir::new().unwrap();
        let snap_dir = TempDir::new().unwrap();
        let ns = "ns1".to_string();
        let wal_cfg = test_config(wal_dir.path());
        let snap_cfg = SnapshotConfig {
            snapshot_dir: snap_dir.path().to_path_buf(),
            ..Default::default()
        };

        // Phase 1: write, checkpoint (snapshot a,b + truncate), then more writes.
        {
            let store = WalStorage::new(InMemoryStorage::new(), wal_cfg.clone())
                .unwrap()
                .with_snapshots(snap_cfg.clone())
                .unwrap();
            store.ensure_namespace(&ns).await.unwrap();
            store.upsert(&ns, vec![tv("a"), tv("b")]).await.unwrap();
            store.snapshot_and_checkpoint().await.unwrap();
            store.upsert(&ns, vec![tv("c")]).await.unwrap();
            store.delete(&ns, &["a".to_string()]).await.unwrap();
        }

        // Phase 2: simulate crash/restart — fresh empty inner, same wal + snapshot dirs.
        let store2 = WalStorage::new(InMemoryStorage::new(), wal_cfg)
            .unwrap()
            .with_snapshots(snap_cfg)
            .unwrap();
        store2.recover_into_inner().await.unwrap();

        let ids: std::collections::HashSet<String> = store2
            .get_all(&ns)
            .await
            .unwrap()
            .into_iter()
            .map(|v| v.id)
            .collect();
        assert!(ids.contains("b"), "b (from snapshot) must survive: {ids:?}");
        assert!(ids.contains("c"), "c (from WAL tail) must survive: {ids:?}");
        assert!(
            !ids.contains("a"),
            "a (deleted post-checkpoint) must stay deleted"
        );
    }

    // DAK-7428: crossing checkpoint_threshold must wake the checkpoint driver so
    // WAL growth is bounded by entry count, not only by the periodic timer.
    #[tokio::test]
    async fn checkpoint_threshold_signals_driver() {
        let dir = TempDir::new().unwrap();
        let mut cfg = test_config(dir.path());
        cfg.checkpoint_threshold = 3;
        let wal = WriteAheadLog::new(cfg).unwrap();
        let entry = || WalEntry::Upsert {
            namespace: "n".to_string(),
            vectors: vec![],
        };
        // First two appends stay below the threshold — no signal yet.
        wal.append(entry()).unwrap();
        wal.append(entry()).unwrap();
        assert!(
            tokio::time::timeout(
                std::time::Duration::from_millis(50),
                wal.checkpoint_needed().notified()
            )
            .await
            .is_err(),
            "must not signal before the threshold is reached"
        );
        // Third append crosses the threshold and stores a wake permit.
        wal.append(entry()).unwrap();
        tokio::time::timeout(
            std::time::Duration::from_millis(100),
            wal.checkpoint_needed().notified(),
        )
        .await
        .expect("crossing checkpoint_threshold must signal the driver");
    }

    #[test]
    fn test_wal_basic_operations() {
        let temp_dir = TempDir::new().unwrap();
        let config = test_config(temp_dir.path());
        let wal = WriteAheadLog::new(config).unwrap();

        // Append some entries
        let lsn1 = wal
            .append(WalEntry::CreateNamespace {
                namespace: "test".to_string(),
            })
            .unwrap();

        let lsn2 = wal
            .append(WalEntry::Upsert {
                namespace: "test".to_string(),
                vectors: vec![SerializedVector {
                    id: "v1".to_string(),
                    values: vec![1.0, 2.0, 3.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            })
            .unwrap();

        assert_eq!(lsn1, 0);
        assert_eq!(lsn2, 1);
        assert_eq!(wal.current_lsn(), 2);
    }

    #[test]
    fn test_wal_recovery() {
        let temp_dir = TempDir::new().unwrap();
        let config = test_config(temp_dir.path());

        // Write some entries
        {
            let wal = WriteAheadLog::new(config.clone()).unwrap();
            wal.append(WalEntry::CreateNamespace {
                namespace: "test".to_string(),
            })
            .unwrap();
            wal.append(WalEntry::Upsert {
                namespace: "test".to_string(),
                vectors: vec![SerializedVector {
                    id: "v1".to_string(),
                    values: vec![1.0, 2.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            })
            .unwrap();
            wal.flush().unwrap();
        }

        // Recover entries
        {
            let wal = WriteAheadLog::new(config).unwrap();
            let entries = wal.recover().unwrap();

            assert_eq!(entries.len(), 2);
            assert!(matches!(entries[0], WalEntry::CreateNamespace { .. }));
            assert!(matches!(entries[1], WalEntry::Upsert { .. }));
        }
    }

    #[test]
    fn test_wal_checkpoint() {
        let temp_dir = TempDir::new().unwrap();
        let config = test_config(temp_dir.path());
        let wal = WriteAheadLog::new(config).unwrap();

        // Write entries and checkpoint
        wal.append(WalEntry::CreateNamespace {
            namespace: "test".to_string(),
        })
        .unwrap();
        let _checkpoint_lsn = wal.checkpoint().unwrap();

        wal.append(WalEntry::Upsert {
            namespace: "test".to_string(),
            vectors: vec![],
        })
        .unwrap();

        let stats = wal.stats();
        assert!(stats.last_checkpoint_lsn > 0);
        assert_eq!(stats.entries_since_checkpoint, 1); // One entry after checkpoint
    }

    #[test]
    fn test_wal_stats() {
        let temp_dir = TempDir::new().unwrap();
        let config = test_config(temp_dir.path());
        let wal = WriteAheadLog::new(config).unwrap();

        for i in 0..5 {
            wal.append(WalEntry::Upsert {
                namespace: "test".to_string(),
                vectors: vec![SerializedVector {
                    id: format!("v{}", i),
                    values: vec![i as f32],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            })
            .unwrap();
        }

        let stats = wal.stats();
        assert_eq!(stats.current_lsn, 5);
        assert_eq!(stats.entries_since_checkpoint, 5);
    }

    #[test]
    fn test_segment_rotation() {
        let temp_dir = TempDir::new().unwrap();
        let config = WalConfig {
            wal_dir: temp_dir.path().to_path_buf(),
            max_segment_size: 100, // Very small to force rotation
            sync_mode: WalSyncMode::EveryWrite,
            checkpoint_threshold: 1000,
        };

        let wal = WriteAheadLog::new(config).unwrap();

        // Write enough to trigger rotation
        for i in 0..10 {
            wal.append(WalEntry::Upsert {
                namespace: "test".to_string(),
                vectors: vec![SerializedVector {
                    id: format!("v{}", i),
                    values: vec![i as f32; 10],
                    metadata: Some("some metadata here".to_string()),
                    ttl_seconds: None,
                    expires_at: None,
                }],
            })
            .unwrap();
        }

        let stats = wal.stats();
        assert!(stats.segment_count > 1);
    }
}