diff --git a/bin/host/src/fetcher/mod.rs b/bin/host/src/fetcher/mod.rs index 3449708f2..eb6e4f73a 100644 --- a/bin/host/src/fetcher/mod.rs +++ b/bin/host/src/fetcher/mod.rs @@ -16,7 +16,7 @@ pub use hint::HintType; /// The [Fetcher] struct is responsible for fetching preimages from a remote source. pub struct Fetcher where - KV: KeyValueStore, + KV: KeyValueStore + ?Sized, { /// Key-value store for preimages. kv_store: Arc>, @@ -32,7 +32,7 @@ where impl Fetcher where - KV: KeyValueStore, + KV: KeyValueStore + ?Sized, { /// Create a new [Fetcher] with the given [KeyValueStore]. pub fn new( @@ -55,9 +55,9 @@ where // Acquire a read lock on the key-value store. let kv_lock = self.kv_store.read().await; - let mut preimage = kv_lock.get(key).cloned(); + let mut preimage = kv_lock.get(key); - // Drop the read lock before beginning the loop. + // Drop the read lock before beginning the retry loop. drop(kv_lock); // Use a loop to keep retrying the prefetch as long as the key is not found @@ -66,7 +66,7 @@ where self.prefetch(hint).await?; let kv_lock = self.kv_store.read().await; - preimage = kv_lock.get(key).cloned(); + preimage = kv_lock.get(key); } preimage.ok_or_else(|| anyhow!("Preimage not found.")) diff --git a/bin/host/src/kv/disk.rs b/bin/host/src/kv/disk.rs new file mode 100644 index 000000000..9357bdb5b --- /dev/null +++ b/bin/host/src/kv/disk.rs @@ -0,0 +1,35 @@ +//! Contains a concrete implementation of the [KeyValueStore] trait that stores data on disk. +//! +//! Data is stored in a directory, with a separate file for each key. The key is the filename, and +//! the value is the raw contents of the file. + +use super::KeyValueStore; +use alloy_primitives::hex; +use std::{fs, path::PathBuf}; + +/// A simple, synchronous key-value store that stores data on disk. +#[derive(Default, Clone, Debug, Eq, PartialEq)] +pub struct DiskKeyValueStore { + data_directory: PathBuf, +} + +impl DiskKeyValueStore { + /// Create a new [DiskKeyValueStore] with the given data directory. + pub fn new(data_directory: PathBuf) -> Self { + Self { data_directory } + } +} + +impl KeyValueStore for DiskKeyValueStore { + fn get(&self, key: alloy_primitives::B256) -> Option> { + let path = self.data_directory.join(format!("{}.bin", hex::encode(key))); + fs::create_dir_all(&self.data_directory).ok()?; + fs::read(path).ok() + } + + fn set(&mut self, key: alloy_primitives::B256, value: Vec) { + let path = self.data_directory.join(format!("{}.bin", hex::encode(key))); + fs::create_dir_all(&self.data_directory).expect("Failed to create directory"); + fs::write(path, value.as_slice()).expect("Failed to write data to disk"); + } +} diff --git a/bin/host/src/kv/mem.rs b/bin/host/src/kv/mem.rs index d7f2a30b7..6459ad715 100644 --- a/bin/host/src/kv/mem.rs +++ b/bin/host/src/kv/mem.rs @@ -1,8 +1,7 @@ //! Contains a concrete implementation of the [KeyValueStore] trait that stores data in memory. -use alloy_primitives::B256; - use super::KeyValueStore; +use alloy_primitives::B256; use std::collections::HashMap; /// A simple, synchronous key-value store that stores data in memory. This is useful for testing and @@ -20,8 +19,8 @@ impl MemoryKeyValueStore { } impl KeyValueStore for MemoryKeyValueStore { - fn get(&self, key: B256) -> Option<&Vec> { - self.store.get(&key) + fn get(&self, key: B256) -> Option> { + self.store.get(&key).cloned() } fn set(&mut self, key: B256, value: Vec) { diff --git a/bin/host/src/kv/mod.rs b/bin/host/src/kv/mod.rs index babec0ee6..6264781c8 100644 --- a/bin/host/src/kv/mod.rs +++ b/bin/host/src/kv/mod.rs @@ -1,14 +1,22 @@ //! This module contains the [KeyValueStore] trait and concrete implementations of it. use alloy_primitives::B256; +use std::sync::Arc; +use tokio::sync::RwLock; mod mem; pub use mem::MemoryKeyValueStore; +mod disk; +pub use disk::DiskKeyValueStore; + +/// A type alias for a shared key-value store. +pub type SharedKeyValueStore = Arc>; + /// Describes the interface of a simple, synchronous key-value store. pub trait KeyValueStore { /// Get the value associated with the given key. - fn get(&self, key: B256) -> Option<&Vec>; + fn get(&self, key: B256) -> Option>; /// Set the value associated with the given key. fn set(&mut self, key: B256, value: Vec); diff --git a/bin/host/src/main.rs b/bin/host/src/main.rs index 2fc587101..d69aaef64 100644 --- a/bin/host/src/main.rs +++ b/bin/host/src/main.rs @@ -1,6 +1,6 @@ use crate::{ cli::{init_tracing_subscriber, HostCli}, - kv::MemoryKeyValueStore, + kv::{DiskKeyValueStore, MemoryKeyValueStore, SharedKeyValueStore}, server::PreimageServer, }; use anyhow::{anyhow, Result}; @@ -49,18 +49,21 @@ async fn start_server(cfg: HostCli) -> Result<()> { let oracle_server = OracleServer::new(preimage_pipe); let hint_reader = HintReader::new(hint_pipe); - // TODO: Optional disk store if `cli.data_dir` is set. - let mem_kv_store = Arc::new(RwLock::new(MemoryKeyValueStore::new())); + let kv_store: SharedKeyValueStore = if let Some(ref data_dir) = cfg.data_dir { + Arc::new(RwLock::new(DiskKeyValueStore::new(data_dir.clone()))) + } else { + Arc::new(RwLock::new(MemoryKeyValueStore::new())) + }; let fetcher = (!cfg.is_offline()).then(|| { let l1_provider = util::http_provider(&cfg.l1_node_address.expect("Provider must be set")); let l2_provider = util::http_provider(&cfg.l2_node_address.expect("Provider must be set")); - Arc::new(RwLock::new(Fetcher::new(mem_kv_store.clone(), l1_provider, l2_provider))) + Arc::new(RwLock::new(Fetcher::new(kv_store.clone(), l1_provider, l2_provider))) }); // Start the server and wait for it to complete. info!("Starting preimage server."); - let server = PreimageServer::new(oracle_server, hint_reader, mem_kv_store, fetcher); + let server = PreimageServer::new(oracle_server, hint_reader, kv_store, fetcher); server.start().await?; info!("Preimage server has exited."); @@ -74,17 +77,20 @@ async fn start_server_and_native_client(cfg: HostCli) -> Result<()> { let oracle_server = OracleServer::new(preimage_pipe); let hint_reader = HintReader::new(hint_pipe); - // TODO: Optional disk store if `cli.data_dir` is set. - let mem_kv_store = Arc::new(RwLock::new(MemoryKeyValueStore::new())); + let kv_store: SharedKeyValueStore = if let Some(ref data_dir) = cfg.data_dir { + Arc::new(RwLock::new(DiskKeyValueStore::new(data_dir.clone()))) + } else { + Arc::new(RwLock::new(MemoryKeyValueStore::new())) + }; let fetcher = (!cfg.is_offline()).then(|| { let l1_provider = util::http_provider(&cfg.l1_node_address.expect("Provider must be set")); let l2_provider = util::http_provider(&cfg.l2_node_address.expect("Provider must be set")); - Arc::new(RwLock::new(Fetcher::new(mem_kv_store.clone(), l1_provider, l2_provider))) + Arc::new(RwLock::new(Fetcher::new(kv_store.clone(), l1_provider, l2_provider))) }); // Create the server and start it. - let server = PreimageServer::new(oracle_server, hint_reader, mem_kv_store, fetcher); + let server = PreimageServer::new(oracle_server, hint_reader, kv_store, fetcher); let server_task = tokio::task::spawn(server.start()); // Start the client program in a separate child process. diff --git a/bin/host/src/server.rs b/bin/host/src/server.rs index 7cc3b8504..daf58fdad 100644 --- a/bin/host/src/server.rs +++ b/bin/host/src/server.rs @@ -13,7 +13,7 @@ pub struct PreimageServer where P: PreimageOracleServer, H: HintReaderServer, - KV: KeyValueStore, + KV: KeyValueStore + ?Sized, { /// The oracle server. oracle_server: P, @@ -30,7 +30,7 @@ impl PreimageServer where P: PreimageOracleServer + Send + Sync + 'static, H: HintReaderServer + Send + Sync + 'static, - KV: KeyValueStore + Send + Sync + 'static, + KV: KeyValueStore + Send + Sync + ?Sized + 'static, { /// Create a new [PreimageServer] with the given [PreimageOracleServer], /// [HintReaderServer], and [KeyValueStore]. Holds onto the file descriptors for the pipes @@ -84,7 +84,6 @@ where .await .get(key.into()) .ok_or_else(|| anyhow!("Preimage not found")) - .cloned() }) } };