1use crate::{Result, StateKey, StateValue, XenithError};
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5
6#[derive(Clone, Debug, Default)]
12pub struct KeyMetadata {
13 pub address: Option<[u8; 20]>,
15 pub slot: Option<[u8; 32]>,
17}
18
19#[async_trait]
35pub trait StateStore: Send + Sync {
36 async fn get(&self, key: &StateKey) -> Result<Option<StateValue>>;
38
39 async fn set(&self, key: &StateKey, value: StateValue) -> Result<()>;
41
42 async fn delete(&self, key: &StateKey) -> Result<()>;
44
45 async fn list_prefix(&self, prefix: &str) -> Result<Vec<StateKey>>;
47
48 async fn get_metadata(&self, key: &StateKey) -> Result<Option<KeyMetadata>>;
50
51 async fn set_metadata(&self, key: &StateKey, meta: KeyMetadata) -> Result<()>;
53}
54
55#[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(); }
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}