edgehdf5_memory/
wal.rs

1//! Write-Ahead Log (WAL) for edgehdf5 agent memory.
2//!
3//! Binary WAL format alongside the main .h5 file enables fast append-only
4//! writes without rewriting the entire HDF5 file on every save.
5
6use std::fs::{File, OpenOptions};
7use std::io::{Read, Seek, SeekFrom, Write};
8use std::path::{Path, PathBuf};
9
10use crate::MemoryError;
11
12const WAL_MAGIC: [u8; 4] = [0x45, 0x48, 0x57, 0x4C]; // "EHWL"
13const WAL_VERSION: u8 = 1;
14
15#[repr(u8)]
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum WalEntryType {
18    Save = 0x01,
19    Tombstone = 0x02,
20    ActivationUpdate = 0x03,
21}
22
23impl WalEntryType {
24    fn from_u8(v: u8) -> Option<Self> {
25        match v {
26            0x01 => Some(Self::Save),
27            0x02 => Some(Self::Tombstone),
28            0x03 => Some(Self::ActivationUpdate),
29            _ => None,
30        }
31    }
32}
33
34#[derive(Debug, Clone)]
35pub struct WalEntry {
36    pub entry_type: WalEntryType,
37    pub timestamp: f64,
38    pub chunk: String,
39    pub embedding: Vec<f32>,
40    pub source_channel: String,
41    pub session_id: String,
42    pub tags: String,
43    /// For tombstone entries: the index of the entry to delete.
44    pub tombstone_index: Option<usize>,
45}
46
47#[derive(Debug)]
48pub struct WalFile {
49    path: PathBuf,
50    file: Option<File>,
51    entry_count: u32,
52}
53
54impl WalFile {
55    /// Open or create a WAL file. If it exists, read the header and entry count.
56    pub fn open(path: &Path) -> Result<Self, MemoryError> {
57        if path.exists() {
58            // Read existing header
59            let mut f = OpenOptions::new()
60                .read(true)
61                .write(true)
62                .append(false)
63                .open(path)?;
64            let mut magic = [0u8; 4];
65            f.read_exact(&mut magic)?;
66            if magic != WAL_MAGIC {
67                return Err(MemoryError::Schema("invalid WAL magic bytes".into()));
68            }
69            let mut ver = [0u8; 1];
70            f.read_exact(&mut ver)?;
71            if ver[0] != WAL_VERSION {
72                return Err(MemoryError::Schema(format!(
73                    "unsupported WAL version {}",
74                    ver[0]
75                )));
76            }
77            let mut count_buf = [0u8; 4];
78            f.read_exact(&mut count_buf)?;
79            let entry_count = u32::from_le_bytes(count_buf);
80            // Seek to end for appending
81            f.seek(SeekFrom::End(0))?;
82            Ok(Self {
83                path: path.to_path_buf(),
84                file: Some(f),
85                entry_count,
86            })
87        } else {
88            // Create new WAL
89            let mut f = File::create(path)?;
90            f.write_all(&WAL_MAGIC)?;
91            f.write_all(&[WAL_VERSION])?;
92            f.write_all(&0u32.to_le_bytes())?;
93            f.flush()?;
94            Ok(Self {
95                path: path.to_path_buf(),
96                file: Some(f),
97                entry_count: 0,
98            })
99        }
100    }
101
102    /// Append a save entry to the WAL.
103    pub fn append_save(&mut self, entry: &WalEntry) -> Result<(), MemoryError> {
104        let f = self.file.as_mut().ok_or_else(|| {
105            MemoryError::Io(std::io::Error::new(
106                std::io::ErrorKind::Other,
107                "WAL file not open",
108            ))
109        })?;
110        // entry_type
111        f.write_all(&[WalEntryType::Save as u8])?;
112        // timestamp
113        f.write_all(&entry.timestamp.to_le_bytes())?;
114        // chunk
115        write_len_prefixed_str(f, &entry.chunk)?;
116        // embedding
117        let emb_len = entry.embedding.len() as u32;
118        f.write_all(&emb_len.to_le_bytes())?;
119        for &val in &entry.embedding {
120            f.write_all(&val.to_le_bytes())?;
121        }
122        // source_channel
123        write_len_prefixed_str(f, &entry.source_channel)?;
124        // session_id
125        write_len_prefixed_str(f, &entry.session_id)?;
126        // tags
127        write_len_prefixed_str(f, &entry.tags)?;
128        f.flush()?;
129
130        self.entry_count += 1;
131        self.write_entry_count()?;
132        Ok(())
133    }
134
135    /// Append a tombstone entry (deletion).
136    pub fn append_tombstone(
137        &mut self,
138        index: usize,
139        timestamp: f64,
140    ) -> Result<(), MemoryError> {
141        let f = self.file.as_mut().ok_or_else(|| {
142            MemoryError::Io(std::io::Error::new(
143                std::io::ErrorKind::Other,
144                "WAL file not open",
145            ))
146        })?;
147        f.write_all(&[WalEntryType::Tombstone as u8])?;
148        f.write_all(&timestamp.to_le_bytes())?;
149        f.write_all(&(index as u32).to_le_bytes())?;
150        f.flush()?;
151
152        self.entry_count += 1;
153        self.write_entry_count()?;
154        Ok(())
155    }
156
157    /// Read all entries from the WAL (for replay on open).
158    pub fn read_entries(path: &Path) -> Result<Vec<WalEntry>, MemoryError> {
159        if !path.exists() {
160            return Ok(Vec::new());
161        }
162        let mut f = File::open(path)?;
163        // Skip header
164        let mut header = [0u8; 9];
165        f.read_exact(&mut header)?;
166        if header[0..4] != WAL_MAGIC {
167            return Err(MemoryError::Schema("invalid WAL magic bytes".into()));
168        }
169        let entry_count = u32::from_le_bytes([header[5], header[6], header[7], header[8]]);
170        let mut entries = Vec::with_capacity(entry_count as usize);
171
172        for _ in 0..entry_count {
173            let mut type_buf = [0u8; 1];
174            f.read_exact(&mut type_buf)?;
175            let entry_type = WalEntryType::from_u8(type_buf[0]).ok_or_else(|| {
176                MemoryError::Schema(format!("unknown WAL entry type: {}", type_buf[0]))
177            })?;
178
179            let mut ts_buf = [0u8; 8];
180            f.read_exact(&mut ts_buf)?;
181            let timestamp = f64::from_le_bytes(ts_buf);
182
183            match entry_type {
184                WalEntryType::Save => {
185                    let chunk = read_len_prefixed_str(&mut f)?;
186                    let embedding = read_embedding(&mut f)?;
187                    let source_channel = read_len_prefixed_str(&mut f)?;
188                    let session_id = read_len_prefixed_str(&mut f)?;
189                    let tags = read_len_prefixed_str(&mut f)?;
190                    entries.push(WalEntry {
191                        entry_type,
192                        timestamp,
193                        chunk,
194                        embedding,
195                        source_channel,
196                        session_id,
197                        tags,
198                        tombstone_index: None,
199                    });
200                }
201                WalEntryType::Tombstone => {
202                    let mut idx_buf = [0u8; 4];
203                    f.read_exact(&mut idx_buf)?;
204                    let idx = u32::from_le_bytes(idx_buf) as usize;
205                    entries.push(WalEntry {
206                        entry_type,
207                        timestamp,
208                        chunk: String::new(),
209                        embedding: Vec::new(),
210                        source_channel: String::new(),
211                        session_id: String::new(),
212                        tags: String::new(),
213                        tombstone_index: Some(idx),
214                    });
215                }
216                WalEntryType::ActivationUpdate => {
217                    // Reserved for future use
218                }
219            }
220        }
221        Ok(entries)
222    }
223
224    /// Truncate the WAL (after merge into .h5).
225    pub fn truncate(&mut self) -> Result<(), MemoryError> {
226        // Close existing handle and recreate
227        self.file = None;
228        let mut f = File::create(&self.path)?;
229        f.write_all(&WAL_MAGIC)?;
230        f.write_all(&[WAL_VERSION])?;
231        f.write_all(&0u32.to_le_bytes())?;
232        f.flush()?;
233        self.file = Some(f);
234        self.entry_count = 0;
235        Ok(())
236    }
237
238    /// Number of pending entries.
239    pub fn pending_count(&self) -> u32 {
240        self.entry_count
241    }
242
243    /// Is the WAL empty?
244    pub fn is_empty(&self) -> bool {
245        self.entry_count == 0
246    }
247
248    /// Update the entry_count in the header (seek to offset 5, write u32 LE).
249    fn write_entry_count(&mut self) -> Result<(), MemoryError> {
250        let f = self.file.as_mut().ok_or_else(|| {
251            MemoryError::Io(std::io::Error::new(
252                std::io::ErrorKind::Other,
253                "WAL file not open",
254            ))
255        })?;
256        let pos = f.stream_position()?;
257        f.seek(SeekFrom::Start(5))?;
258        f.write_all(&self.entry_count.to_le_bytes())?;
259        f.flush()?;
260        f.seek(SeekFrom::Start(pos))?;
261        Ok(())
262    }
263}
264
265/// Replay WAL entries into a MemoryCache.
266pub fn replay_into_cache(entries: &[WalEntry], cache: &mut crate::cache::MemoryCache) {
267    for entry in entries {
268        match entry.entry_type {
269            WalEntryType::Save => {
270                cache.push(
271                    entry.chunk.clone(),
272                    entry.embedding.clone(),
273                    entry.source_channel.clone(),
274                    entry.timestamp,
275                    entry.session_id.clone(),
276                    entry.tags.clone(),
277                );
278            }
279            WalEntryType::Tombstone => {
280                if let Some(idx) = entry.tombstone_index {
281                    cache.mark_deleted(idx);
282                }
283            }
284            WalEntryType::ActivationUpdate => {}
285        }
286    }
287}
288
289// --- Binary helpers ---
290
291fn write_len_prefixed_str(f: &mut File, s: &str) -> Result<(), MemoryError> {
292    let bytes = s.as_bytes();
293    f.write_all(&(bytes.len() as u32).to_le_bytes())?;
294    f.write_all(bytes)?;
295    Ok(())
296}
297
298fn read_len_prefixed_str(f: &mut File) -> Result<String, MemoryError> {
299    let mut len_buf = [0u8; 4];
300    f.read_exact(&mut len_buf)?;
301    let len = u32::from_le_bytes(len_buf) as usize;
302    let mut buf = vec![0u8; len];
303    f.read_exact(&mut buf)?;
304    String::from_utf8(buf).map_err(|e| MemoryError::Schema(format!("invalid UTF-8 in WAL: {e}")))
305}
306
307fn read_embedding(f: &mut File) -> Result<Vec<f32>, MemoryError> {
308    let mut len_buf = [0u8; 4];
309    f.read_exact(&mut len_buf)?;
310    let count = u32::from_le_bytes(len_buf) as usize;
311    let mut vals = Vec::with_capacity(count);
312    for _ in 0..count {
313        let mut val_buf = [0u8; 4];
314        f.read_exact(&mut val_buf)?;
315        vals.push(f32::from_le_bytes(val_buf));
316    }
317    Ok(vals)
318}
319
320// --- Tests ---
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use tempfile::TempDir;
326
327    fn make_wal_entry(chunk: &str, embedding: &[f32]) -> WalEntry {
328        WalEntry {
329            entry_type: WalEntryType::Save,
330            timestamp: 1234567.89,
331            chunk: chunk.to_string(),
332            embedding: embedding.to_vec(),
333            source_channel: "test-channel".to_string(),
334            session_id: "sess-001".to_string(),
335            tags: "tag1,tag2".to_string(),
336            tombstone_index: None,
337        }
338    }
339
340    #[test]
341    fn test_wal_create_and_header() {
342        let dir = TempDir::new().unwrap();
343        let wal_path = dir.path().join("test.h5.wal");
344        let wal = WalFile::open(&wal_path).unwrap();
345        assert_eq!(wal.pending_count(), 0);
346        assert!(wal.is_empty());
347        drop(wal);
348
349        // Verify raw bytes on disk
350        let bytes = std::fs::read(&wal_path).unwrap();
351        assert_eq!(&bytes[0..4], &WAL_MAGIC);
352        assert_eq!(bytes[4], WAL_VERSION);
353        assert_eq!(&bytes[5..9], &0u32.to_le_bytes());
354    }
355
356    #[test]
357    fn test_wal_append_and_read() {
358        let dir = TempDir::new().unwrap();
359        let wal_path = dir.path().join("test.h5.wal");
360        {
361            let mut wal = WalFile::open(&wal_path).unwrap();
362            wal.append_save(&make_wal_entry("first", &[1.0, 2.0]))
363                .unwrap();
364            wal.append_save(&make_wal_entry("second", &[3.0, 4.0]))
365                .unwrap();
366            wal.append_save(&make_wal_entry("third", &[5.0, 6.0]))
367                .unwrap();
368            assert_eq!(wal.pending_count(), 3);
369        }
370
371        let entries = WalFile::read_entries(&wal_path).unwrap();
372        assert_eq!(entries.len(), 3);
373        assert_eq!(entries[0].chunk, "first");
374        assert_eq!(entries[0].embedding, vec![1.0, 2.0]);
375        assert_eq!(entries[1].chunk, "second");
376        assert_eq!(entries[2].chunk, "third");
377        assert_eq!(entries[2].embedding, vec![5.0, 6.0]);
378    }
379
380    #[test]
381    fn test_wal_truncate() {
382        let dir = TempDir::new().unwrap();
383        let wal_path = dir.path().join("test.h5.wal");
384        let mut wal = WalFile::open(&wal_path).unwrap();
385        for i in 0..5 {
386            wal.append_save(&make_wal_entry(&format!("entry {i}"), &[i as f32]))
387                .unwrap();
388        }
389        assert_eq!(wal.pending_count(), 5);
390
391        wal.truncate().unwrap();
392        assert_eq!(wal.pending_count(), 0);
393        assert!(wal.is_empty());
394
395        let entries = WalFile::read_entries(&wal_path).unwrap();
396        assert!(entries.is_empty());
397    }
398
399    #[test]
400    fn test_wal_append_tombstone() {
401        let dir = TempDir::new().unwrap();
402        let wal_path = dir.path().join("test.h5.wal");
403        {
404            let mut wal = WalFile::open(&wal_path).unwrap();
405            wal.append_tombstone(42, 9999.0).unwrap();
406            assert_eq!(wal.pending_count(), 1);
407        }
408
409        let entries = WalFile::read_entries(&wal_path).unwrap();
410        assert_eq!(entries.len(), 1);
411        assert_eq!(entries[0].entry_type, WalEntryType::Tombstone);
412        assert_eq!(entries[0].tombstone_index, Some(42));
413        assert!((entries[0].timestamp - 9999.0).abs() < 1e-6);
414    }
415
416    #[test]
417    fn test_wal_binary_roundtrip() {
418        let dir = TempDir::new().unwrap();
419        let wal_path = dir.path().join("test.h5.wal");
420        let unicode_chunk = "Hello 世界! 🌍 émojis & ünïcödé";
421        let embedding = vec![0.1, -0.2, 3.14159, f32::MAX, f32::MIN_POSITIVE];
422        {
423            let mut wal = WalFile::open(&wal_path).unwrap();
424            let entry = WalEntry {
425                entry_type: WalEntryType::Save,
426                timestamp: std::f64::consts::PI,
427                chunk: unicode_chunk.to_string(),
428                embedding: embedding.clone(),
429                source_channel: "channel/with/slashes".to_string(),
430                session_id: "sess-öö-123".to_string(),
431                tags: "α,β,γ".to_string(),
432                tombstone_index: None,
433            };
434            wal.append_save(&entry).unwrap();
435        }
436
437        let entries = WalFile::read_entries(&wal_path).unwrap();
438        assert_eq!(entries.len(), 1);
439        let e = &entries[0];
440        assert_eq!(e.entry_type, WalEntryType::Save);
441        assert!((e.timestamp - std::f64::consts::PI).abs() < 1e-15);
442        assert_eq!(e.chunk, unicode_chunk);
443        assert_eq!(e.embedding, embedding);
444        assert_eq!(e.source_channel, "channel/with/slashes");
445        assert_eq!(e.session_id, "sess-öö-123");
446        assert_eq!(e.tags, "α,β,γ");
447    }
448
449    #[test]
450    fn test_wal_empty_on_create() {
451        let dir = TempDir::new().unwrap();
452        let wal_path = dir.path().join("test.h5.wal");
453        let wal = WalFile::open(&wal_path).unwrap();
454        assert_eq!(wal.pending_count(), 0);
455        assert!(wal.is_empty());
456    }
457
458    // --- Integration tests (WAL + HDF5Memory) ---
459
460    use crate::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
461
462    fn make_config(dir: &TempDir) -> MemoryConfig {
463        let mut config = MemoryConfig::new(dir.path().join("test.h5"), "agent-test", 4);
464        config.wal_enabled = true;
465        config
466    }
467
468    fn make_entry(chunk: &str, embedding: &[f32]) -> MemoryEntry {
469        MemoryEntry {
470            chunk: chunk.to_string(),
471            embedding: embedding.to_vec(),
472            source_channel: "test".to_string(),
473            timestamp: 1000000.0,
474            session_id: "session-1".to_string(),
475            tags: "tag1,tag2".to_string(),
476        }
477    }
478
479    #[test]
480    fn test_save_with_wal() {
481        let dir = TempDir::new().unwrap();
482        let config = make_config(&dir);
483        let h5_path = config.path.clone();
484        let mut mem = HDF5Memory::create(config).unwrap();
485
486        // Get initial .h5 size (empty file)
487        let initial_size = std::fs::metadata(&h5_path).unwrap().len();
488
489        mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0])).unwrap();
490        mem.save(make_entry("b", &[0.0, 1.0, 0.0, 0.0])).unwrap();
491        mem.save(make_entry("c", &[0.0, 0.0, 1.0, 0.0])).unwrap();
492
493        // Cache has 3 entries
494        assert_eq!(mem.count(), 3);
495
496        // .h5 file should NOT have been updated (still initial size)
497        let after_size = std::fs::metadata(&h5_path).unwrap().len();
498        assert_eq!(initial_size, after_size, ".h5 should not grow with WAL enabled");
499
500        // .wal file should exist
501        let wal_path = h5_path.with_extension("h5.wal");
502        assert!(wal_path.exists(), ".wal file should exist");
503        assert_eq!(mem.wal_pending_count(), 3);
504    }
505
506    #[test]
507    fn test_wal_auto_merge() {
508        let dir = TempDir::new().unwrap();
509        let mut config = make_config(&dir);
510        config.wal_max_entries = 5;
511        let h5_path = config.path.clone();
512        let mut mem = HDF5Memory::create(config).unwrap();
513
514        // Save 5 entries (at threshold but not over)
515        for i in 0..5 {
516            mem.save(make_entry(&format!("entry {i}"), &[i as f32, 0.0, 0.0, 0.0]))
517                .unwrap();
518        }
519        // WAL should still have 5 pending (not yet merged, threshold is >=)
520        assert_eq!(mem.wal_pending_count(), 5);
521
522        // 6th entry triggers auto-merge (pending > wal_max_entries)
523        mem.save(make_entry("entry 5", &[5.0, 0.0, 0.0, 0.0]))
524            .unwrap();
525
526        // After auto-merge: WAL should be empty, cache still has all entries
527        assert_eq!(mem.wal_pending_count(), 0);
528        assert_eq!(mem.count(), 6);
529
530        // WAL file should be truncated (only header)
531        let wal_path = h5_path.with_extension("h5.wal");
532        let entries = WalFile::read_entries(&wal_path).unwrap();
533        assert!(entries.is_empty(), "WAL should be empty after auto-merge");
534    }
535
536    #[test]
537    fn test_wal_flush_explicit() {
538        let dir = TempDir::new().unwrap();
539        let config = make_config(&dir);
540        let h5_path = config.path.clone();
541        let mut mem = HDF5Memory::create(config).unwrap();
542
543        mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0])).unwrap();
544        mem.save(make_entry("b", &[0.0, 1.0, 0.0, 0.0])).unwrap();
545        mem.save(make_entry("c", &[0.0, 0.0, 1.0, 0.0])).unwrap();
546
547        assert_eq!(mem.wal_pending_count(), 3);
548
549        mem.flush_wal().unwrap();
550
551        // WAL should be empty after explicit flush
552        assert_eq!(mem.wal_pending_count(), 0);
553        // Cache should still have 3
554        assert_eq!(mem.count(), 3);
555        // WAL file on disk should be empty
556        let wal_path = h5_path.with_extension("h5.wal");
557        let entries = WalFile::read_entries(&wal_path).unwrap();
558        assert!(entries.is_empty());
559    }
560
561    #[test]
562    fn test_wal_replay_on_open() {
563        // Test WAL replay using read_entries + replay_into_cache directly,
564        // since the HDF5 read path is independent of WAL functionality.
565        let dir = TempDir::new().unwrap();
566        let config = make_config(&dir);
567        let h5_path = config.path.clone();
568
569        {
570            let mut mem = HDF5Memory::create(config).unwrap();
571            mem.save(make_entry("replay-a", &[1.0, 0.0, 0.0, 0.0]))
572                .unwrap();
573            mem.save(make_entry("replay-b", &[0.0, 1.0, 0.0, 0.0]))
574                .unwrap();
575            mem.save(make_entry("replay-c", &[0.0, 0.0, 1.0, 0.0]))
576                .unwrap();
577            assert_eq!(mem.wal_pending_count(), 3);
578            // Drop without flushing — WAL has 3 entries
579        }
580
581        // Verify WAL file has the entries
582        let wal_path = h5_path.with_extension("h5.wal");
583        assert!(wal_path.exists());
584        let entries = WalFile::read_entries(&wal_path).unwrap();
585        assert_eq!(entries.len(), 3);
586        assert_eq!(entries[0].chunk, "replay-a");
587        assert_eq!(entries[1].chunk, "replay-b");
588        assert_eq!(entries[2].chunk, "replay-c");
589
590        // Replay into a fresh cache (simulates what open() does)
591        let mut cache = crate::cache::MemoryCache::new(4);
592        super::replay_into_cache(&entries, &mut cache);
593        assert_eq!(cache.len(), 3);
594        assert_eq!(cache.chunks[0], "replay-a");
595        assert_eq!(cache.chunks[1], "replay-b");
596        assert_eq!(cache.chunks[2], "replay-c");
597        assert_eq!(cache.count_active(), 3);
598    }
599
600    #[test]
601    fn test_tick_session_merges_wal() {
602        let dir = TempDir::new().unwrap();
603        let config = make_config(&dir);
604        let h5_path = config.path.clone();
605        let mut mem = HDF5Memory::create(config).unwrap();
606
607        mem.save(make_entry("tick-a", &[1.0, 0.0, 0.0, 0.0]))
608            .unwrap();
609        mem.save(make_entry("tick-b", &[0.0, 1.0, 0.0, 0.0]))
610            .unwrap();
611        mem.save(make_entry("tick-c", &[0.0, 0.0, 1.0, 0.0]))
612            .unwrap();
613
614        assert_eq!(mem.wal_pending_count(), 3);
615
616        mem.tick_session().unwrap();
617
618        // WAL should be empty after tick_session merges
619        assert_eq!(mem.wal_pending_count(), 0);
620        // Cache should still have 3 entries
621        assert_eq!(mem.count(), 3);
622        // WAL file on disk should be empty
623        let wal_path = h5_path.with_extension("h5.wal");
624        let entries = WalFile::read_entries(&wal_path).unwrap();
625        assert!(entries.is_empty());
626    }
627
628    #[test]
629    fn test_wal_disabled() {
630        let dir = TempDir::new().unwrap();
631        let mut config = make_config(&dir);
632        config.wal_enabled = false;
633        let h5_path = config.path.clone();
634        let mut mem = HDF5Memory::create(config).unwrap();
635
636        mem.save(make_entry("no-wal", &[1.0, 0.0, 0.0, 0.0]))
637            .unwrap();
638
639        // With WAL disabled, save goes through flush() directly (old behavior)
640        assert_eq!(mem.count(), 1);
641        // No WAL file should exist
642        let wal_path = h5_path.with_extension("h5.wal");
643        assert!(!wal_path.exists(), "no .wal file when WAL disabled");
644        assert_eq!(mem.wal_pending_count(), 0);
645    }
646}
edgehdf5_memory/wal.rs

edgehdf5_memory/
wal.rs
