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