use crate::BatchId;
pub trait SnapshotStore<S> {
type Error: std::error::Error;
fn load(
&self,
id: &BatchId,
) -> impl std::future::Future<Output = Result<Option<S>, Self::Error>> + Send;
fn persist(
&self,
id: &BatchId,
snapshot: S,
) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send;
fn remove(
&self,
id: &BatchId,
) -> impl std::future::Future<Output = Result<bool, Self::Error>> + Send;
fn contains(
&self,
id: &BatchId,
) -> impl std::future::Future<Output = Result<bool, Self::Error>> + Send;
}
#[cfg(test)]
mod tests {
use super::*;
use alloy_primitives::B256;
use std::collections::HashMap;
use std::convert::Infallible;
use std::sync::Mutex;
#[derive(Debug, Default)]
struct InMemorySnapshotStore<S> {
entries: Mutex<HashMap<BatchId, S>>,
}
impl<S> InMemorySnapshotStore<S> {
fn new() -> Self {
Self {
entries: Mutex::new(HashMap::new()),
}
}
fn len(&self) -> usize {
self.entries.lock().expect("poisoned").len()
}
}
impl<S: Clone + Send + Sync> SnapshotStore<S> for InMemorySnapshotStore<S> {
type Error = Infallible;
async fn load(&self, id: &BatchId) -> Result<Option<S>, Self::Error> {
Ok(self.entries.lock().expect("poisoned").get(id).cloned())
}
async fn persist(&self, id: &BatchId, snapshot: S) -> Result<(), Self::Error> {
self.entries.lock().expect("poisoned").insert(*id, snapshot);
Ok(())
}
async fn remove(&self, id: &BatchId) -> Result<bool, Self::Error> {
Ok(self.entries.lock().expect("poisoned").remove(id).is_some())
}
async fn contains(&self, id: &BatchId) -> Result<bool, Self::Error> {
Ok(self.entries.lock().expect("poisoned").contains_key(id))
}
}
fn id(byte: u8) -> BatchId {
B256::repeat_byte(byte)
}
#[tokio::test]
async fn load_misses_on_cold_store() {
let store: InMemorySnapshotStore<u64> = InMemorySnapshotStore::new();
assert_eq!(store.load(&id(1)).await.unwrap(), None);
assert!(!store.contains(&id(1)).await.unwrap());
}
#[tokio::test]
async fn persist_then_load_round_trips() {
let store = InMemorySnapshotStore::new();
store.persist(&id(2), 42u64).await.unwrap();
assert!(store.contains(&id(2)).await.unwrap());
assert_eq!(store.load(&id(2)).await.unwrap(), Some(42));
assert_eq!(store.load(&id(3)).await.unwrap(), None);
}
#[tokio::test]
async fn persist_overwrites_existing_entry() {
let store = InMemorySnapshotStore::new();
store.persist(&id(4), 1u64).await.unwrap();
store.persist(&id(4), 2u64).await.unwrap();
assert_eq!(store.load(&id(4)).await.unwrap(), Some(2));
assert_eq!(store.len(), 1);
}
#[tokio::test]
async fn remove_reports_prior_presence() {
let store = InMemorySnapshotStore::new();
store.persist(&id(5), 7u64).await.unwrap();
assert!(store.remove(&id(5)).await.unwrap());
assert_eq!(store.load(&id(5)).await.unwrap(), None);
assert!(!store.remove(&id(5)).await.unwrap());
}
#[tokio::test]
async fn entries_are_isolated_by_batch_id() {
let store = InMemorySnapshotStore::new();
store.persist(&id(6), 60u64).await.unwrap();
store.persist(&id(7), 70u64).await.unwrap();
assert_eq!(store.load(&id(6)).await.unwrap(), Some(60));
assert_eq!(store.load(&id(7)).await.unwrap(), Some(70));
assert!(store.remove(&id(6)).await.unwrap());
assert_eq!(store.load(&id(6)).await.unwrap(), None);
assert_eq!(store.load(&id(7)).await.unwrap(), Some(70));
assert_eq!(store.len(), 1);
}
}