Add DatabaseOptions to let callers specify the metadata file path (#141)

Also add a way to create a new database, without failing if it already
exists.

Author: and-cb

benches/benchmark_common.rs

@@ -60,7 +60,7 @@ pub fn get_base_database(
 
     let main_file_name_path = dir.path().join("triedb");
     let meta_file_name_path = dir.path().join("triedb.meta");
-    let db = Database::create(main_file_name_path.to_str().unwrap()).unwrap();
+    let db = Database::create_new(&main_file_name_path).unwrap();
 
     setup_database(&db, fallback_eoa_size, fallback_contract_size, fallback_storage_per_contract)
         .unwrap();

src/database.rs

@@ -8,7 +8,11 @@ use crate::{
 };
 use alloy_primitives::B256;
 use parking_lot::Mutex;
-use std::{io, path::Path};
+use std::{
+    fs::File,
+    io,
+    path::{Path, PathBuf},
+};
 
 #[derive(Debug)]
 pub struct Database {
@@ -17,6 +21,15 @@ pub struct Database {
     metrics: DatabaseMetrics,
 }
 
+#[must_use]
+#[derive(Default, Debug)]
+pub struct DatabaseOptions {
+    create: bool,
+    create_new: bool,
+    wipe: bool,
+    meta_path: Option<PathBuf>,
+}
+
 #[derive(Debug)]
 pub enum Error {
     PageError(PageError),
@@ -30,36 +43,97 @@ pub enum OpenError {
     IO(io::Error),
 }
 
-impl Database {
-    pub fn create(path: impl AsRef<Path>) -> Result<Self, OpenError> {
-        let db_file_path = path.as_ref();
+impl DatabaseOptions {
+    /// Sets the option to create a new database, or open it if it already exists.
+    ///
+    /// The semantics of this method are equivalent to [`std::fs::OpenOptions::create()`].
+    pub fn create(&mut self, create: bool) -> &mut Self {
+        self.create = create;
+        self
+    }
 
-        let mut meta_file_path = db_file_path.to_path_buf();
-        meta_file_path.as_mut_os_string().push(".meta");
+    /// Sets the option to create a new database, failing if it already exists.
+    ///
+    /// The semantics of this method are equivalent to [`std::fs::OpenOptions::create_new()`].
+    ///
+    /// If `create_new(true)` is set, then [`create()`](Self::create) and [`wipe()`](Self::wipe)
+    /// are ignored.
+    pub fn create_new(&mut self, create_new: bool) -> &mut Self {
+        self.create_new = create_new;
+        self
+    }
 
-        let meta_manager =
-            MetadataManager::open(meta_file_path).map_err(OpenError::MetadataError)?;
-        let page_manager = PageManager::options()
-            .create_new(true)
-            .open(db_file_path)
-            .map_err(OpenError::PageError)?;
+    /// Sets the option to erase the database if it gets successfully opened.
+    pub fn wipe(&mut self, wipe: bool) -> &mut Self {
+        self.wipe = wipe;
+        self
+    }
 
-        Ok(Self::new(StorageEngine::new(page_manager, meta_manager)))
+    /// Specifies the path of the metadata file.
+    ///
+    /// By default, the metadata file path is generated by appending `".meta"` to the database file
+    /// path.
+    pub fn meta_path(&mut self, meta_path: impl Into<PathBuf>) -> &mut Self {
+        self.meta_path = Some(meta_path.into());
+        self
     }
 
-    pub fn open(path: impl AsRef<Path>) -> Result<Self, OpenError> {
-        let db_file_path = path.as_ref();
+    /// Opens the database file at the given path.
+    pub fn open(&self, db_path: impl AsRef<Path>) -> Result<Database, OpenError> {
+        let db_path = db_path.as_ref();
+        let meta_path = self.meta_path.clone().unwrap_or_else(|| {
+            let mut meta_path = db_path.to_path_buf();
+            meta_path.as_mut_os_string().push(".meta");
+            meta_path
+        });
 
-        let mut meta_file_path = db_file_path.to_path_buf();
-        meta_file_path.as_mut_os_string().push(".meta");
-        let meta_manager =
-            MetadataManager::open(meta_file_path).map_err(OpenError::MetadataError)?;
+        Database::open_with_options(db_path, meta_path, self)
+    }
+}
+
+impl Database {
+    pub fn open(db_path: impl AsRef<Path>) -> Result<Self, OpenError> {
+        Self::options().open(db_path)
+    }
+
+    pub fn create_new(db_path: impl AsRef<Path>) -> Result<Self, OpenError> {
+        Self::options().create_new(true).open(db_path)
+    }
+
+    pub fn options() -> DatabaseOptions {
+        DatabaseOptions::default()
+    }
+
+    fn open_with_options(
+        db_path: impl AsRef<Path>,
+        meta_path: impl AsRef<Path>,
+        opts: &DatabaseOptions,
+    ) -> Result<Self, OpenError> {
+        let db_path = db_path.as_ref();
+        let meta_path = meta_path.as_ref();
+
+        let meta_file = File::options()
+            .read(true)
+            .write(true)
+            .create(opts.create)
+            .create_new(opts.create_new)
+            .truncate(false)
+            .open(meta_path)
+            .map_err(OpenError::IO)?;
+        let mut meta_manager =
+            MetadataManager::from_file(meta_file).map_err(OpenError::MetadataError)?;
+
+        if opts.wipe {
+            meta_manager.wipe().map_err(OpenError::IO)?
+        }
 
         let page_count = meta_manager.active_slot().page_count();
         let page_manager = PageManager::options()
-            .create(false)
+            .create(opts.create)
+            .create_new(opts.create_new)
+            .wipe(opts.wipe)
             .page_count(page_count)
-            .open(db_file_path)
+            .open(db_path)
             .map_err(OpenError::PageError)?;
 
         Ok(Self::new(StorageEngine::new(page_manager, meta_manager)))
@@ -176,13 +250,67 @@ mod tests {
     use crate::{account::Account, path::AddressPath};
     use alloy_primitives::{address, Address, U256};
     use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
+    use std::fs;
     use tempdir::TempDir;
 
+    #[test]
+    fn test_open() {
+        let tmp_dir = TempDir::new("test_db").expect("temporary dir creation failed");
+        let db_path = tmp_dir.path().join("db-file");
+        let auto_meta_path = tmp_dir.path().join("db-file.meta");
+        let custom_meta_path = tmp_dir.path().join("meta-file");
+
+        // Try to open a non-existing database
+        Database::options()
+            .open(&db_path)
+            .expect_err("opening a non-existing database should have failed");
+        assert!(!db_path.exists());
+        assert!(!auto_meta_path.exists());
+
+        // Open with create(true)
+        Database::options()
+            .create(true)
+            .open(&db_path)
+            .expect("database creation should have succeeded");
+        assert!(db_path.exists());
+        assert!(auto_meta_path.exists());
+
+        // Open again with create_new(true)
+        Database::options()
+            .create_new(true)
+            .open(&db_path)
+            .expect_err("database creation should have failed");
+        assert!(db_path.exists());
+        assert!(auto_meta_path.exists());
+
+        // Remove the files and open with create_new(true)
+        fs::remove_file(&db_path).expect("database file removal failed");
+        fs::remove_file(&auto_meta_path).expect("metadata file removal failed");
+        Database::options()
+            .create_new(true)
+            .open(&db_path)
+            .expect("database creation should have succeeded");
+        assert!(db_path.exists());
+        assert!(auto_meta_path.exists());
+
+        // Remove all files, and open with a custom metadata path
+        fs::remove_file(&db_path).expect("database file removal failed");
+        fs::remove_file(&auto_meta_path).expect("metadata file removal failed");
+        Database::options()
+            .create(true)
+            .meta_path(&custom_meta_path)
+            .open(&db_path)
+            .expect("database creation should have succeeded");
+        assert!(db_path.exists());
+        assert!(custom_meta_path.exists());
+        assert!(!auto_meta_path.exists());
+    }
+
     #[test]
     fn test_set_get_account() {
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let db = Database::create(file_path).unwrap();
+        let db = Database::create_new(file_path).unwrap();
 
         let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
 
@@ -223,7 +351,7 @@ mod tests {
         // create the database on disk. currently this will create a database with 0 pages
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let _db = Database::create(&file_path).unwrap();
+        let _db = Database::create_new(&file_path).unwrap();
 
         // WHEN: the database is opened
         let db = Database::open(&file_path).unwrap();
@@ -239,7 +367,7 @@ mod tests {
     fn test_data_persistence() {
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let db = Database::create(&file_path).unwrap();
+        let db = Database::create_new(&file_path).unwrap();
 
         let address1 = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
         let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY);
@@ -302,7 +430,7 @@ mod tests {
         // Create a new database and verify it has no pages
         let tmp_dir = TempDir::new("test_db").unwrap();
         let file_path = tmp_dir.path().join("test.db");
-        let db = Database::create(file_path).unwrap();
+        let db = Database::create_new(file_path).unwrap();
         assert_eq!(db.storage_engine.page_manager.size(), 0);
 
         // Add 1000 accounts

src/meta/mod.rs

@@ -577,6 +577,14 @@ impl MetadataManager {
         Ok(())
     }
 
+    /// Erases the metadata contents.
+    pub fn wipe(&mut self) -> io::Result<()> {
+        let size = Self::SIZE_MIN;
+        self.file.set_len(size)?;
+        self.mmap.as_mut()[size as usize..].fill(0);
+        Ok(())
+    }
+
     /// Saves the metadata to the storage device.
     pub fn sync(&self) -> io::Result<()> {
         if cfg!(not(miri)) {

src/page/manager/options.rs

@@ -61,6 +61,15 @@ impl PageManagerOptions {
         self
     }
 
+    /// Causes the file length to be set to 0 after opening it.
+    ///
+    /// Note that if `wipe(true)` is set, then setting [`page_count()`](Self::page_count) with any
+    /// number greater than `0` will cause the open to fail.
+    pub fn wipe(&mut self, wipe: bool) -> &mut Self {
+        self.open_options.truncate(wipe);
+        self
+    }
+
     /// Opens the file at `path` with the options specified by `self`.
     pub fn open(&self, path: impl AsRef<Path>) -> Result<PageManager, PageError> {
         PageManager::open_with_options(self, path)

tests/ethereum_execution_spec.rs

@@ -34,7 +34,7 @@ fn run_ethereum_execution_spec_state_tests() {
                 .as_str()
                 .replace("/", "_")[0..min(test_case_name.len(), 100)];
             let file_path = tmp_dir.path().join(database_file_name).to_str().unwrap().to_owned();
-            let test_database = Database::create(file_path.as_str()).unwrap();
+            let test_database = Database::create_new(file_path).unwrap();
 
             // will track accounts and storage that need to be deleted. this is essentially the
             // "diff" between the pre state and post state.