diff --git a/src/storage/journal.rs b/src/storage/journal.rs index 55ad087..6833bd6 100644 --- a/src/storage/journal.rs +++ b/src/storage/journal.rs @@ -69,8 +69,9 @@ //! [`MemoryStorage`](super::MemoryStorage) instead (the default). use core::marker::PhantomData; +use core::sync::atomic::{AtomicU64, Ordering}; use std::fs::{File, OpenOptions}; -use std::io::{BufRead, BufReader, Seek, SeekFrom, Write}; +use std::io::{BufRead, BufReader, Read as _, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -112,6 +113,8 @@ pub struct JournalStorage { path: PathBuf, /// Serialise in-process writes so we only hold the file lock briefly. write_lock: Mutex<()>, + /// Byte offset of last-read position for incremental refresh. + file_offset: AtomicU64, _marker: PhantomData, } @@ -131,6 +134,7 @@ impl JournalStorage { memory: MemoryStorage::new(), path: path.as_ref().to_path_buf(), write_lock: Mutex::new(()), + file_offset: AtomicU64::new(0), _marker: PhantomData, } } @@ -146,11 +150,12 @@ impl JournalStorage { /// exists but cannot be read or parsed. pub fn open(path: impl AsRef) -> crate::Result { let path = path.as_ref().to_path_buf(); - let trials = load_trials_from_file(&path)?; + let (trials, offset) = load_trials_from_file(&path)?; Ok(Self { memory: MemoryStorage::with_trials(trials), path, write_lock: Mutex::new(()), + file_offset: AtomicU64::new(offset), _marker: PhantomData, }) } @@ -180,6 +185,11 @@ impl JournalStorage { file.sync_data() .map_err(|e| crate::Error::Storage(e.to_string()))?; + let pos = file + .stream_position() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + self.file_offset.store(pos, Ordering::SeqCst); + file.unlock() .map_err(|e| crate::Error::Storage(e.to_string()))?; @@ -203,36 +213,97 @@ impl Storage for JournalStorag } fn refresh(&self) -> bool { - let Ok(loaded) = load_trials_from_file::(&self.path) else { + let Ok(file) = File::open(&self.path) else { return false; }; - let mut guard = self.memory.trials_arc().write(); - if loaded.len() > guard.len() { - if let Some(max_id) = loaded.iter().map(|t| t.id).max() { - self.memory.bump_next_id(max_id + 1); - } - *guard = loaded; - true - } else { - false + + if file.lock_shared().is_err() { + return false; } + + let offset = self.file_offset.load(Ordering::SeqCst); + + let file_size = if let Ok(m) = file.metadata() { + m.len() + } else { + let _ = file.unlock(); + return false; + }; + + if file_size <= offset { + let _ = file.unlock(); + return false; + } + + let mut buf = String::new(); + let mut handle = &file; + if handle.seek(SeekFrom::Start(offset)).is_err() { + let _ = file.unlock(); + return false; + } + if handle.read_to_string(&mut buf).is_err() { + let _ = file.unlock(); + return false; + } + + let _ = file.unlock(); + + let bytes_read = buf.len() as u64; + let mut new_trials = Vec::new(); + + for line in buf.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let trial: CompletedTrial = match serde_json::from_str(line) { + Ok(t) => t, + Err(_) => return false, + }; + if trial.validate().is_err() { + return false; + } + new_trials.push(trial); + } + + if new_trials.is_empty() { + self.file_offset + .store(offset + bytes_read, Ordering::SeqCst); + return false; + } + + let mut guard = self.memory.trials_arc().write(); + if let Some(max_id) = new_trials.iter().map(|t| t.id).max() { + self.memory.bump_next_id(max_id + 1); + } + guard.extend(new_trials); + self.file_offset + .store(offset + bytes_read, Ordering::SeqCst); + true } } -/// Read all trials from a JSONL file. Returns an empty vec if the -/// file does not exist. +/// Read all trials from a JSONL file. Returns an empty vec (and +/// offset 0) if the file does not exist. The returned `u64` is the +/// file size at the time of reading, suitable for initialising the +/// incremental-refresh offset. fn load_trials_from_file( path: &Path, -) -> crate::Result>> { +) -> crate::Result<(Vec>, u64)> { let file = match File::open(path) { Ok(f) => f, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok((Vec::new(), 0)), Err(e) => return Err(crate::Error::Storage(e.to_string())), }; file.lock_shared() .map_err(|e| crate::Error::Storage(e.to_string()))?; + let file_size = file + .metadata() + .map_err(|e| crate::Error::Storage(e.to_string()))? + .len(); + let reader = BufReader::new(&file); let mut trials = Vec::new(); @@ -251,5 +322,5 @@ fn load_trials_from_file( file.unlock() .map_err(|e| crate::Error::Storage(e.to_string()))?; - Ok(trials) + Ok((trials, file_size)) } diff --git a/tests/journal_tests.rs b/tests/journal_tests.rs index 6d5f473..4c25e0e 100644 --- a/tests/journal_tests.rs +++ b/tests/journal_tests.rs @@ -3,6 +3,8 @@ use std::collections::HashMap; use std::sync::Arc; +use std::io::Write; + use optimizer::parameter::{FloatParam, Parameter}; use optimizer::sampler::CompletedTrial; use optimizer::sampler::random::RandomSampler; @@ -317,3 +319,57 @@ fn accepts_valid_journal_with_distributions() { std::fs::remove_file(&path).ok(); } + +#[test] +fn refresh_skips_own_writes() { + let path = temp_path(); + let storage = JournalStorage::new(&path); + + for i in 0..5 { + storage.push(sample_trial(i, i as f64)); + // Our own push advanced the offset, so refresh should find nothing new. + assert!(!storage.refresh(), "refresh returned true after push {i}"); + } + + assert_eq!(storage.trials_arc().read().len(), 5); + std::fs::remove_file(&path).ok(); +} + +#[test] +fn refresh_picks_up_external_writes() { + let path = temp_path(); + let storage = JournalStorage::new(&path); + + // Push 3 trials through the storage (advances offset). + for i in 0..3 { + storage.push(sample_trial(i, i as f64)); + } + assert_eq!(storage.trials_arc().read().len(), 3); + + // Simulate an external process appending 2 more lines directly. + { + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + for i in 3..5u64 { + let trial = sample_trial(i, i as f64); + let line = serde_json::to_string(&trial).unwrap(); + writeln!(file, "{line}").unwrap(); + } + file.sync_all().unwrap(); + } + + // refresh() should pick up the 2 external trials. + assert!(storage.refresh(), "refresh should detect external writes"); + assert_eq!(storage.trials_arc().read().len(), 5); + + // A second refresh should be a no-op. + assert!( + !storage.refresh(), + "second refresh should return false (no new data)" + ); + assert_eq!(storage.trials_arc().read().len(), 5); + + std::fs::remove_file(&path).ok(); +}