Skip to main content

xenith_core/
store.rs

1use crate::{Result, StateKey, StateValue, XenithError};
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5
6/// Optional on-chain location metadata for a [`StateKey`].
7///
8/// Used by `SyncEngine` with [`crate::ReadStrategy::Quorum`] to know
9/// which EVM contract address and storage slot to read when verifying on-chain
10/// agreement across multiple chains.
11#[derive(Clone, Debug, Default)]
12pub struct KeyMetadata {
13    /// EVM contract address whose storage slot holds this state.
14    pub address: Option<[u8; 20]>,
15    /// 32-byte storage slot index within `address`.
16    pub slot: Option<[u8; 32]>,
17}
18
19/// Pluggable persistence layer for synced state.
20///
21/// Implementations must be `Send + Sync`; use `Arc<dyn StateStore>` for shared access.
22///
23/// # Example
24///
25/// ```rust,no_run
26/// use xenith_core::{InMemoryStore, StateStore, StateKey};
27/// use std::sync::Arc;
28///
29/// # async fn example() {
30/// let store: Arc<dyn StateStore> = Arc::new(InMemoryStore::default());
31/// let keys = store.list_prefix("uniswap").await.unwrap();
32/// # }
33/// ```
34#[async_trait]
35pub trait StateStore: Send + Sync {
36    /// Retrieve the value for `key`, or `None` if it has not been set.
37    async fn get(&self, key: &StateKey) -> Result<Option<StateValue>>;
38
39    /// Insert or overwrite the value for `key`.
40    async fn set(&self, key: &StateKey, value: StateValue) -> Result<()>;
41
42    /// Remove the value for `key`. No-op if the key does not exist.
43    async fn delete(&self, key: &StateKey) -> Result<()>;
44
45    /// Return all keys whose string representation starts with `prefix`.
46    async fn list_prefix(&self, prefix: &str) -> Result<Vec<StateKey>>;
47
48    /// Retrieve the [`KeyMetadata`] for `key`, or `None` if not set.
49    async fn get_metadata(&self, key: &StateKey) -> Result<Option<KeyMetadata>>;
50
51    /// Store [`KeyMetadata`] for `key`.
52    async fn set_metadata(&self, key: &StateKey, meta: KeyMetadata) -> Result<()>;
53}
54
55/// In-process, heap-backed implementation of [`StateStore`].
56///
57/// Suitable for testing and single-process use. Backed by a `Mutex<HashMap>`
58/// wrapped in `Arc` so it is cheaply cloneable across tasks.
59///
60/// # Example
61///
62/// ```
63/// use xenith_core::{InMemoryStore, StateStore, StateKey, StateValue, StateVersion, ChainId};
64/// use bytes::Bytes;
65///
66/// # #[tokio::main(flavor = "current_thread")]
67/// # async fn main() {
68/// let store = InMemoryStore::default();
69/// let key = StateKey::new("proto", "pool", "0x1");
70/// let val = StateValue {
71///     data: Bytes::new(),
72///     version: StateVersion { timestamp_ms: 1, sequence: 0, source_chain: 1 },
73///     updated_at: 0,
74///     source_chain: ChainId::from(1),
75/// };
76/// store.set(&key, val.clone()).await.unwrap();
77/// assert_eq!(store.get(&key).await.unwrap(), Some(val));
78/// # }
79/// ```
80#[derive(Clone, Default)]
81pub struct InMemoryStore {
82    inner: Arc<Mutex<HashMap<String, StateValue>>>,
83    meta: Arc<Mutex<HashMap<String, KeyMetadata>>>,
84}
85
86#[async_trait]
87impl StateStore for InMemoryStore {
88    async fn get(&self, key: &StateKey) -> Result<Option<StateValue>> {
89        let map = self
90            .inner
91            .lock()
92            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
93        Ok(map.get(key.as_ref()).cloned())
94    }
95
96    async fn set(&self, key: &StateKey, value: StateValue) -> Result<()> {
97        let mut map = self
98            .inner
99            .lock()
100            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
101        map.insert(key.as_ref().to_owned(), value);
102        Ok(())
103    }
104
105    async fn delete(&self, key: &StateKey) -> Result<()> {
106        let mut map = self
107            .inner
108            .lock()
109            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
110        map.remove(key.as_ref());
111        Ok(())
112    }
113
114    async fn list_prefix(&self, prefix: &str) -> Result<Vec<StateKey>> {
115        let map = self
116            .inner
117            .lock()
118            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
119        let mut keys: Vec<StateKey> = map
120            .keys()
121            .filter(|k| k.starts_with(prefix))
122            .map(|k| StateKey::from_raw(k.clone()))
123            .collect();
124        keys.sort_by(|a, b| a.as_ref().cmp(b.as_ref()));
125        Ok(keys)
126    }
127
128    async fn get_metadata(&self, key: &StateKey) -> Result<Option<KeyMetadata>> {
129        let map = self
130            .meta
131            .lock()
132            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
133        Ok(map.get(key.as_ref()).cloned())
134    }
135
136    async fn set_metadata(&self, key: &StateKey, meta: KeyMetadata) -> Result<()> {
137        let mut map = self
138            .meta
139            .lock()
140            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
141        map.insert(key.as_ref().to_owned(), meta);
142        Ok(())
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use crate::ChainId;
150    use bytes::Bytes;
151
152    fn val(ts: u64) -> StateValue {
153        use crate::StateVersion;
154        StateValue {
155            data: Bytes::from_static(b"x"),
156            version: StateVersion {
157                timestamp_ms: ts,
158                sequence: 0,
159                source_chain: 1,
160            },
161            updated_at: 0,
162            source_chain: ChainId(1),
163        }
164    }
165
166    #[tokio::test]
167    async fn set_then_get_returns_value() {
168        let store = InMemoryStore::default();
169        let key = StateKey::new("proto", "pool", "0x1");
170        store.set(&key, val(1)).await.unwrap();
171        assert_eq!(store.get(&key).await.unwrap(), Some(val(1)));
172    }
173
174    #[tokio::test]
175    async fn get_missing_key_returns_none() {
176        let store = InMemoryStore::default();
177        let key = StateKey::new("proto", "pool", "missing");
178        assert_eq!(store.get(&key).await.unwrap(), None);
179    }
180
181    #[tokio::test]
182    async fn delete_removes_key() {
183        let store = InMemoryStore::default();
184        let key = StateKey::new("proto", "pool", "0x2");
185        store.set(&key, val(1)).await.unwrap();
186        store.delete(&key).await.unwrap();
187        assert_eq!(store.get(&key).await.unwrap(), None);
188    }
189
190    #[tokio::test]
191    async fn delete_missing_key_is_noop() {
192        let store = InMemoryStore::default();
193        let key = StateKey::new("proto", "pool", "ghost");
194        store.delete(&key).await.unwrap(); // must not error
195    }
196
197    #[tokio::test]
198    async fn list_prefix_returns_matching_keys() {
199        let store = InMemoryStore::default();
200        let a = StateKey::new("uniswap", "pool", "0xaaa");
201        let b = StateKey::new("uniswap", "pool", "0xbbb");
202        let other = StateKey::new("aave", "reserve", "0xaaa");
203        store.set(&a, val(1)).await.unwrap();
204        store.set(&b, val(2)).await.unwrap();
205        store.set(&other, val(3)).await.unwrap();
206
207        let keys = store.list_prefix("uniswap").await.unwrap();
208        assert_eq!(keys.len(), 2);
209        assert!(keys.contains(&a));
210        assert!(keys.contains(&b));
211        assert!(!keys.contains(&other));
212    }
213
214    #[tokio::test]
215    async fn list_prefix_empty_when_no_match() {
216        let store = InMemoryStore::default();
217        store
218            .set(&StateKey::new("aave", "x", "1"), val(1))
219            .await
220            .unwrap();
221        assert!(store.list_prefix("uniswap").await.unwrap().is_empty());
222    }
223}