Skip to main content

sova_store/
lib.rs

1//! Byte-oriented key-value store for Sova plugins (sessions, cache, CSRF, rate-limit).
2//!
3//! Trait is stable (memory + file + sql + redis + redb backends).
4//! Enable feature `unstable-store` for backwards-compatible feature flags.
5//! **Not in sova-core** — wire with `app.state(store.namespace("sess"))`.
6
7mod 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    /// Atomic increment; creates the key at `by` when missing. Returns new value.
37    fn incr<'a>(&'a self, key: &'a str, by: i64, ttl: Option<Duration>) -> BoxFuture<'a, u64>;
38    /// Remove keys starting with `prefix`. Returns how many were removed.
39    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/// App-wide store handle — `app.state(AppStore::memory())` or via facade `SharedStore` plugin.
50/// Session / meta / rate-limit can reuse namespaced views of the same backend.
51#[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/// Scoped handle — prefixes all keys; [`clear_prefix`](KvStore::clear_prefix) stays inside.
86#[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/// In-process KvStore (sharded Mutex maps + TTL).
133#[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    /// Default shard count: `max(1, available_parallelism * 2)`.
146    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    /// Explicit shard count (minimum 1).
155    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
304/// Shared conformance suite for any [`KvStore`].
305pub 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}