surfpool_core/storage/
mod.rs

1#[cfg(any(feature = "postgres", feature = "sqlite"))]
2mod diesel_common;
3mod fifo_map;
4mod hash_map;
5mod overlay;
6#[cfg(feature = "postgres")]
7mod postgres;
8#[cfg(feature = "sqlite")]
9mod sqlite;
10pub use hash_map::HashMap as StorageHashMap;
11pub use overlay::OverlayStorage;
12#[cfg(feature = "postgres")]
13pub use postgres::PostgresStorage;
14#[cfg(feature = "sqlite")]
15pub use sqlite::SqliteStorage;
16pub use surfpool_types::FifoMap as StorageFifoMap;
17
18use crate::error::SurfpoolError;
19
20pub fn new_kv_store<K, V>(
21    database_url: &Option<&str>,
22    table_name: &str,
23    surfnet_id: &str,
24) -> StorageResult<Box<dyn Storage<K, V>>>
25where
26    K: serde::Serialize
27        + serde::de::DeserializeOwned
28        + Send
29        + Sync
30        + 'static
31        + Clone
32        + Eq
33        + std::hash::Hash,
34    V: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static + Clone,
35{
36    new_kv_store_with_default(database_url, table_name, surfnet_id, || {
37        Box::new(StorageHashMap::new())
38    })
39}
40
41pub fn new_kv_store_with_default<K, V, F>(
42    database_url: &Option<&str>,
43    table_name: &str,
44    surfnet_id: &str,
45    default_storage_constructor: F,
46) -> StorageResult<Box<dyn Storage<K, V>>>
47where
48    K: serde::Serialize
49        + serde::de::DeserializeOwned
50        + Send
51        + Sync
52        + 'static
53        + Clone
54        + Eq
55        + std::hash::Hash,
56    V: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static + Clone,
57    F: FnOnce() -> Box<dyn Storage<K, V>>,
58{
59    match database_url {
60        Some(url) => {
61            #[cfg(feature = "postgres")]
62            if url.starts_with("postgres://") || url.starts_with("postgresql://") {
63                let storage = PostgresStorage::connect(url, table_name, surfnet_id)?;
64                Ok(Box::new(storage))
65            } else {
66                #[cfg(feature = "sqlite")]
67                {
68                    let storage = SqliteStorage::connect(url, table_name, surfnet_id)?;
69                    Ok(Box::new(storage))
70                }
71                #[cfg(not(feature = "sqlite"))]
72                {
73                    Err(StorageError::InvalidPostgresUrl(url.to_string()))
74                }
75            }
76
77            #[cfg(not(feature = "postgres"))]
78            if url.starts_with("postgres://") || url.starts_with("postgresql://") {
79                Err(StorageError::PostgresNotEnabled)
80            } else {
81                #[cfg(feature = "sqlite")]
82                {
83                    let storage = SqliteStorage::connect(
84                        database_url.unwrap_or(":memory:"),
85                        table_name,
86                        surfnet_id,
87                    )?;
88                    Ok(Box::new(storage))
89                }
90                #[cfg(not(feature = "sqlite"))]
91                {
92                    Err(StorageError::SqliteNotEnabled)
93                }
94            }
95        }
96        _ => {
97            let storage = default_storage_constructor();
98            Ok(storage)
99        }
100    }
101}
102
103#[derive(Debug, thiserror::Error)]
104pub enum StorageError {
105    #[error("Sqlite storage is not enabled in this build")]
106    SqliteNotEnabled,
107    #[error(
108        "Postgres storage is not enabled in this build. To use PostgreSQL, build surfpool from source with the `postgres` feature flag."
109    )]
110    PostgresNotEnabled,
111    #[error("Invalid Postgres database URL: {0}")]
112    InvalidPostgresUrl(String),
113    #[error("Failed to get pooled connection for '{0}' database: {1}")]
114    PooledConnectionError(String, #[source] surfpool_db::diesel::r2d2::PoolError),
115    #[error("Failed to serialize key for '{0}' database: {1}")]
116    SerializeKeyError(String, serde_json::Error),
117    #[error("Failed to serialize value for '{0}' database: {1}")]
118    SerializeValueError(String, serde_json::Error),
119    #[error("Failed to deserialize value in '{0}' database: {1}")]
120    DeserializeValueError(String, serde_json::Error),
121    #[error("Failed to acquire lock for database")]
122    LockError,
123    #[error("Query failed for table '{0}' in '{1}' database: {2}")]
124    QueryError(String, String, #[source] QueryExecuteError),
125}
126
127impl StorageError {
128    pub fn create_table(
129        table_name: &str,
130        db_type: &str,
131        e: surfpool_db::diesel::result::Error,
132    ) -> Self {
133        StorageError::QueryError(
134            table_name.to_string(),
135            db_type.to_string(),
136            QueryExecuteError::CreateTableError(e),
137        )
138    }
139    pub fn store(
140        table_name: &str,
141        db_type: &str,
142        store_key: &str,
143        e: surfpool_db::diesel::result::Error,
144    ) -> Self {
145        StorageError::QueryError(
146            table_name.to_string(),
147            db_type.to_string(),
148            QueryExecuteError::StoreError(store_key.to_string(), e),
149        )
150    }
151    pub fn get(
152        table_name: &str,
153        db_type: &str,
154        get_key: &str,
155        e: surfpool_db::diesel::result::Error,
156    ) -> Self {
157        StorageError::QueryError(
158            table_name.to_string(),
159            db_type.to_string(),
160            QueryExecuteError::GetError(get_key.to_string(), e),
161        )
162    }
163    pub fn delete(
164        table_name: &str,
165        db_type: &str,
166        delete_key: &str,
167        e: surfpool_db::diesel::result::Error,
168    ) -> Self {
169        StorageError::QueryError(
170            table_name.to_string(),
171            db_type.to_string(),
172            QueryExecuteError::DeleteError(delete_key.to_string(), e),
173        )
174    }
175    pub fn get_all_keys(
176        table_name: &str,
177        db_type: &str,
178        e: surfpool_db::diesel::result::Error,
179    ) -> Self {
180        StorageError::QueryError(
181            table_name.to_string(),
182            db_type.to_string(),
183            QueryExecuteError::GetAllKeysError(e),
184        )
185    }
186    pub fn get_all_key_value_pairs(
187        table_name: &str,
188        db_type: &str,
189        e: surfpool_db::diesel::result::Error,
190    ) -> Self {
191        StorageError::QueryError(
192            table_name.to_string(),
193            db_type.to_string(),
194            QueryExecuteError::GetAllKeyValuePairsError(e),
195        )
196    }
197    pub fn count(table_name: &str, db_type: &str, e: surfpool_db::diesel::result::Error) -> Self {
198        StorageError::QueryError(
199            table_name.to_string(),
200            db_type.to_string(),
201            QueryExecuteError::CountError(e),
202        )
203    }
204}
205
206#[derive(Debug, thiserror::Error)]
207pub enum QueryExecuteError {
208    #[error("Failed to create table: {0}")]
209    CreateTableError(#[source] surfpool_db::diesel::result::Error),
210    #[error("Failed to store value for key '{0}': {1}")]
211    StoreError(String, #[source] surfpool_db::diesel::result::Error),
212    #[error("Failed to get value for key '{0}': {1}")]
213    GetError(String, #[source] surfpool_db::diesel::result::Error),
214    #[error("Failed to delete value for key '{0}': {1}")]
215    DeleteError(String, #[source] surfpool_db::diesel::result::Error),
216    #[error("Failed to get all keys: {0}")]
217    GetAllKeysError(#[source] surfpool_db::diesel::result::Error),
218    #[error("Failed to get all key-value pairs: {0}")]
219    GetAllKeyValuePairsError(#[source] surfpool_db::diesel::result::Error),
220    #[error("Failed to count entries: {0}")]
221    CountError(#[source] surfpool_db::diesel::result::Error),
222}
223
224pub type StorageResult<T> = Result<T, StorageError>;
225
226impl From<StorageError> for jsonrpc_core::Error {
227    fn from(err: StorageError) -> Self {
228        SurfpoolError::from(err).into()
229    }
230}
231
232pub trait Storage<K, V>: Send + Sync {
233    fn store(&mut self, key: K, value: V) -> StorageResult<()>;
234    fn clear(&mut self) -> StorageResult<()>;
235    fn get(&self, key: &K) -> StorageResult<Option<V>>;
236    fn take(&mut self, key: &K) -> StorageResult<Option<V>>;
237    fn keys(&self) -> StorageResult<Vec<K>>;
238    fn into_iter(&self) -> StorageResult<Box<dyn Iterator<Item = (K, V)> + '_>>;
239    fn contains_key(&self, key: &K) -> StorageResult<bool> {
240        Ok(self.get(key)?.is_some())
241    }
242
243    /// Returns the number of entries in the storage.
244    fn count(&self) -> StorageResult<u64>;
245
246    /// Explicitly shutdown the storage, performing any cleanup like WAL checkpoint.
247    /// This should be called before the application exits to ensure data is persisted.
248    /// Default implementation does nothing.
249    fn shutdown(&self) {}
250
251    // Enable cloning of boxed trait objects
252    fn clone_box(&self) -> Box<dyn Storage<K, V>>;
253
254    /// Returns `Some` if this storage is an overlay-style wrapper (`OverlayStorage`) whose
255    /// buffered writes/deletes can be drained for atomic commit semantics. Default `None`.
256    /// Used by the atomic Jito bundle commit path to flush a sandbox SVM's overlay storages
257    /// back onto the original VM's underlying storage on bundle success.
258    fn as_overlay(&self) -> Option<&dyn OverlayLike<K, V>> {
259        None
260    }
261}
262
263/// Trait implemented by `OverlayStorage<K, V>` to expose its buffered writes/deletes so that a
264/// caller (e.g. atomic bundle commit) can drain them back onto a target storage.
265pub trait OverlayLike<K, V>: Send + Sync {
266    /// Returns the current in-memory overlay state: pending writes, pending deletes
267    /// (tombstones), and whether the base was logically cleared.
268    fn extract_overlay(&self) -> StorageResult<OverlayDelta<K, V>>;
269}
270
271/// Captured overlay state for atomic replay onto a target storage.
272pub struct OverlayDelta<K, V> {
273    pub writes: Vec<(K, V)>,
274    pub deletes: Vec<K>,
275    pub base_cleared: bool,
276}
277
278// Implement Clone for Box<dyn Storage<K, V>>
279impl<K, V> Clone for Box<dyn Storage<K, V>> {
280    fn clone(&self) -> Self {
281        self.clone_box()
282    }
283}
284
285// Separate trait for construction - this doesn't need to be dyn-compatible
286pub trait StorageConstructor<K, V>: Storage<K, V> + Clone {
287    fn connect(database_url: &str, table_name: &str, surfnet_id: &str) -> StorageResult<Self>
288    where
289        Self: Sized;
290}
291
292#[cfg(test)]
293pub mod tests {
294    use std::os::unix::fs::PermissionsExt;
295
296    use crossbeam_channel::Receiver;
297    use surfpool_types::{SimnetEvent, SvmFeatureConfig};
298    use uuid::Uuid;
299
300    use crate::surfnet::{
301        GeyserEvent,
302        svm::{SurfnetSvm, SurfnetSvmConfig},
303    };
304
305    /// Environment variable for PostgreSQL database URL used in tests
306    pub const POSTGRES_TEST_URL_ENV: &str = "SURFPOOL_TEST_POSTGRES_URL";
307
308    /// Generates a random surfnet_id
309    pub fn random_surfnet_id() -> String {
310        let uuid = Uuid::new_v4();
311        uuid.to_string()
312    }
313
314    pub enum TestType {
315        NoDb,
316        InMemorySqlite,
317        OnDiskSqlite(String),
318        /// PostgreSQL with a random surfnet_id for test isolation
319        #[cfg(feature = "postgres")]
320        Postgres {
321            url: String,
322            surfnet_id: String,
323        },
324    }
325
326    impl TestType {
327        pub fn initialize_svm(&self) -> (SurfnetSvm, Receiver<SimnetEvent>, Receiver<GeyserEvent>) {
328            self.initialize_svm_with_features(SvmFeatureConfig::default())
329        }
330
331        /// Like [`initialize_svm`], but constructs the SVM with a custom
332        /// [`SvmFeatureConfig`] applied at build time.
333        pub fn initialize_svm_with_features(
334            &self,
335            feature_config: SvmFeatureConfig,
336        ) -> (SurfnetSvm, Receiver<SimnetEvent>, Receiver<GeyserEvent>) {
337            match &self {
338                TestType::NoDb => SurfnetSvm::new(SurfnetSvmConfig {
339                    feature_config,
340                    ..SurfnetSvmConfig::default()
341                })
342                .unwrap(),
343                TestType::InMemorySqlite => SurfnetSvm::new_with_db(
344                    Some(":memory:"),
345                    SurfnetSvmConfig {
346                        surfnet_id: "0".to_string(),
347                        feature_config,
348                        ..SurfnetSvmConfig::default()
349                    },
350                )
351                .unwrap(),
352                TestType::OnDiskSqlite(db_path) => SurfnetSvm::new_with_db(
353                    Some(db_path.as_ref()),
354                    SurfnetSvmConfig {
355                        surfnet_id: "0".to_string(),
356                        feature_config,
357                        ..SurfnetSvmConfig::default()
358                    },
359                )
360                .unwrap(),
361                #[cfg(feature = "postgres")]
362                TestType::Postgres { url, surfnet_id } => SurfnetSvm::new_with_db(
363                    Some(url.as_ref()),
364                    SurfnetSvmConfig {
365                        surfnet_id: surfnet_id.clone(),
366                        feature_config,
367                        ..SurfnetSvmConfig::default()
368                    },
369                )
370                .unwrap(),
371            }
372        }
373
374        pub fn sqlite() -> Self {
375            let database_url = crate::storage::tests::create_tmp_sqlite_storage();
376            TestType::OnDiskSqlite(database_url)
377        }
378
379        pub fn no_db() -> Self {
380            TestType::NoDb
381        }
382
383        pub fn in_memory() -> Self {
384            TestType::InMemorySqlite
385        }
386
387        /// Creates a PostgreSQL test type with a random surfnet_id for test isolation.
388        /// The database URL is read from the SURFPOOL_TEST_POSTGRES_URL environment variable.
389        /// Panics if the environment variable is not set.
390        #[cfg(feature = "postgres")]
391        pub fn postgres() -> Self {
392            let url = std::env::var(POSTGRES_TEST_URL_ENV).unwrap_or_else(|_| {
393                panic!(
394                    "PostgreSQL test URL not set. Set the {} environment variable.",
395                    POSTGRES_TEST_URL_ENV
396                )
397            });
398            let surfnet_id = random_surfnet_id();
399            println!(
400                "Created PostgreSQL test connection with surfnet_id: {}",
401                surfnet_id
402            );
403            TestType::Postgres { url, surfnet_id }
404        }
405
406        /// Creates a PostgreSQL test type with a random surfnet_id for test isolation.
407        /// Returns None if the SURFPOOL_TEST_POSTGRES_URL environment variable is not set.
408        #[cfg(feature = "postgres")]
409        pub fn postgres_if_available() -> Option<Self> {
410            std::env::var(POSTGRES_TEST_URL_ENV).ok().map(|url| {
411                let surfnet_id = random_surfnet_id();
412                println!(
413                    "Created PostgreSQL test connection with surfnet_id: {}",
414                    surfnet_id
415                );
416                TestType::Postgres { url, surfnet_id }
417            })
418        }
419    }
420
421    impl Drop for TestType {
422        fn drop(&mut self) {
423            if let TestType::OnDiskSqlite(db_path) = self {
424                // Delete file at db_path when TestType goes out of scope
425                let _ = std::fs::remove_file(db_path);
426            }
427            // Note: PostgreSQL data is isolated by surfnet_id and doesn't need cleanup
428            // The random surfnet_id ensures test isolation without table cleanup
429        }
430    }
431
432    pub fn create_tmp_sqlite_storage() -> String {
433        // let temp_dir = tempfile::tempdir().expect("Failed to create temp dir for SqliteStorage");
434        let write_permissions = std::fs::Permissions::from_mode(0o600);
435        let file = tempfile::Builder::new()
436            .permissions(write_permissions)
437            .suffix(".sqlite")
438            .tempfile()
439            .expect("Failed to create temp file for SqliteStorage");
440        let database_url = file.path().to_path_buf();
441
442        // Use a simple path without creating the file beforehand
443        // Let SQLite create the database file itself
444        let database_url = database_url.to_str().unwrap().to_string();
445        println!("Created temporary Sqlite database at: {}", database_url);
446        database_url
447    }
448}
surfpool_core/storage/mod.rs

surfpool_core/storage/
mod.rs
