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;
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    /// Atomic increment; creates the key at `by` when missing. Returns new value.
38    fn incr<'a>(&'a self, key: &'a str, by: i64, ttl: Option<Duration>) -> BoxFuture<'a, u64>;
39    /// Remove keys starting with `prefix`. Returns how many were removed.
40    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/// App-wide store handle — `app.state(AppStore::memory())` or via facade `SharedStore` plugin.
51/// Session / meta / rate-limit can reuse namespaced views of the same backend.
52#[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/// Scoped handle — prefixes all keys; [`clear_prefix`](KvStore::clear_prefix) stays inside.
87#[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/// In-process KvStore (sharded Mutex maps + TTL).
134#[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    /// Default shard count: `max(1, available_parallelism * 2)`.
147    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    /// Explicit shard count (minimum 1).
156    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
305/// Shared conformance suite for any [`KvStore`].
306pub 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}