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