feat: adding wal crash recovery optimizations (#1)

Author: wolfv

src/main.rs

@@ -233,8 +233,16 @@ async fn main() -> anyhow::Result<()> {
         let map_fn = Arc::new(vindex::JsonKeysMapFn::new(&args.vindex_key_field));
 
         let vi = if let Some(wal_path) = &args.vindex_wal_path {
-            tracing::info!("Vindex WAL path: {}", wal_path);
-            vindex::VerifiableIndex::with_wal(map_fn, wal_path)?
+            // Get the integrated_size from the database for WAL validation
+            // This ensures we truncate the WAL to match the database state after a crash
+            let log_state = db.get_log_state().await?;
+            let expected_tree_size = log_state.integrated_size.value();
+            tracing::info!(
+                "Vindex WAL path: {}, expected tree size from DB: {}",
+                wal_path,
+                expected_tree_size
+            );
+            vindex::VerifiableIndex::with_wal(map_fn, wal_path, expected_tree_size)?
         } else {
             vindex::VerifiableIndex::new(map_fn)
         };

src/vindex/mod.rs

@@ -21,11 +21,11 @@ use crate::error::Result;
 use crate::types::LogIndex;
 pub use prefix_tree::{LookupProof, PrefixTree, ProofNode};
 use sha2::{Digest, Sha256};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::io::Write as _;
 use std::path::Path;
 use std::sync::{Arc, RwLock};
-pub use wal::{WalReader, WalWriter};
+pub use wal::{validate_and_truncate_wal, WalReader, WalWriter};
 
 /// A 32-byte SHA256 hash used as a key in the index.
 pub type IndexKey = [u8; 32];
@@ -145,21 +145,49 @@ impl VerifiableIndex {
 
     /// Create a new verifiable index with WAL persistence.
     ///
-    /// If the WAL file exists, it will be replayed to rebuild the index.
-    pub fn with_wal(map_fn: Arc<dyn MapFn>, wal_path: impl AsRef<Path>) -> Result<Self> {
+    /// If the WAL file exists, it will be validated against the expected tree size
+    /// and truncated if necessary (to handle crash recovery), then replayed to
+    /// rebuild the index.
+    ///
+    /// # Arguments
+    /// * `map_fn` - The function to extract keys from entry data
+    /// * `wal_path` - Path to the WAL file
+    /// * `expected_tree_size` - The integrated tree size from the database; used to
+    ///   truncate the WAL if it ran ahead before a crash
+    pub fn with_wal(
+        map_fn: Arc<dyn MapFn>,
+        wal_path: impl AsRef<Path>,
+        expected_tree_size: u64,
+    ) -> Result<Self> {
         let wal_path = wal_path.as_ref();
 
+        // Validate and truncate WAL to match expected tree size
+        // This is critical for crash recovery: if the WAL was flushed but the database
+        // wasn't updated before a crash, we truncate the WAL to avoid duplicates
+        let actual_wal_size = validate_and_truncate_wal(wal_path, expected_tree_size)?;
+        tracing::info!(
+            "WAL validated: expected_tree_size={}, actual_wal_size={}",
+            expected_tree_size,
+            actual_wal_size
+        );
+
         // Create or open WAL
         let mut tree_size = 0u64;
         let mut index: HashMap<IndexKey, Vec<LogIndex>> = HashMap::new();
         let mut prefix_tree = PrefixTree::new();
 
+        // Track seen (idx, key) pairs to prevent duplicates (defense in depth)
+        let mut seen: HashSet<(u64, IndexKey)> = HashSet::new();
+
         // Replay existing WAL if it exists
         if wal_path.exists() {
             let mut reader = WalReader::open(wal_path)?;
             while let Some((idx, keys)) = reader.next_entry()? {
                 for key in keys {
-                    index.entry(key).or_default().push(idx);
+                    // Deduplicate: only add if we haven't seen this (idx, key) pair
+                    if seen.insert((idx.value(), key)) {
+                        index.entry(key).or_default().push(idx);
+                    }
                 }
                 tree_size = tree_size.max(idx.value() + 1);
             }
@@ -169,6 +197,12 @@ impl VerifiableIndex {
                 let value_hash = compute_value_hash(indices);
                 prefix_tree.insert(key, value_hash);
             }
+
+            tracing::info!(
+                "WAL replayed: {} unique keys, tree_size={}",
+                index.len(),
+                tree_size
+            );
         }
 
         let wal_writer = WalWriter::open(wal_path)?;

src/vindex/wal.rs

@@ -7,11 +7,14 @@
 //!
 //! Each line represents one log entry and the keys it maps to.
 //! Empty lines (entries with no keys) are skipped.
+//!
+//! On startup, the WAL is validated and truncated to match the expected
+//! tree size from the database. This prevents duplicate entries after a crash.
 
 use crate::error::{Error, Result};
 use crate::types::LogIndex;
 use std::fs::{File, OpenOptions};
-use std::io::{BufRead, BufReader, BufWriter, Write};
+use std::io::{BufRead, BufReader, BufWriter, Read, Write};
 use std::path::Path;
 
 use super::IndexKey;
@@ -60,11 +63,18 @@ impl WalWriter {
         Ok(())
     }
 
-    /// Flush the WAL to disk.
+    /// Flush the WAL to disk with fsync for durability.
     pub fn flush(&mut self) -> Result<()> {
         self.writer
             .flush()
             .map_err(|e| Error::Internal(format!("failed to flush WAL: {}", e)))?;
+
+        // fsync to ensure data is persisted to disk (not just OS buffer)
+        self.writer
+            .get_ref()
+            .sync_data()
+            .map_err(|e| Error::Internal(format!("failed to sync WAL to disk: {}", e)))?;
+
         Ok(())
     }
 }
@@ -106,6 +116,106 @@ impl WalReader {
     }
 }
 
+/// Validate and truncate the WAL file to match the expected tree size.
+///
+/// This function reads the WAL backwards to find the last complete entry,
+/// and truncates any entries with index >= expected_tree_size.
+/// This is critical for crash recovery: if the WAL was flushed but the database
+/// wasn't updated before a crash, we need to truncate the WAL to match the
+/// database state to avoid duplicate entries.
+///
+/// Returns the actual tree size found in the WAL (may be less than expected if WAL is behind).
+pub fn validate_and_truncate_wal(path: impl AsRef<Path>, expected_tree_size: u64) -> Result<u64> {
+    let path = path.as_ref();
+
+    if !path.exists() {
+        return Ok(0);
+    }
+
+    let mut file = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .open(path)
+        .map_err(|e| Error::Internal(format!("failed to open WAL for validation: {}", e)))?;
+
+    let file_size = file
+        .metadata()
+        .map_err(|e| Error::Internal(format!("failed to get WAL metadata: {}", e)))?
+        .len();
+
+    if file_size == 0 {
+        return Ok(0);
+    }
+
+    // Read the entire file to find all entries and their positions
+    let mut content = String::new();
+    file.read_to_string(&mut content)
+        .map_err(|e| Error::Internal(format!("failed to read WAL: {}", e)))?;
+
+    // Find all line boundaries and parse entries
+    let mut last_valid_pos: u64 = 0;
+    let mut max_valid_idx: Option<u64> = None;
+    let mut current_pos: u64 = 0;
+
+    for line in content.lines() {
+        let line_len = line.len() as u64 + 1; // +1 for newline
+
+        if line.trim().is_empty() {
+            current_pos += line_len;
+            continue;
+        }
+
+        match parse_wal_line(line) {
+            Ok((idx, _)) => {
+                let idx_val = idx.value();
+                if expected_tree_size == 0 || idx_val < expected_tree_size {
+                    // This entry is within bounds
+                    last_valid_pos = current_pos + line_len;
+                    max_valid_idx = Some(match max_valid_idx {
+                        Some(prev) => prev.max(idx_val),
+                        None => idx_val,
+                    });
+                } else {
+                    // Entry is beyond expected tree size - stop here
+                    tracing::warn!(
+                        "WAL entry {} >= expected tree size {}, truncating",
+                        idx_val,
+                        expected_tree_size
+                    );
+                    break;
+                }
+            }
+            Err(e) => {
+                tracing::warn!(
+                    "Failed to parse WAL line, truncating at position {}: {}",
+                    current_pos,
+                    e
+                );
+                break;
+            }
+        }
+
+        current_pos += line_len;
+    }
+
+    // Truncate file if needed
+    if last_valid_pos < file_size {
+        tracing::info!(
+            "Truncating WAL from {} to {} bytes (removing {} bytes)",
+            file_size,
+            last_valid_pos,
+            file_size - last_valid_pos
+        );
+        file.set_len(last_valid_pos)
+            .map_err(|e| Error::Internal(format!("failed to truncate WAL: {}", e)))?;
+        file.sync_all()
+            .map_err(|e| Error::Internal(format!("failed to sync truncated WAL: {}", e)))?;
+    }
+
+    // Return the tree size based on max index found + 1 (since indices are 0-based)
+    Ok(max_valid_idx.map(|idx| idx + 1).unwrap_or(0))
+}
+
 /// Parse a single WAL line.
 fn parse_wal_line(line: &str) -> Result<(LogIndex, Vec<IndexKey>)> {
     let line = line.trim();
@@ -205,4 +315,101 @@ mod tests {
         assert_eq!(idx.value(), 123);
         assert_eq!(keys.len(), 0);
     }
+
+    #[test]
+    fn test_validate_and_truncate_wal_no_truncation_needed() {
+        let temp_file = NamedTempFile::new().unwrap();
+        let path = temp_file.path();
+
+        // Write entries 0-4
+        {
+            let mut writer = WalWriter::open(path).unwrap();
+            let key = [1u8; 32];
+            for i in 0..5 {
+                writer.append(LogIndex::new(i), &[key]).unwrap();
+            }
+            writer.flush().unwrap();
+        }
+
+        // Validate with expected size 5 - no truncation needed
+        let actual_size = validate_and_truncate_wal(path, 5).unwrap();
+        assert_eq!(actual_size, 5);
+
+        // Verify all entries still present
+        let mut reader = WalReader::open(path).unwrap();
+        for i in 0..5 {
+            let (idx, _) = reader.next_entry().unwrap().unwrap();
+            assert_eq!(idx.value(), i);
+        }
+        assert!(reader.next_entry().unwrap().is_none());
+    }
+
+    #[test]
+    fn test_validate_and_truncate_wal_truncates_excess() {
+        let temp_file = NamedTempFile::new().unwrap();
+        let path = temp_file.path();
+
+        // Write entries 0-9
+        {
+            let mut writer = WalWriter::open(path).unwrap();
+            let key = [1u8; 32];
+            for i in 0..10 {
+                writer.append(LogIndex::new(i), &[key]).unwrap();
+            }
+            writer.flush().unwrap();
+        }
+
+        // Validate with expected size 5 - should truncate entries 5-9
+        let actual_size = validate_and_truncate_wal(path, 5).unwrap();
+        assert_eq!(actual_size, 5);
+
+        // Verify only entries 0-4 remain
+        let mut reader = WalReader::open(path).unwrap();
+        for i in 0..5 {
+            let (idx, _) = reader.next_entry().unwrap().unwrap();
+            assert_eq!(idx.value(), i);
+        }
+        assert!(reader.next_entry().unwrap().is_none());
+    }
+
+    #[test]
+    fn test_validate_and_truncate_wal_empty_file() {
+        let temp_file = NamedTempFile::new().unwrap();
+        let path = temp_file.path();
+
+        // Create empty file
+        File::create(path).unwrap();
+
+        let actual_size = validate_and_truncate_wal(path, 5).unwrap();
+        assert_eq!(actual_size, 0);
+    }
+
+    #[test]
+    fn test_validate_and_truncate_wal_nonexistent_file() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let path = temp_dir.path().join("nonexistent.wal");
+
+        let actual_size = validate_and_truncate_wal(&path, 5).unwrap();
+        assert_eq!(actual_size, 0);
+    }
+
+    #[test]
+    fn test_validate_and_truncate_wal_wal_behind_expected() {
+        let temp_file = NamedTempFile::new().unwrap();
+        let path = temp_file.path();
+
+        // Write only entries 0-2 (WAL is behind expected)
+        {
+            let mut writer = WalWriter::open(path).unwrap();
+            let key = [1u8; 32];
+            for i in 0..3 {
+                writer.append(LogIndex::new(i), &[key]).unwrap();
+            }
+            writer.flush().unwrap();
+        }
+
+        // Validate with expected size 10 - WAL is behind, no truncation
+        let actual_size = validate_and_truncate_wal(path, 10).unwrap();
+        assert_eq!(actual_size, 3);
+    }
 }