From ad3444bf838f1b365370d7c8c96fb28dca081eae Mon Sep 17 00:00:00 2001 From: Hemmi Shinichi Date: Mon, 28 Sep 2026 08:37:30 +0900 Subject: [PATCH] Introduce builders for the storage implementations Replace the `XxxStorageOptions`/`new_with_option(s)` construction pattern with builders following the API style of `std::thread::Builder`: settings are configured by chaining setter methods and the storage is created with `build`. - `InMemoryStorageBuilder`: `InMemoryStorage::builder().apply_discard(true).build()` - `SQLite3StorageBuilder`: `SQLite3Storage::builder(file_path).apply_discard(true).build()` - `JournalStorageBuilder`: `JournalStorage::builder(backend).apply_discard(true).build()` `InMemoryStorageOptions`, `SQLite3StorageOptions`, `JournalStorageOptions` and `new_with_option`/`new_with_options` are removed. `new` is kept for backward compatibility and now delegates to the builders. The pyo3 bindings construct storages through the builders as well. --- rustuna_core/src/storage.rs | 98 ++++++++++++++-------- rustuna_pyo3/src/storage/in_memory.rs | 16 ++-- rustuna_pyo3/src/storage/journal.rs | 13 +-- rustuna_pyo3/src/storage/sqlite3.rs | 13 +-- rustuna_storage/src/cache.rs | 8 +- rustuna_storage/src/journal/storage.rs | 104 +++++++++++++++--------- rustuna_storage/src/sqlite3.rs | 107 ++++++++++++++++--------- 7 files changed, 230 insertions(+), 129 deletions(-) diff --git a/rustuna_core/src/storage.rs b/rustuna_core/src/storage.rs index 7d595ccb..bc55c4cd 100644 --- a/rustuna_core/src/storage.rs +++ b/rustuna_core/src/storage.rs @@ -160,12 +160,6 @@ pub trait Storage: Send + Sync { fn may_omit_trials(&self) -> bool; } -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -/// Options for [`InMemoryStorage`]. -pub struct InMemoryStorageOptions { - pub apply_discard: bool, -} - /// In-memory storage implementation used by default in Rust code and tests. /// /// This implementation keeps all studies, trials, and caches in process memory. @@ -180,25 +174,45 @@ pub struct InMemoryStorage { // Supports the state-counting API required when `discard_trials` removes trials. // Storing only discarded trials' states avoids tracking state transitions and keeps this simple. discarded_state_counts: HashMap<(u32, TrialState), u32>, - option: InMemoryStorageOptions, + apply_discard: bool, } -impl InMemoryStorage { - /// Creates an empty in-memory storage. - pub fn new() -> InMemoryStorage { - InMemoryStorage { - studies: vec![], - trials: HashMap::new(), - trial_id_number_map: TrialIdNumberHashMap::new(), - study_caches: HashMap::new(), - next_study_id: 0, - next_trial_id: 0, - discarded_state_counts: HashMap::new(), - option: InMemoryStorageOptions::default(), +/// Builder for [`InMemoryStorage`], following the API style of [`std::thread::Builder`]. +/// +/// # Examples +/// +/// ``` +/// use rustuna_core::storage::InMemoryStorage; +/// +/// let storage = InMemoryStorage::builder().apply_discard(true).build(); +/// # let _ = storage; +/// ``` +#[derive(Debug)] +pub struct InMemoryStorageBuilder { + apply_discard: bool, +} +impl Default for InMemoryStorageBuilder { + fn default() -> Self { + Self::new() + } +} +impl InMemoryStorageBuilder { + /// Creates a builder with the default configuration. + pub fn new() -> Self { + Self { + apply_discard: false, } } - /// Creates an empty in-memory storage. - pub fn new_with_option(option: InMemoryStorageOptions) -> InMemoryStorage { + /// Sets whether [`Storage::discard_trials`] removes trials from this storage. + /// + /// When this is `false`, discarding is a no-op. When this is `true`, discarded trials + /// are omitted from subsequent reads. + pub fn apply_discard(self, apply_discard: bool) -> Self { + Self { apply_discard } + } + + /// Builds the storage. + pub fn build(self) -> InMemoryStorage { InMemoryStorage { studies: vec![], trials: HashMap::new(), @@ -207,9 +221,33 @@ impl InMemoryStorage { next_study_id: 0, next_trial_id: 0, discarded_state_counts: HashMap::new(), - option, + apply_discard: self.apply_discard, } } +} + +impl InMemoryStorage { + /// Returns a builder for creating a storage with an explicit configuration. + /// + /// This is the counterpart of [`std::thread::Builder`]: settings are configured by + /// chaining methods and the storage is created with [`InMemoryStorageBuilder::build`]. + /// + /// # Examples + /// + /// ``` + /// use rustuna_core::storage::InMemoryStorage; + /// + /// let storage = InMemoryStorage::builder().apply_discard(true).build(); + /// # let _ = storage; + /// ``` + pub fn builder() -> InMemoryStorageBuilder { + InMemoryStorageBuilder::new() + } + + /// Creates an empty in-memory storage. + pub fn new() -> InMemoryStorage { + Self::builder().build() + } pub fn insert_study_with_id( &mut self, @@ -622,7 +660,7 @@ impl Storage for InMemoryStorage { } fn discard_trials(&mut self, trial_ids: &[u32]) -> Result<()> { - if !self.option.apply_discard { + if !self.apply_discard { return Ok(()); } for trial_id in trial_ids { @@ -647,7 +685,7 @@ impl Storage for InMemoryStorage { } fn may_omit_trials(&self) -> bool { - self.option.apply_discard + self.apply_discard } } @@ -803,9 +841,7 @@ mod tests { #[test] fn delete_study_removes_discarded_trial_mappings() -> Result<()> { - let mut storage = InMemoryStorage::new_with_option(InMemoryStorageOptions { - apply_discard: true, - }); + let mut storage = InMemoryStorage::builder().apply_discard(true).build(); let study_id = storage .create_new_study("study", vec![Direction::Minimize])? .id; @@ -910,9 +946,7 @@ mod tests { #[test] fn discard_trials_omits_trials() -> Result<()> { - let mut storage = InMemoryStorage::new_with_option(InMemoryStorageOptions { - apply_discard: true, - }); + let mut storage = InMemoryStorage::builder().apply_discard(true).build(); let study_id = storage .create_new_study("study", vec![Direction::Minimize])? .id; @@ -933,9 +967,7 @@ mod tests { #[test] fn get_n_trials_counts_states() -> Result<()> { - let mut storage = InMemoryStorage::new_with_option(InMemoryStorageOptions { - apply_discard: true, - }); + let mut storage = InMemoryStorage::builder().apply_discard(true).build(); let study_id = storage .create_new_study("study", vec![Direction::Minimize])? .id; diff --git a/rustuna_pyo3/src/storage/in_memory.rs b/rustuna_pyo3/src/storage/in_memory.rs index 947323c4..71607785 100644 --- a/rustuna_pyo3/src/storage/in_memory.rs +++ b/rustuna_pyo3/src/storage/in_memory.rs @@ -2,7 +2,7 @@ use std::sync::{Arc, RwLock}; use pyo3::prelude::*; -use rustuna_core::storage::{InMemoryStorage, InMemoryStorageOptions, Storage}; +use rustuna_core::storage::{InMemoryStorage, Storage}; use crate::distribution::PyDistribution; use crate::storage::binding::StorageBinding; @@ -18,15 +18,17 @@ pub struct PyInMemoryStorage { impl Default for PyInMemoryStorage { fn default() -> Self { - Self::new(InMemoryStorageOptions::default()) + Self::new(false) } } impl PyInMemoryStorage { - pub fn new(option: InMemoryStorageOptions) -> Self { - let binding = StorageBinding::new(Arc::new(RwLock::new(InMemoryStorage::new_with_option( - option, - )))); + pub fn new(apply_discard: bool) -> Self { + let binding = StorageBinding::new(Arc::new(RwLock::new( + InMemoryStorage::builder() + .apply_discard(apply_discard) + .build(), + ))); PyInMemoryStorage { binding } } @@ -40,7 +42,7 @@ impl PyInMemoryStorage { #[new] #[pyo3(signature = (*, apply_discard = false))] fn py_new(apply_discard: bool) -> Self { - PyInMemoryStorage::new(InMemoryStorageOptions { apply_discard }) + PyInMemoryStorage::new(apply_discard) } fn create_new_study( diff --git a/rustuna_pyo3/src/storage/journal.rs b/rustuna_pyo3/src/storage/journal.rs index e4b32f13..c07b636d 100644 --- a/rustuna_pyo3/src/storage/journal.rs +++ b/rustuna_pyo3/src/storage/journal.rs @@ -5,7 +5,7 @@ use pyo3::prelude::*; use rustuna_core::storage::Storage; use rustuna_storage::journal::file::JournalFileBackend; -use rustuna_storage::journal::storage::{JournalStorage, JournalStorageOptions}; +use rustuna_storage::journal::storage::JournalStorage; use crate::distribution::PyDistribution; use crate::storage::binding::StorageBinding; @@ -33,11 +33,12 @@ impl PyJournalFileStorage { let backend = JournalFileBackend::new(file_path, None).map_err(|e| { PyRuntimeError::new_err(format!("Failed to create journal file: {e:?}")) })?; - let storage = JournalStorage::new_with_options( - Box::new(backend), - JournalStorageOptions { apply_discard }, - ) - .map_err(|e| PyRuntimeError::new_err(format!("Failed to create journal storage: {e:?}")))?; + let storage = JournalStorage::builder(Box::new(backend)) + .apply_discard(apply_discard) + .build() + .map_err(|e| { + PyRuntimeError::new_err(format!("Failed to create journal storage: {e:?}")) + })?; let binding = StorageBinding::new(Arc::new(RwLock::new(storage))); Ok(PyJournalFileStorage { binding }) } diff --git a/rustuna_pyo3/src/storage/sqlite3.rs b/rustuna_pyo3/src/storage/sqlite3.rs index 01dc6659..3a3ba020 100644 --- a/rustuna_pyo3/src/storage/sqlite3.rs +++ b/rustuna_pyo3/src/storage/sqlite3.rs @@ -5,7 +5,7 @@ use pyo3::prelude::*; use rustuna_core::storage::Storage; use rustuna_storage::cache::CachedStorage; -use rustuna_storage::sqlite3::{SQLite3Storage, SQLite3StorageOptions}; +use rustuna_storage::sqlite3::SQLite3Storage; use crate::distribution::PyDistribution; use crate::storage::binding::StorageBinding; @@ -30,11 +30,12 @@ impl PySQLite3Storage { #[new] #[pyo3(signature = (file_path, *, create_database = true, apply_discard = false))] fn py_new(file_path: &str, create_database: bool, apply_discard: bool) -> PyResult { - let backend = - SQLite3Storage::new_with_option(file_path, SQLite3StorageOptions { apply_discard }) - .map_err(|e| { - PyRuntimeError::new_err(format!("Failed to open the SQLite3 file: {e:?}")) - })?; + let backend = SQLite3Storage::builder(file_path) + .apply_discard(apply_discard) + .build() + .map_err(|e| { + PyRuntimeError::new_err(format!("Failed to open the SQLite3 file: {e:?}")) + })?; if create_database { backend.create_database().map_err(|e| { PyRuntimeError::new_err(format!("Failed to create the database: {e:?}")) diff --git a/rustuna_storage/src/cache.rs b/rustuna_storage/src/cache.rs index bcd0f805..9dff8caa 100644 --- a/rustuna_storage/src/cache.rs +++ b/rustuna_storage/src/cache.rs @@ -79,9 +79,11 @@ pub trait CachedStorageBackend: Send + Sync { ) -> Result<()>; /// Whether reads from this backend omit discarded trials. /// - /// This mirrors `InMemoryStorageOptions::apply_discard` and - /// `JournalStorageOptions::apply_discard`: [`Self::discard_trials`] persists the discard - /// regardless of this flag, which only decides whether reads apply it. + /// This mirrors the `apply_discard` setting of the storage builders (e.g. + /// [`InMemoryStorageBuilder`](rustuna_core::storage::InMemoryStorageBuilder) and + /// [`JournalStorageBuilder`](crate::journal::storage::JournalStorageBuilder)): + /// [`Self::discard_trials`] persists the discard regardless of this flag, which only + /// decides whether reads apply it. fn apply_discard(&self) -> bool { false } diff --git a/rustuna_storage/src/journal/storage.rs b/rustuna_storage/src/journal/storage.rs index b0fa486d..4b7a51a7 100644 --- a/rustuna_storage/src/journal/storage.rs +++ b/rustuna_storage/src/journal/storage.rs @@ -19,14 +19,6 @@ use rustuna_core::{Error, ErrorKind, Result}; use super::{JournalBackend, JournalLog, JournalOperation}; use crate::datetime::{journal_datetime_to_naive_utc, naive_utc_to_aware_utc, now_aware_utc}; -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -/// Options for [`JournalStorage`]. -pub struct JournalStorageOptions { - /// If `true`, discarded trials are omitted when replaying the journal. - /// Discard logs are written regardless of this option. - pub apply_discard: bool, -} - /// Storage implementation backed by an append-only journal log. /// /// Similar in spirit to Optuna's `JournalStorage`, this storage writes every state-changing @@ -36,25 +28,72 @@ pub struct JournalStorage { replay: JournalReplayState, } -impl JournalStorage { - /// Creates a journal storage and synchronizes its in-memory replay state from the backend. - pub fn new(backend: Box) -> Result { - Self::new_with_options(backend, JournalStorageOptions::default()) +/// Builder for [`JournalStorage`], following the API style of [`std::thread::Builder`]. +/// +/// # Examples +/// +/// ```no_run +/// use rustuna_storage::journal::file::JournalFileBackend; +/// use rustuna_storage::journal::storage::JournalStorage; +/// +/// # fn main() -> rustuna_core::Result<()> { +/// let path = std::env::temp_dir().join("journal.log"); +/// let backend = JournalFileBackend::new(&path, None)?; +/// let storage = JournalStorage::builder(Box::new(backend)) +/// .apply_discard(true) +/// .build()?; +/// # let _ = storage; +/// # Ok(()) +/// # } +/// ``` +pub struct JournalStorageBuilder { + backend: Box, + apply_discard: bool, +} +impl JournalStorageBuilder { + /// Creates a builder backed by `backend`. + pub fn new(backend: Box) -> Self { + Self { + backend, + apply_discard: false, + } } - /// Creates a journal storage and synchronizes its in-memory replay state from the backend. - pub fn new_with_options( - backend: Box, - options: JournalStorageOptions, - ) -> Result { + /// Sets whether discarded trials are omitted when replaying the journal. + /// + /// Discard logs are written regardless of this setting. + pub fn apply_discard(self, apply_discard: bool) -> Self { + Self { + apply_discard, + ..self + } + } + + /// Builds the storage and synchronizes its in-memory replay state from the backend. + pub fn build(self) -> Result { let worker_id_prefix = format!("{}-{}-", unique_prefix(), std::process::id()); let mut storage = JournalStorage { - backend, - replay: JournalReplayState::new(worker_id_prefix, options.apply_discard), + backend: self.backend, + replay: JournalReplayState::new(worker_id_prefix, self.apply_discard), }; storage.sync_with_backend()?; Ok(storage) } +} + +impl JournalStorage { + /// Returns a builder for creating a storage with an explicit configuration. + /// + /// This is the counterpart of [`std::thread::Builder`]: settings are configured by + /// chaining methods and the storage is created with [`JournalStorageBuilder::build`]. + pub fn builder(backend: Box) -> JournalStorageBuilder { + JournalStorageBuilder::new(backend) + } + + /// Creates a journal storage and synchronizes its in-memory replay state from the backend. + pub fn new(backend: Box) -> Result { + Self::builder(backend).build() + } fn worker_id(&self) -> String { format!( @@ -2330,12 +2369,9 @@ mod tests { } let backend = InMemoryJournalBackend { logs: logs.clone() }; - let mut reloaded = JournalStorage::new_with_options( - Box::new(backend), - JournalStorageOptions { - apply_discard: true, - }, - )?; + let mut reloaded = JournalStorage::builder(Box::new(backend)) + .apply_discard(true) + .build()?; let trials = reloaded.get_trials(study_id)?; assert!(trials[0].is_none()); Ok(()) @@ -2450,12 +2486,9 @@ mod tests { storage.discard_trials(&[trial_id])?; let backend = InMemoryJournalBackend { logs: logs.clone() }; - let mut storage2 = JournalStorage::new_with_options( - Box::new(backend), - JournalStorageOptions { - apply_discard: false, - }, - )?; + let mut storage2 = JournalStorage::builder(Box::new(backend)) + .apply_discard(false) + .build()?; let trials = storage2.get_trials(study_id)?; assert_eq!(trials.len(), 1); assert_eq!(trials[0].as_ref().unwrap().id, trial_id); @@ -2467,12 +2500,9 @@ mod tests { fn get_n_trials_counts_states_including_discarded_trials() -> Result<()> { let logs = Arc::new(Mutex::new(Vec::new())); let backend = InMemoryJournalBackend { logs }; - let mut storage = JournalStorage::new_with_options( - Box::new(backend), - JournalStorageOptions { - apply_discard: true, - }, - )?; + let mut storage = JournalStorage::builder(Box::new(backend)) + .apply_discard(true) + .build()?; let study_id = storage.create_new_study("s", vec![Direction::Minimize])?.id; let running_trial_id = storage.create_new_trial(study_id)?.id; let complete_trial_id = storage.create_new_trial(study_id)?.id; diff --git a/rustuna_storage/src/sqlite3.rs b/rustuna_storage/src/sqlite3.rs index 3a92b8da..db03b6ea 100644 --- a/rustuna_storage/src/sqlite3.rs +++ b/rustuna_storage/src/sqlite3.rs @@ -15,16 +15,6 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Mutex; -/// Options for [`SQLite3Storage`]. -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -pub struct SQLite3StorageOptions { - /// If `true`, discarded trials are omitted from subsequent reads. - /// - /// As in `JournalStorageOptions`, this only gates reads: `discard_trials` marks the trials - /// in the database regardless of this option. - pub apply_discard: bool, -} - /// SQLite-backed storage backend. /// /// This backend persists studies and trials in a local SQLite database and is typically wrapped @@ -32,7 +22,7 @@ pub struct SQLite3StorageOptions { /// `rustuna_core`. pub struct SQLite3Storage { conn: Mutex, - options: SQLite3StorageOptions, + apply_discard: bool, has_discarded_at_column: AtomicBool, } @@ -45,40 +35,82 @@ const TRIALS_DISCARDED_AT_INDEX_SQL: &str = type TrialRow = (u32, u32, String, Option, Option); -impl SQLite3Storage { - /// Opens a SQLite database file. - pub fn new(file_path: &str) -> Result { - Self::new_with_option(file_path, SQLite3StorageOptions::default()) +/// Builder for [`SQLite3Storage`], following the API style of [`std::thread::Builder`]. +/// +/// # Examples +/// +/// ``` +/// use rustuna_storage::sqlite3::SQLite3Storage; +/// +/// let storage = SQLite3Storage::builder(":memory:") +/// .apply_discard(true) +/// .build() +/// .unwrap(); +/// # let _ = storage; +/// ``` +pub struct SQLite3StorageBuilder { + file_path: String, + apply_discard: bool, +} +impl SQLite3StorageBuilder { + /// Creates a builder that opens `file_path`. + pub fn new(file_path: &str) -> Self { + Self { + file_path: file_path.to_owned(), + apply_discard: false, + } } - /// Opens a SQLite database file with the given options. + /// Sets whether discarded trials are omitted from subsequent reads. /// - /// When `apply_discard` is enabled, call [`Self::validate_discard_support`] after - /// [`Self::create_database`] to reject databases whose schema predates the discard column. - pub fn new_with_option( - file_path: &str, - options: SQLite3StorageOptions, - ) -> Result { - let conn = Connection::open(file_path).map_err(|e| { + /// This only gates reads: `discard_trials` marks the trials in the database regardless + /// of this setting. When it is enabled, call [`SQLite3Storage::validate_discard_support`] + /// after [`SQLite3Storage::create_database`] to reject databases whose schema predates + /// the discard column. + pub fn apply_discard(self, apply_discard: bool) -> Self { + Self { + apply_discard, + ..self + } + } + + /// Builds the storage. + pub fn build(self) -> Result { + let conn = Connection::open(&self.file_path).map_err(|e| { Error::with_reason( ErrorKind::StorageError, - format!("Failed to open {file_path}: {e}"), + format!("Failed to open {}: {e}", self.file_path), ) })?; - let has_discarded_at_column = Self::has_discarded_at_column(&conn)?; + let has_discarded_at_column = SQLite3Storage::has_discarded_at_column(&conn)?; Ok(SQLite3Storage { conn: Mutex::new(conn), - options, + apply_discard: self.apply_discard, has_discarded_at_column: AtomicBool::new(has_discarded_at_column), }) } +} + +impl SQLite3Storage { + /// Returns a builder for creating a storage with an explicit configuration. + /// + /// This is the counterpart of [`std::thread::Builder`]: settings are configured by + /// chaining methods and the storage is created with [`SQLite3StorageBuilder::build`]. + pub fn builder(file_path: &str) -> SQLite3StorageBuilder { + SQLite3StorageBuilder::new(file_path) + } + + /// Opens a SQLite database file. + pub fn new(file_path: &str) -> Result { + Self::builder(file_path).build() + } /// Returns an error when discards were requested but the database cannot record them. /// /// [`Self::create_database`] migrates the column in, so this only fails for databases opened /// without initialization. pub fn validate_discard_support(&self) -> Result<()> { - if self.options.apply_discard && !self.has_discarded_at_column.load(Ordering::Acquire) { + if self.apply_discard && !self.has_discarded_at_column.load(Ordering::Acquire) { return Err(Error::with_reason( ErrorKind::StorageError, "apply_discard requires the Rustuna-specific `discarded_at` column on the \ @@ -239,7 +271,7 @@ impl SQLite3Storage { impl CachedStorageBackend for SQLite3Storage { fn apply_discard(&self) -> bool { - self.options.apply_discard + self.apply_discard } fn discard_trials(&mut self, trial_ids: &[u32]) -> Result<()> { @@ -1497,7 +1529,7 @@ impl CachedStorageBackend for SQLite3Storage { let select_columns = "SELECT trial_id, number, state, datetime_start, datetime_complete FROM trials"; let discard_condition = - if self.options.apply_discard && self.has_discarded_at_column.load(Ordering::Acquire) { + if self.apply_discard && self.has_discarded_at_column.load(Ordering::Acquire) { " AND discarded_at IS NULL" } else { "" @@ -2156,7 +2188,7 @@ mod tests { use rustuna_core::study::{create_study, Direction}; fn init_storage() -> Result { - init_storage_with_option(SQLite3StorageOptions::default()) + init_storage_with_option(false) } /// Reads SQLite's own idea of the current UTC time, truncated to whole seconds. @@ -2229,8 +2261,10 @@ mod tests { Ok(()) } - fn init_storage_with_option(options: SQLite3StorageOptions) -> Result { - let storage = SQLite3Storage::new_with_option(":memory:", options)?; + fn init_storage_with_option(apply_discard: bool) -> Result { + let storage = SQLite3Storage::builder(":memory:") + .apply_discard(apply_discard) + .build()?; storage.create_database()?; storage.validate_discard_support()?; Ok(storage) @@ -2304,9 +2338,7 @@ mod tests { #[test] fn discard_trials_are_omitted_by_cached_storage() -> Result<()> { - let backend = init_storage_with_option(SQLite3StorageOptions { - apply_discard: true, - })?; + let backend = init_storage_with_option(true)?; let mut storage = CachedStorage::new(Box::new(backend)); assert!(storage.may_omit_trials()); @@ -2328,8 +2360,9 @@ mod tests { } fn open_file_storage(path: &str, apply_discard: bool) -> Result { - let backend = - SQLite3Storage::new_with_option(path, SQLite3StorageOptions { apply_discard })?; + let backend = SQLite3Storage::builder(path) + .apply_discard(apply_discard) + .build()?; backend.create_database()?; backend.validate_discard_support()?; Ok(CachedStorage::new(Box::new(backend)))