1mod cache;
8#[cfg(feature = "store-crypto")]
9mod encrypted;
10#[cfg(feature = "file")]
11mod file;
12#[cfg(feature = "redb")]
13mod redb;
14#[cfg(feature = "redis")]
15mod redis;
16#[cfg(feature = "sql")]
17mod sql;
18
19use bytes::Bytes;
20use std::collections::HashMap;
21use std::collections::hash_map::DefaultHasher;
22use std::future::Future;
23use std::hash::{Hash, Hasher};
24use std::pin::Pin;
25use std::sync::Arc;
26use std::time::{Duration, Instant};
27use tokio::sync::Mutex;
28
29pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
30
31pub use cache::{Cache, CacheError};
32
33pub trait KvStore: Send + Sync + 'static {
34 fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Option<Bytes>>;
35 fn set<'a>(&'a self, key: &'a str, val: Bytes, ttl: Option<Duration>) -> BoxFuture<'a, ()>;
36 fn remove<'a>(&'a self, key: &'a str) -> BoxFuture<'a, ()>;
37 fn incr<'a>(&'a self, key: &'a str, by: i64, ttl: Option<Duration>) -> BoxFuture<'a, u64>;
39 fn clear_prefix<'a>(&'a self, prefix: &'a str) -> BoxFuture<'a, u64>;
41}
42
43pub fn namespace(store: Arc<dyn KvStore>, name: &str) -> Namespace {
44 Namespace {
45 store,
46 prefix: format!("{name}:"),
47 }
48}
49
50#[derive(Clone)]
53pub struct AppStore {
54 pub inner: Arc<dyn KvStore>,
55}
56
57impl AppStore {
58 pub fn new(store: Arc<dyn KvStore>) -> Self {
59 Self { inner: store }
60 }
61
62 pub fn memory() -> Self {
63 Self::new(Arc::new(MemoryStore::new()))
64 }
65
66 pub fn namespaced(&self, name: &str) -> Arc<dyn KvStore> {
67 Arc::new(namespace(Arc::clone(&self.inner), name))
68 }
69}
70
71#[cfg(feature = "store-crypto")]
72pub use encrypted::{encrypted, encrypted_ns, AppKey, Encrypted};
73
74#[cfg(feature = "file")]
75pub use file::{Durability, FileStore};
76
77#[cfg(feature = "redb")]
78pub use redb::RedbStore;
79
80#[cfg(feature = "redis")]
81pub use redis::RedisStore;
82
83#[cfg(feature = "sql")]
84pub use sql::SqlStore;
85
86#[derive(Clone)]
88pub struct Namespace {
89 store: Arc<dyn KvStore>,
90 prefix: String,
91}
92
93impl Namespace {
94 fn full(&self, key: &str) -> String {
95 format!("{}{key}", self.prefix)
96 }
97}
98
99impl KvStore for Namespace {
100 fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Option<Bytes>> {
101 let k = self.full(key);
102 Box::pin(async move { self.store.get(&k).await })
103 }
104
105 fn set<'a>(&'a self, key: &'a str, val: Bytes, ttl: Option<Duration>) -> BoxFuture<'a, ()> {
106 let k = self.full(key);
107 Box::pin(async move { self.store.set(&k, val, ttl).await })
108 }
109
110 fn remove<'a>(&'a self, key: &'a str) -> BoxFuture<'a, ()> {
111 let k = self.full(key);
112 Box::pin(async move { self.store.remove(&k).await })
113 }
114
115 fn incr<'a>(&'a self, key: &'a str, by: i64, ttl: Option<Duration>) -> BoxFuture<'a, u64> {
116 let k = self.full(key);
117 Box::pin(async move { self.store.incr(&k, by, ttl).await })
118 }
119
120 fn clear_prefix<'a>(&'a self, prefix: &'a str) -> BoxFuture<'a, u64> {
121 let p = self.full(prefix);
122 Box::pin(async move { self.store.clear_prefix(&p).await })
123 }
124}
125
126struct Entry {
127 val: Bytes,
128 exp: Option<Instant>,
129}
130
131type ShardMap = Arc<Mutex<HashMap<String, Entry>>>;
132
133#[derive(Clone)]
135pub struct MemoryStore {
136 shards: Arc<[ShardMap]>,
137}
138
139impl Default for MemoryStore {
140 fn default() -> Self {
141 Self::new()
142 }
143}
144
145impl MemoryStore {
146 pub fn new() -> Self {
148 let n = std::thread::available_parallelism()
149 .map(|p| p.get() * 2)
150 .unwrap_or(2)
151 .max(1);
152 Self::with_shards(n)
153 }
154
155 pub fn with_shards(n: usize) -> Self {
157 let n = n.max(1);
158 let shards: Vec<_> = (0..n)
159 .map(|_| Arc::new(Mutex::new(HashMap::new())))
160 .collect();
161 Self {
162 shards: shards.into(),
163 }
164 }
165
166 fn shard(&self, key: &str) -> &ShardMap {
167 &self.shards[shard_index(key, self.shards.len())]
168 }
169
170 fn alive(e: &Entry, now: Instant) -> bool {
171 e.exp.map(|t| t > now).unwrap_or(true)
172 }
173}
174
175fn shard_index(key: &str, n: usize) -> usize {
176 let mut h = DefaultHasher::new();
177 key.hash(&mut h);
178 (h.finish() as usize) % n
179}
180
181fn trunc_key(key: &str) -> String {
182 const MAX: usize = 120;
183 if key.len() <= MAX {
184 key.to_string()
185 } else {
186 format!("{}…", &key[..MAX - 1])
187 }
188}
189
190impl KvStore for MemoryStore {
191 fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Option<Bytes>> {
192 let shard = Arc::clone(self.shard(key));
193 let key = key.to_string();
194 Box::pin(async move {
195 let started = Instant::now();
196 let mut map = shard.lock().await;
197 let now = Instant::now();
198 let val = match map.get(&key) {
199 Some(e) if Self::alive(e, now) => Some(e.val.clone()),
200 Some(_) => {
201 map.remove(&key);
202 None
203 }
204 None => None,
205 };
206 let hit = val.is_some();
207 let n = val.as_ref().map(|b| b.len() as u64);
208 tracing::debug!(
209 target: "sova.store",
210 op = "get",
211 backend = "memory",
212 key = %trunc_key(&key),
213 hit,
214 bytes = n,
215 duration_ms = started.elapsed().as_secs_f64() * 1000.0,
216 request_id = sova_core::current_request_id().as_deref().unwrap_or(""),
217 "sova.store"
218 );
219 val
220 })
221 }
222
223 fn set<'a>(&'a self, key: &'a str, val: Bytes, ttl: Option<Duration>) -> BoxFuture<'a, ()> {
224 let shard = Arc::clone(self.shard(key));
225 let key = key.to_string();
226 Box::pin(async move {
227 let started = Instant::now();
228 let n = val.len() as u64;
229 let mut map = shard.lock().await;
230 map.insert(
231 key.clone(),
232 Entry {
233 val,
234 exp: ttl.map(|d| Instant::now() + d),
235 },
236 );
237 tracing::debug!(
238 target: "sova.store",
239 op = "set",
240 backend = "memory",
241 key = %trunc_key(&key),
242 bytes = n,
243 duration_ms = started.elapsed().as_secs_f64() * 1000.0,
244 request_id = sova_core::current_request_id().as_deref().unwrap_or(""),
245 "sova.store"
246 );
247 })
248 }
249
250 fn remove<'a>(&'a self, key: &'a str) -> BoxFuture<'a, ()> {
251 let shard = Arc::clone(self.shard(key));
252 let key = key.to_string();
253 Box::pin(async move {
254 shard.lock().await.remove(&key);
255 })
256 }
257
258 fn incr<'a>(&'a self, key: &'a str, by: i64, ttl: Option<Duration>) -> BoxFuture<'a, u64> {
259 let shard = Arc::clone(self.shard(key));
260 let key = key.to_string();
261 Box::pin(async move {
262 let mut map = shard.lock().await;
263 let now = Instant::now();
264 let cur = match map.get(&key) {
265 Some(e) if Self::alive(e, now) => {
266 let s = std::str::from_utf8(&e.val).unwrap_or("0");
267 s.parse::<i64>().unwrap_or(0)
268 }
269 _ => 0,
270 };
271 let next = (cur + by).max(0) as u64;
272 map.insert(
273 key,
274 Entry {
275 val: Bytes::from(next.to_string()),
276 exp: ttl.map(|d| now + d),
277 },
278 );
279 next
280 })
281 }
282
283 fn clear_prefix<'a>(&'a self, prefix: &'a str) -> BoxFuture<'a, u64> {
284 let prefix = prefix.to_string();
285 let shards: Vec<_> = self.shards.iter().cloned().collect();
286 Box::pin(async move {
287 let mut total = 0u64;
288 for shard in shards {
289 let mut map = shard.lock().await;
290 let keys: Vec<_> = map
291 .keys()
292 .filter(|k| k.starts_with(&prefix))
293 .cloned()
294 .collect();
295 total += keys.len() as u64;
296 for k in keys {
297 map.remove(&k);
298 }
299 }
300 total
301 })
302 }
303}
304
305pub mod conformance {
307 use super::*;
308 use std::sync::Arc;
309
310 pub async fn run(store: Arc<dyn KvStore>) {
311 get_set_ttl(store.clone()).await;
312 namespace_isolation(store.clone()).await;
313 incr_atomic(store.clone()).await;
314 clear_prefix_scoped(store).await;
315 }
316
317 async fn get_set_ttl(store: Arc<dyn KvStore>) {
318 store
319 .set("a", Bytes::from_static(b"1"), Some(Duration::from_millis(50)))
320 .await;
321 assert_eq!(store.get("a").await.as_deref(), Some(b"1".as_slice()));
322 tokio::time::sleep(Duration::from_millis(80)).await;
323 assert!(store.get("a").await.is_none());
324 }
325
326 async fn namespace_isolation(store: Arc<dyn KvStore>) {
327 let a = namespace(store.clone(), "a");
328 let b = namespace(store.clone(), "b");
329 a.set("k", Bytes::from_static(b"A"), None).await;
330 b.set("k", Bytes::from_static(b"B"), None).await;
331 assert_eq!(a.get("k").await.as_deref(), Some(b"A".as_slice()));
332 assert_eq!(b.get("k").await.as_deref(), Some(b"B".as_slice()));
333 a.clear_prefix("").await;
334 assert!(a.get("k").await.is_none());
335 assert_eq!(b.get("k").await.as_deref(), Some(b"B".as_slice()));
336 }
337
338 async fn incr_atomic(store: Arc<dyn KvStore>) {
339 store.remove("c").await;
340 let mut handles = Vec::new();
341 for _ in 0..50 {
342 let s = store.clone();
343 handles.push(tokio::spawn(async move {
344 s.incr("c", 1, None).await;
345 }));
346 }
347 for h in handles {
348 h.await.unwrap();
349 }
350 assert_eq!(store.get("c").await.unwrap().as_ref(), b"50");
351 }
352
353 async fn clear_prefix_scoped(store: Arc<dyn KvStore>) {
354 store.set("p:1", Bytes::from_static(b"x"), None).await;
355 store.set("p:2", Bytes::from_static(b"y"), None).await;
356 store.set("q:1", Bytes::from_static(b"z"), None).await;
357 assert_eq!(store.clear_prefix("p:").await, 2);
358 assert!(store.get("p:1").await.is_none());
359 assert_eq!(store.get("q:1").await.as_deref(), Some(b"z".as_slice()));
360 }
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
368 async fn memory_conformance() {
369 conformance::run(Arc::new(MemoryStore::new())).await;
370 }
371
372 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
373 async fn sharded_clear_prefix_across_shards() {
374 let store = MemoryStore::with_shards(4);
375 for i in 0..32 {
376 store
377 .set(
378 &format!("shard:{i}"),
379 Bytes::from_static(b"v"),
380 None,
381 )
382 .await;
383 }
384 store
385 .set("other:1", Bytes::from_static(b"z"), None)
386 .await;
387 assert_eq!(store.clear_prefix("shard:").await, 32);
388 assert!(store.get("shard:0").await.is_none());
389 assert_eq!(
390 store.get("other:1").await.as_deref(),
391 Some(b"z".as_slice())
392 );
393 }
394
395 #[tokio::test]
396 async fn cache_remember_memory() {
397 let store = AppStore::memory();
398 let cache = store.cache();
399 let v = cache
400 .remember("k", None, || async { Ok::<_, CacheError>(42u32) })
401 .await
402 .unwrap();
403 assert_eq!(v, 42);
404 let hit = cache.get_json::<u32>("k").await;
405 assert_eq!(hit, Some(42));
406 }
407}