refactor: replace native_db with memdb and add tests

crates/memdb/src/db.rs

@@ -0,0 +1,598 @@
+use crate::error::{DbError, Result};
+use crate::query::{Page, Paginated};
+use crate::wal::{self, WalEntry, WalOp};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::BTreeMap;
+use std::fs::File;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+// ─── Durability ──────────────────────────────────────────────────────────
+
+/// Controls when WAL data is fsynced to disk.
+#[derive(Clone, Debug)]
+pub enum Durability {
+    /// fsync every write — safest, ~250 ops/s.
+    Full,
+    /// Buffer up to `max_ops` writes, then fsync once.
+    /// Call `MemDb::flush()` before shutdown to commit any remaining
+    /// buffered writes.
+    Batch { max_ops: usize },
+    /// Never fsync — fastest (~10k+ ops/s), zero durability.
+    /// Useful for ephemeral caches or when `snapshot()` handles persistence.
+    Off,
+}
+
+impl Default for Durability {
+    fn default() -> Self {
+        Durability::Full
+    }
+}
+
+impl Durability {
+    /// Convenience: batch up to `max_ops` writes per fsync.
+    pub fn batch(max_ops: usize) -> Self {
+        assert!(max_ops > 0);
+        Durability::Batch { max_ops }
+    }
+}
+
+// ─── In-memory state ─────────────────────────────────────────────────────
+
+/// Complete in-memory state, used for deserializing snapshot.json.
+#[derive(Deserialize, Default)]
+struct Snapshot {
+    /// The highest WAL seq covered by this snapshot.
+    last_seq: u64,
+    /// All collection data.
+    data: BTreeMap<String, BTreeMap<String, Value>>,
+}
+
+/// Borrowed snapshot for zero-copy serialization — avoids cloning the
+/// entire dataset when writing a snapshot file.
+#[derive(Serialize)]
+struct SnapshotRef<'a> {
+    last_seq: u64,
+    data: &'a BTreeMap<String, BTreeMap<String, Value>>,
+}
+
+/// Runtime state. All writes are serialized under this lock.
+struct Inner {
+    last_seq: u64,
+    data: BTreeMap<String, BTreeMap<String, Value>>,
+    /// Open WAL file handle, reused across writes.
+    wal_file: Option<File>,
+    /// WAL file path, used when snapshot truncation needs to reopen the handle.
+    wal_path: PathBuf,
+    snapshot_path: PathBuf,
+    durability: Durability,
+    /// Buffered WAL entries not yet flushed to disk (Batch mode).
+    pending: Vec<WalEntry>,
+    /// When the first entry was added to the current batch.
+    pending_since: Option<Instant>,
+}
+
+impl Inner {
+    /// Execute a batch of ops under the lock: allocate seq → apply to memory
+    /// → write WAL (fsync behaviour depends on Durability).
+    fn commit(&mut self, ops: Vec<WalOp>) -> Result<u64> {
+        if ops.is_empty() {
+            return Ok(0);
+        }
+        self.last_seq += 1;
+        let seq = self.last_seq;
+        let entry = WalEntry {
+            seq,
+            ops,
+        };
+        // Always apply to memory first — clients can read their own writes
+        // immediately regardless of durability mode.
+        for op in &entry.ops {
+            apply_op(&mut self.data, op.clone());
+        }
+        // WAL path depends on Durability.
+        match self.durability {
+            Durability::Full => {
+                if let Some(ref mut f) = self.wal_file {
+                    wal::write_entry(f, &entry)?;
+                    wal::sync_wal(f)?;
+                }
+            }
+            Durability::Batch { .. } => {
+                self.push_pending(entry);
+                if self.pending.len() >= self.batch_threshold() {
+                    self.flush_pending()?;
+                }
+            }
+            Durability::Off => {
+                if let Some(ref mut f) = self.wal_file {
+                    wal::write_entry(f, &entry)?;
+                }
+            }
+        }
+        Ok(seq)
+    }
+
+    fn batch_threshold(&self) -> usize {
+        match self.durability {
+            Durability::Batch { max_ops } => max_ops,
+            _ => 0,
+        }
+    }
+
+    fn push_pending(&mut self, entry: WalEntry) {
+        if self.pending.is_empty() {
+            self.pending_since = Some(Instant::now());
+        }
+        self.pending.push(entry);
+    }
+
+    /// Write all buffered entries to WAL and fsync once.
+    fn flush_pending(&mut self) -> Result<usize> {
+        let count = self.pending.len();
+        if count == 0 {
+            return Ok(0);
+        }
+        if let Some(ref mut f) = self.wal_file {
+            for entry in &self.pending {
+                wal::write_entry(f, entry)?;
+            }
+            wal::sync_wal(f)?;
+        }
+        self.pending.clear();
+        self.pending_since = None;
+        Ok(count)
+    }
+
+}
+
+fn apply_op(data: &mut BTreeMap<String, BTreeMap<String, Value>>, op: WalOp) {
+    match op {
+        WalOp::Insert { collection, key, value } => {
+            data.entry(collection).or_default().insert(key, value);
+        }
+        WalOp::Upsert { collection, key, value } => {
+            data.entry(collection).or_default().insert(key, value);
+        }
+        WalOp::Delete { collection, key } => {
+            if let Some(col) = data.get_mut(&collection) {
+                col.remove(&key);
+            }
+        }
+    }
+}
+
+// ─── MemDb ───────────────────────────────────────────────────────────────
+
+#[derive(Clone)]
+pub struct MemDb {
+    /// All writes (WAL + memory) are serialized under this lock.
+    inner: Arc<Mutex<Inner>>,
+}
+
+impl MemDb {
+    /// Open the database with `Durability::Full` (backward-compatible).
+    pub fn open(data_dir: impl AsRef<Path>) -> Result<Self> {
+        Self::open_with(data_dir, Durability::Full)
+    }
+
+    /// Open the database with a specific durability policy.
+    pub fn open_with(data_dir: impl AsRef<Path>, durability: Durability) -> Result<Self> {
+        let dir = data_dir.as_ref();
+        std::fs::create_dir_all(dir)?;
+
+        let snapshot_path = dir.join("snapshot.json");
+        let wal_path = dir.join("wal.jsonl");
+
+        // 1. Read snapshot.
+        let mut snapshot = if snapshot_path.exists() {
+            let bytes = std::fs::read(&snapshot_path)?;
+            serde_json::from_slice::<Snapshot>(&bytes)?
+        } else {
+            Snapshot::default()
+        };
+
+        let after_seq = snapshot.last_seq;
+
+        // 2. Replay WAL entries with seq > last_seq.
+        let entries = wal::read_after(&wal_path, after_seq)?;
+        let replayed = entries.len();
+        for entry in entries {
+            for op in entry.ops {
+                apply_op(&mut snapshot.data, op);
+            }
+            snapshot.last_seq = snapshot.last_seq.max(entry.seq);
+        }
+
+        if replayed > 0 {
+            eprintln!("[memdb] replayed {replayed} WAL entries after seq={after_seq}");
+        }
+
+        // 3. Open WAL file handle for subsequent writes.
+        let wal_file = std::fs::OpenOptions::new()
+            .create(true)
+            .append(true)
+            .open(&wal_path)?;
+
+        Ok(Self {
+            inner: Arc::new(Mutex::new(Inner {
+                last_seq: snapshot.last_seq,
+                data: snapshot.data,
+                wal_file: Some(wal_file),
+                wal_path,
+                snapshot_path,
+                durability,
+                pending: vec![],
+                pending_since: None,
+            })),
+        })
+    }
+
+    /// Pure in-memory mode (for tests, no persistence).
+    pub fn in_memory() -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(Inner {
+                last_seq: 0,
+                data: BTreeMap::new(),
+                wal_file: None,
+                wal_path: PathBuf::from("/dev/null"),
+                snapshot_path: PathBuf::from("/dev/null"),
+                durability: Durability::Off,
+                pending: vec![],
+                pending_since: None,
+            })),
+        }
+    }
+
+    /// Flush any buffered WAL entries to disk.
+    /// Important in `Durability::Batch` mode before shutdown — without this
+    /// call the last buffered batch may be lost on crash.
+    pub fn flush(&self) -> Result<usize> {
+        self.inner.lock().unwrap().flush_pending()
+    }
+
+    /// Return the number of buffered entries not yet flushed to disk.
+    pub fn pending_writes(&self) -> usize {
+        self.inner.lock().unwrap().pending.len()
+    }
+
+    /// Trigger a manual snapshot:
+    /// 1. Flush pending WAL entries (so crash recovery sees them).
+    /// 2. Read (last_seq, data) atomically under the lock.
+    /// 3. Write snapshot file outside the lock (non-blocking for writers).
+    /// 4. Atomic rename ensures snapshot file is never partial.
+    /// 5. Re-lock and truncate the WAL only when no writes raced in between.
+    pub fn snapshot(&self) -> Result<()> {
+        // Flush pending so every committed write is in the WAL before we
+        // potentially truncate it.
+        self.flush()?;
+
+        // Serialize inside the lock — borrows data directly (zero-copy),
+        // then write to disk outside the lock so writers aren't blocked.
+        let (bytes, last_seq, snapshot_path) = {
+            let inner = self.inner.lock().unwrap();
+            let path = inner.snapshot_path.clone();
+            if path == Path::new("/dev/null") {
+                return Ok(());
+            }
+            let snap = SnapshotRef {
+                last_seq: inner.last_seq,
+                data: &inner.data,
+            };
+            (serde_json::to_vec_pretty(&snap)?, inner.last_seq, path)
+        };
+
+        // Write snapshot outside lock so writers are not blocked.
+        let tmp = snapshot_path.with_extension("tmp");
+        std::fs::write(&tmp, bytes)?;
+        std::fs::rename(&tmp, &snapshot_path)?;
+
+        // Re-lock. Only truncate the WAL if no writes have committed since
+        // the snapshot was taken.
+        let mut inner = self.inner.lock().unwrap();
+        if inner.last_seq == last_seq {
+            if let Some(f) = inner.wal_file.take() {
+                drop(f);
+                // Truncate the file to zero.
+                std::fs::OpenOptions::new()
+                    .write(true)
+                    .truncate(true)
+                    .open(&inner.wal_path)?;
+                // Reopen in append mode for future writes.
+                inner.wal_file = Some(
+                    std::fs::OpenOptions::new()
+                        .create(true)
+                        .append(true)
+                        .open(&inner.wal_path)?,
+                );
+            }
+        }
+
+        eprintln!("[memdb] snapshot saved at seq={last_seq}");
+        Ok(())
+    }
+
+    /// Start a background snapshot worker that fires at the given interval.
+    pub fn start_snapshot_worker(&self, interval: Duration) -> tokio::task::JoinHandle<()> {
+        let db = self.clone();
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(interval);
+            ticker.tick().await; // skip immediate first tick
+            loop {
+                ticker.tick().await;
+                if let Err(e) = db.snapshot() {
+                    eprintln!("[memdb] snapshot error: {e}");
+                }
+            }
+        })
+    }
+
+    /// Start a background flush worker for Batch durability mode.
+    /// Guarantees that buffered writes are flushed at least every `interval`.
+    pub fn start_flush_worker(&self, interval: Duration) -> tokio::task::JoinHandle<()> {
+        let db = self.clone();
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(interval);
+            ticker.tick().await;
+            loop {
+                ticker.tick().await;
+                if let Err(e) = db.flush() {
+                    eprintln!("[memdb] flush error: {e}");
+                }
+            }
+        })
+    }
+
+    /// Get a handle to the named collection.
+    pub fn collection(&self, name: &'static str) -> Collection {
+        Collection {
+            db: self.clone(),
+            name,
+        }
+    }
+
+    /// Begin a cross-collection atomic transaction.
+    pub fn transaction(&self) -> Transaction {
+        Transaction {
+            db: self.clone(),
+            ops: vec![],
+        }
+    }
+
+    fn commit(&self, ops: Vec<WalOp>) -> Result<()> {
+        self.inner.lock().unwrap().commit(ops)?;
+        Ok(())
+    }
+}
+
+// ─── Collection ──────────────────────────────────────────────────────────
+
+#[derive(Clone)]
+pub struct Collection {
+    db: MemDb,
+    name: &'static str,
+}
+
+impl Collection {
+    // ── Write operations ──────────────────────────────────────────────────
+
+    /// Insert a record. Returns `DuplicateKey` if the key already exists.
+    pub fn insert<T: Serialize>(&self, key: impl Into<String>, value: &T) -> Result<()> {
+        let key = key.into();
+        let value = serde_json::to_value(value)?;
+        let mut inner = self.db.inner.lock().unwrap();
+        if inner.data.get(self.name).and_then(|c| c.get(&key)).is_some() {
+            return Err(DbError::DuplicateKey(self.name.to_string(), key));
+        }
+        inner.commit(vec![WalOp::Insert {
+            collection: self.name.to_string(),
+            key,
+            value,
+        }])?;
+        Ok(())
+    }
+
+    /// Insert or overwrite a record.
+    pub fn upsert<T: Serialize>(&self, key: impl Into<String>, value: &T) -> Result<()> {
+        let key = key.into();
+        let value = serde_json::to_value(value)?;
+        self.db.commit(vec![WalOp::Upsert {
+            collection: self.name.to_string(),
+            key,
+            value,
+        }])
+    }
+
+    /// Read-modify-write, performed atomically under the lock.
+    /// Returns the updated value.
+    pub fn update<T>(&self, key: impl Into<String>, f: impl FnOnce(T) -> T) -> Result<T>
+    where
+        T: Serialize + for<'de> Deserialize<'de>,
+    {
+        let key = key.into();
+        let mut inner = self.db.inner.lock().unwrap();
+        let old_val = inner
+            .data
+            .get(self.name)
+            .and_then(|c| c.get(&key))
+            .ok_or_else(|| DbError::NotFound(self.name.to_string(), key.clone()))?
+            .clone();
+        let old: T = serde_json::from_value(old_val)?;
+        let updated = f(old);
+        let new_val = serde_json::to_value(&updated)?;
+        inner.commit(vec![WalOp::Upsert {
+            collection: self.name.to_string(),
+            key,
+            value: new_val,
+        }])?;
+        Ok(updated)
+    }
+
+    /// Delete a record. Returns whether the record existed.
+    pub fn delete(&self, key: impl Into<String>) -> Result<bool> {
+        let key = key.into();
+        let mut inner = self.db.inner.lock().unwrap();
+        let existed = inner.data.get(self.name).and_then(|c| c.get(&key)).is_some();
+        if existed {
+            inner.commit(vec![WalOp::Delete {
+                collection: self.name.to_string(),
+                key,
+            }])?;
+        }
+        Ok(existed)
+    }
+
+    // ── Read operations (in-memory, no WAL involvement) ───────────────────
+
+    /// Look up a record by primary key.
+    pub fn get<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<Option<T>> {
+        let inner = self.db.inner.lock().unwrap();
+        match inner.data.get(self.name).and_then(|c| c.get(key)) {
+            Some(v) => Ok(Some(serde_json::from_value(v.clone())?)),
+            None => Ok(None),
+        }
+    }
+
+    /// Look up a record by primary key; return an error if not found.
+    pub fn get_required<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T> {
+        self.get(key)?
+            .ok_or_else(|| DbError::NotFound(self.name.to_string(), key.to_string()))
+    }
+
+    /// Full scan with a predicate filter.
+    pub fn filter<T, F>(&self, predicate: F) -> Result<Vec<T>>
+    where
+        T: for<'de> Deserialize<'de>,
+        F: Fn(&T) -> bool,
+    {
+        let inner = self.db.inner.lock().unwrap();
+        let Some(col) = inner.data.get(self.name) else {
+            return Ok(vec![]);
+        };
+        let mut results = vec![];
+        for v in col.values() {
+            let item: T = serde_json::from_value(v.clone())?;
+            if predicate(&item) {
+                results.push(item);
+            }
+        }
+        Ok(results)
+    }
+
+    /// List all records in key order.
+    pub fn list_all<T: for<'de> Deserialize<'de>>(&self) -> Result<Vec<T>> {
+        let inner = self.db.inner.lock().unwrap();
+        let Some(col) = inner.data.get(self.name) else {
+            return Ok(vec![]);
+        };
+        col.values()
+            .map(|v| serde_json::from_value(v.clone()).map_err(DbError::from))
+            .collect()
+    }
+
+    /// Prefix scan — returns records whose key starts with the given prefix,
+    /// in key order.
+    pub fn scan_prefix<T: for<'de> Deserialize<'de>>(&self, prefix: &str) -> Result<Vec<T>> {
+        let inner = self.db.inner.lock().unwrap();
+        let Some(col) = inner.data.get(self.name) else {
+            return Ok(vec![]);
+        };
+        col.range(prefix.to_string()..)
+            .take_while(|(k, _)| k.starts_with(prefix))
+            .map(|(_, v)| serde_json::from_value(v.clone()).map_err(DbError::from))
+            .collect()
+    }
+
+    /// Paginated query in key order.
+    pub fn paginate<T: for<'de> Deserialize<'de>>(
+        &self,
+        page: &Page,
+        desc: bool,
+    ) -> Result<Paginated<T>> {
+        let inner = self.db.inner.lock().unwrap();
+        let Some(col) = inner.data.get(self.name) else {
+            return Ok(Paginated::empty(page));
+        };
+        let total = col.len() as u64;
+        let offset = page.offset();
+        let items: Vec<T> = if desc {
+            col.values()
+                .rev()
+                .skip(offset)
+                .take(page.page_size as usize)
+                .map(|v| serde_json::from_value(v.clone()).map_err(DbError::from))
+                .collect::<Result<_>>()?
+        } else {
+            col.values()
+                .skip(offset)
+                .take(page.page_size as usize)
+                .map(|v| serde_json::from_value(v.clone()).map_err(DbError::from))
+                .collect::<Result<_>>()?
+        };
+        Ok(Paginated::new(page, total, items))
+    }
+
+    /// Return the number of records in this collection.
+    pub fn count(&self) -> usize {
+        let inner = self.db.inner.lock().unwrap();
+        inner.data.get(self.name).map(|c| c.len()).unwrap_or(0)
+    }
+
+    /// Check whether a key exists.
+    pub fn exists(&self, key: &str) -> bool {
+        let inner = self.db.inner.lock().unwrap();
+        inner.data.get(self.name).and_then(|c| c.get(key)).is_some()
+    }
+}
+
+// ─── Transaction ─────────────────────────────────────────────────────────
+
+pub struct Transaction {
+    db: MemDb,
+    ops: Vec<WalOp>,
+}
+
+impl Transaction {
+    pub fn insert<T: Serialize>(
+        mut self,
+        collection: &str,
+        key: impl Into<String>,
+        value: &T,
+    ) -> Result<Self> {
+        self.ops.push(WalOp::Insert {
+            collection: collection.to_string(),
+            key: key.into(),
+            value: serde_json::to_value(value)?,
+        });
+        Ok(self)
+    }
+
+    pub fn upsert<T: Serialize>(
+        mut self,
+        collection: &str,
+        key: impl Into<String>,
+        value: &T,
+    ) -> Result<Self> {
+        self.ops.push(WalOp::Upsert {
+            collection: collection.to_string(),
+            key: key.into(),
+            value: serde_json::to_value(value)?,
+        });
+        Ok(self)
+    }
+
+    pub fn delete(mut self, collection: &str, key: impl Into<String>) -> Self {
+        self.ops.push(WalOp::Delete {
+            collection: collection.to_string(),
+            key: key.into(),
+        });
+        self
+    }
+
+    /// Atomically commit all buffered ops as a single WAL entry.
+    pub fn commit(self) -> Result<()> {
+        self.db.commit(self.ops)
+    }
+}

crates/memdb/src/error.rs

@@ -0,0 +1,18 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum DbError {
+    #[error("io error: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error("serialization error: {0}")]
+    Serde(#[from] serde_json::Error),
+
+    #[error("record not found: collection={0}, key={1}")]
+    NotFound(String, String),
+
+    #[error("duplicate key: collection={0}, key={1}")]
+    DuplicateKey(String, String),
+}
+
+pub type Result<T> = std::result::Result<T, DbError>;

crates/memdb/src/lib.rs

@@ -0,0 +1,8 @@
+pub mod db;
+pub mod error;
+pub mod query;
+pub mod wal;
+
+pub use db::{Collection, Durability, MemDb, Transaction};
+pub use error::DbError;
+pub use query::{Page, Paginated};

crates/memdb/src/query.rs

@@ -0,0 +1,48 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug)]
+pub struct Page {
+    pub page: u64,
+    pub page_size: u64,
+}
+
+impl Page {
+    pub fn new(page: u64, page_size: u64) -> Self {
+        assert!(page >= 1 && page_size >= 1, "page and page_size must be >= 1");
+        Self { page, page_size }
+    }
+
+    pub fn offset(&self) -> usize {
+        ((self.page - 1) * self.page_size) as usize
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Paginated<T> {
+    pub page: u64,
+    pub page_size: u64,
+    pub total: u64,
+    pub total_pages: u64,
+    pub items: Vec<T>,
+}
+
+impl<T> Paginated<T> {
+    pub fn new(page: &Page, total: u64, items: Vec<T>) -> Self {
+        let total_pages = if total == 0 {
+            0
+        } else {
+            (total + page.page_size - 1) / page.page_size
+        };
+        Self {
+            page: page.page,
+            page_size: page.page_size,
+            total,
+            total_pages,
+            items,
+        }
+    }
+
+    pub fn empty(page: &Page) -> Self {
+        Self::new(page, 0, vec![])
+    }
+}

crates/memdb/src/wal.rs

@@ -0,0 +1,74 @@
+use crate::error::Result;
+use serde::{Deserialize, Serialize};
+use std::io::{BufRead, Write};
+use std::path::Path;
+
+/// A single WAL operation.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "op")]
+pub enum WalOp {
+    Insert {
+        collection: String,
+        key: String,
+        value: serde_json::Value,
+    },
+    Upsert {
+        collection: String,
+        key: String,
+        value: serde_json::Value,
+    },
+    Delete {
+        collection: String,
+        key: String,
+    },
+}
+
+/// A WAL entry carrying a monotonically increasing sequence number.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WalEntry {
+    pub seq: u64,
+    /// One entry may carry multiple ops (transaction batch).
+    pub ops: Vec<WalOp>,
+}
+
+/// Write a WAL entry line to the file (no fsync — caller decides when to
+/// sync for durability).
+pub fn write_entry(file: &mut std::fs::File, entry: &WalEntry) -> Result<()> {
+    let line = serde_json::to_string(entry)?;
+    writeln!(file, "{}", line)?;
+    Ok(())
+}
+
+/// Force buffered WAL data to disk.
+pub fn sync_wal(file: &std::fs::File) -> Result<()> {
+    file.sync_data()?;
+    Ok(())
+}
+
+/// Read all valid WAL entries whose seq > `after_seq`.
+/// Corrupted lines (e.g. partial write after power loss) are skipped
+/// with a warning and do not prevent startup.
+pub fn read_after(path: &Path, after_seq: u64) -> Result<Vec<WalEntry>> {
+    if !path.exists() {
+        return Ok(vec![]);
+    }
+    let file = std::fs::File::open(path)?;
+    let reader = std::io::BufReader::new(file);
+    let mut entries = vec![];
+    for line in reader.lines() {
+        let line = line?;
+        let trimmed = line.trim();
+        if trimmed.is_empty() {
+            continue;
+        }
+        match serde_json::from_str::<WalEntry>(trimmed) {
+            Ok(entry) if entry.seq > after_seq => entries.push(entry),
+            // seq <= after_seq — already covered by snapshot, skip.
+            Ok(_) => {}
+            Err(e) => {
+                eprintln!("[wal] skipping corrupted entry: {e}");
+            }
+        }
+    }
+    Ok(entries)
+}