diff --git a/src/storage/journal.rs b/src/storage/journal.rs index 6b0fb87..33c7653 100644 --- a/src/storage/journal.rs +++ b/src/storage/journal.rs @@ -115,12 +115,19 @@ impl Storage for JournalStorag self.memory.trials_arc() } + fn next_trial_id(&self) -> u64 { + self.memory.next_trial_id() + } + fn refresh(&self) -> bool { let Ok(loaded) = load_trials_from_file::(&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 { diff --git a/src/storage/memory.rs b/src/storage/memory.rs index d511faf..d9c2d29 100644 --- a/src/storage/memory.rs +++ b/src/storage/memory.rs @@ -1,3 +1,4 @@ +use core::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use parking_lot::RwLock; @@ -10,6 +11,7 @@ use crate::sampler::CompletedTrial; /// This is a thin wrapper around `Arc>>>`. pub struct MemoryStorage { trials: Arc>>>, + next_id: AtomicU64, } impl MemoryStorage { @@ -18,16 +20,24 @@ impl MemoryStorage { pub fn new() -> Self { Self { trials: Arc::new(RwLock::new(Vec::new())), + next_id: AtomicU64::new(0), } } /// Creates an in-memory store pre-populated with `trials`. #[must_use] pub fn with_trials(trials: Vec>) -> Self { + let next_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1); Self { trials: Arc::new(RwLock::new(trials)), + next_id: AtomicU64::new(next_id), } } + + /// Ensures the ID counter is at least `min_value`. + pub(crate) fn bump_next_id(&self, min_value: u64) { + self.next_id.fetch_max(min_value, Ordering::SeqCst); + } } impl Default for MemoryStorage { @@ -44,4 +54,8 @@ impl Storage for MemoryStorage { fn trials_arc(&self) -> &Arc>>> { &self.trials } + + fn next_trial_id(&self) -> u64 { + self.next_id.fetch_add(1, Ordering::SeqCst) + } } diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 65b5648..fc9202b 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -42,6 +42,12 @@ pub trait Storage: Send + Sync { /// lock for efficient, allocation-free access. fn trials_arc(&self) -> &Arc>>>; + /// Atomically returns the next unique trial ID. + /// + /// Each call increments an internal counter so that consecutive + /// calls always produce distinct IDs. + fn next_trial_id(&self) -> u64; + /// Reload from an external source (e.g. a file written by another /// process). Returns `true` if the in-memory buffer was updated. /// diff --git a/src/storage/sqlite.rs b/src/storage/sqlite.rs index b1378fe..075155d 100644 --- a/src/storage/sqlite.rs +++ b/src/storage/sqlite.rs @@ -98,6 +98,10 @@ impl Storage for SqliteStorage self.memory.trials_arc() } + fn next_trial_id(&self) -> u64 { + self.memory.next_trial_id() + } + fn refresh(&self) -> bool { let conn = self.conn.lock(); let Ok(loaded) = load_all::(&conn) else { @@ -105,6 +109,9 @@ impl Storage for SqliteStorage }; 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 { diff --git a/src/study.rs b/src/study.rs index f2cd352..c6d9507 100644 --- a/src/study.rs +++ b/src/study.rs @@ -6,7 +6,6 @@ use core::fmt; use core::future::Future; use core::marker::PhantomData; use core::ops::ControlFlow; -use core::sync::atomic::{AtomicU64, Ordering}; use core::time::Duration; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; @@ -52,8 +51,6 @@ where pruner: Arc, /// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)). storage: Arc>, - /// Counter for generating unique trial IDs. - next_trial_id: AtomicU64, /// Optional factory for creating sampler-aware trials. /// Set automatically for `Study` so that `create_trial()` and all /// optimization methods use the sampler without requiring `_with_sampler` suffixes. @@ -240,26 +237,18 @@ where let storage: Arc> = Arc::new(storage); let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner); - let next_id = storage - .trials_arc() - .read() - .iter() - .map(|t| t.id) - .max() - .map_or(0, |id| id + 1); - Self { direction, sampler, pruner, storage, - next_trial_id: AtomicU64::new(next_id), trial_factory, enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } } /// Returns the optimization direction. + #[must_use] pub fn direction(&self) -> Direction { self.direction } @@ -316,7 +305,6 @@ where sampler, pruner, storage, - next_trial_id: AtomicU64::new(0), trial_factory, enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } @@ -344,6 +332,7 @@ where } /// Returns a reference to the study's pruner. + #[must_use] pub fn pruner(&self) -> &dyn Pruner { &*self.pruner } @@ -423,7 +412,7 @@ where /// Generates the next unique trial ID. pub(crate) fn next_trial_id(&self) -> u64 { - self.next_trial_id.fetch_add(1, Ordering::SeqCst) + self.storage.next_trial_id() } /// Creates a new trial with a unique ID. @@ -448,13 +437,9 @@ where /// let trial2 = study.create_trial(); /// assert_eq!(trial2.id(), 1); /// ``` + #[must_use] pub fn create_trial(&self) -> Trial { - if self.storage.refresh() { - let trials = self.storage.trials_arc().read(); - if let Some(max_id) = trials.iter().map(|t| t.id).max() { - self.next_trial_id.fetch_max(max_id + 1, Ordering::SeqCst); - } - } + self.storage.refresh(); let id = self.next_trial_id(); let mut trial = if let Some(factory) = &self.trial_factory { @@ -565,6 +550,7 @@ where /// let value = x_val * x_val; /// study.tell(trial, Ok::<_, &str>(value)); /// ``` + #[must_use] pub fn ask(&self) -> Trial { self.create_trial() } @@ -649,6 +635,7 @@ where /// println!("Trial {} has value {:?}", completed.id, completed.value); /// } /// ``` + #[must_use] pub fn trials(&self) -> Vec> where V: Clone, @@ -675,11 +662,13 @@ where /// study.complete_trial(trial, 0.5); /// assert_eq!(study.n_trials(), 1); /// ``` + #[must_use] pub fn n_trials(&self) -> usize { self.storage.trials_arc().read().len() } /// Returns the number of pruned trials. + #[must_use] pub fn n_pruned_trials(&self) -> usize { self.storage .trials_arc() @@ -821,6 +810,7 @@ where /// Only includes completed trials (not failed or pruned). /// /// If fewer than `n` completed trials exist, returns all of them. + #[must_use] pub fn top_trials(&self, n: usize) -> Vec> where V: Clone, @@ -1987,6 +1977,7 @@ where /// /// This clones the internal trial list, so it is suitable for /// analysis and iteration but not for hot paths. + #[must_use] pub fn iter(&self) -> std::vec::IntoIter> { self.trials().into_iter() } @@ -2234,6 +2225,7 @@ impl Study { since = "0.2.0", note = "use `create_trial()` instead — it now uses the sampler automatically for Study" )] + #[must_use] pub fn create_trial_with_sampler(&self) -> Trial { self.create_trial() } @@ -2336,20 +2328,11 @@ impl Study { let storage: Arc> = Arc::new(storage); let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner); - let next_id = storage - .trials_arc() - .read() - .iter() - .map(|t| t.id) - .max() - .map_or(0, |id| id + 1); - Self { direction, sampler, pruner, storage, - next_trial_id: AtomicU64::new(next_id), trial_factory, enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } @@ -2451,20 +2434,11 @@ impl StudyBuilder { let storage: Arc> = Arc::from(storage); let trial_factory = Study::make_trial_factory(&sampler, &storage, &pruner); - let next_id = storage - .trials_arc() - .read() - .iter() - .map(|t| t.id) - .max() - .map_or(0, |id| id + 1); - Study { direction: self.direction, sampler, pruner, storage, - next_trial_id: AtomicU64::new(next_id), trial_factory, enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } @@ -2590,11 +2564,13 @@ impl Study { /// Returns an I/O error if the file cannot be created or written. pub fn save(&self, path: impl AsRef) -> std::io::Result<()> { let path = path.as_ref(); + let trials = self.trials(); + let next_trial_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1); let snapshot = StudySnapshot { version: 1, direction: self.direction, - trials: self.trials(), - next_trial_id: self.next_trial_id.load(Ordering::Relaxed), + trials, + next_trial_id, metadata: HashMap::new(), }; @@ -2660,12 +2636,12 @@ impl = serde_json::from_reader(file) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - let study = Study::new(snapshot.direction); - *study.storage.trials_arc().write() = snapshot.trials; - study - .next_trial_id - .store(snapshot.next_trial_id, Ordering::Relaxed); - Ok(study) + let storage = crate::storage::MemoryStorage::with_trials(snapshot.trials); + Ok(Self::with_sampler_and_storage( + snapshot.direction, + RandomSampler::new(), + storage, + )) } }