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