Skip to main content

pjson_rs/infrastructure/adapters/
generic_store.rs

1//! Generic in-memory store for thread-safe key-value storage
2//!
3//! Provides a reusable foundation for GAT repository implementations.
4//! Uses lock-free DashMap for concurrent access per infrastructure guidelines.
5//!
6//! # Concurrency Model
7//!
8//! This store uses [`DashMap`] which provides lock-free concurrent access through
9//! sharded hash maps. Each shard has its own lock, enabling high concurrency for
10//! operations on different keys.
11//!
12//! # Iteration Consistency Guarantees
13//!
14//! **DashMap provides weakly consistent iteration:**
15//!
16//! - Individual items are always consistent (no torn reads)
17//! - Items added during iteration may or may not be included
18//! - Items removed during iteration may or may not be included
19//! - Overall result represents a "fuzzy" snapshot of the store
20//!
21//! This is a fundamental trade-off for lock-free performance. For operations
22//! requiring strong consistency:
23//!
24//! - Use single-key lookups (`get`, `contains_key`) for authoritative checks
25//! - Accept eventual consistency for bulk queries (filter, iteration)
26//!
27//! The weakly consistent iteration model enables lock-free concurrent access
28//! without the overhead of MVCC or snapshot isolation, which is appropriate
29//! for in-memory session management where eventual consistency is acceptable.
30
31use dashmap::DashMap;
32use std::{hash::Hash, sync::Arc};
33
34/// Maximum pre-allocation size for filter result vectors.
35///
36/// Prevents excessive memory allocation when result_limit is very large.
37/// Actual allocation is min(result_limit, MAX_PREALLOC_SIZE).
38const MAX_PREALLOC_SIZE: usize = 1024;
39
40/// Generic thread-safe in-memory store
41///
42/// Uses `DashMap` for lock-free concurrent access with sharded hash maps.
43/// Arc wrapper enables cheap cloning for shared ownership across async tasks.
44///
45/// # Iteration Consistency
46///
47/// This store uses DashMap which provides **weakly consistent iteration**:
48/// - Individual items are always consistent (no torn reads)
49/// - Items added during iteration may or may not be included
50/// - Items removed during iteration may or may not be included
51/// - Overall result represents a "fuzzy" snapshot of the store
52///
53/// For operations requiring strong consistency:
54/// - Use single-key lookups (`get`, `contains_key`)
55/// - Accept eventual consistency for bulk queries
56///
57/// This trade-off enables lock-free concurrent access without the overhead
58/// of MVCC or snapshot isolation.
59#[derive(Debug)]
60pub struct InMemoryStore<K, V>
61where
62    K: Eq + Hash + Clone + Send + Sync,
63    V: Clone + Send + Sync,
64{
65    data: Arc<DashMap<K, V>>,
66}
67
68impl<K, V> InMemoryStore<K, V>
69where
70    K: Eq + Hash + Clone + Send + Sync,
71    V: Clone + Send + Sync,
72{
73    /// Create empty store
74    pub fn new() -> Self {
75        Self {
76            data: Arc::new(DashMap::new()),
77        }
78    }
79
80    /// Get number of entries
81    pub fn count(&self) -> usize {
82        self.data.len()
83    }
84
85    /// Remove all entries
86    pub fn clear(&self) {
87        self.data.clear();
88    }
89
90    /// Get all keys
91    ///
92    /// # Consistency
93    ///
94    /// Returns a weakly consistent snapshot. Keys added or removed during
95    /// iteration may or may not be included.
96    pub fn all_keys(&self) -> Vec<K> {
97        self.data.iter().map(|entry| entry.key().clone()).collect()
98    }
99
100    /// Get value by key
101    ///
102    /// # Consistency
103    ///
104    /// Single-key lookups are always consistent and provide the most recent
105    /// committed value for the key.
106    pub fn get(&self, key: &K) -> Option<V> {
107        self.data.get(key).map(|entry| entry.value().clone())
108    }
109
110    /// Insert or update value
111    pub fn insert(&self, key: K, value: V) -> Option<V> {
112        self.data.insert(key, value)
113    }
114
115    /// Remove entry by key
116    pub fn remove(&self, key: &K) -> Option<V> {
117        self.data.remove(key).map(|(_k, v)| v)
118    }
119
120    /// Filter values by predicate
121    ///
122    /// # Consistency
123    ///
124    /// Results are weakly consistent. Items added or removed during iteration
125    /// may or may not be included. For authoritative checks, use single-key
126    /// lookups (`get`, `contains_key`).
127    pub fn filter<P>(&self, predicate: P) -> Vec<V>
128    where
129        P: Fn(&V) -> bool,
130    {
131        self.data
132            .iter()
133            .filter(|entry| predicate(entry.value()))
134            .map(|entry| entry.value().clone())
135            .collect()
136    }
137
138    /// Filter with bounded results and scan limit
139    ///
140    /// Returns at most `result_limit` items matching predicate.
141    /// Stops iteration after scanning `scan_limit` items.
142    ///
143    /// # Consistency
144    ///
145    /// Results are weakly consistent. Items added or removed during iteration
146    /// may or may not be included. For authoritative checks, use single-key
147    /// lookups (`get`, `contains_key`).
148    ///
149    /// # Returns
150    ///
151    /// A tuple of (results, limit_reached) where:
152    /// - `results`: Vec of matching items (at most `result_limit` items)
153    /// - `limit_reached`: true if either scan_limit or result_limit was hit,
154    ///   meaning the query stopped before examining all items. Results are
155    ///   still valid but potentially incomplete.
156    ///
157    /// # Example
158    ///
159    /// ```ignore
160    /// use super::limits::{MAX_SCAN_LIMIT, MAX_RESULTS_LIMIT};
161    ///
162    /// let (results, truncated) = store.filter_limited(
163    ///     |v| v.is_active(),
164    ///     MAX_RESULTS_LIMIT,
165    ///     MAX_SCAN_LIMIT,
166    /// );
167    ///
168    /// if truncated {
169    ///     // Results may be incomplete
170    /// }
171    /// ```
172    pub fn filter_limited<P>(
173        &self,
174        predicate: P,
175        result_limit: usize,
176        scan_limit: usize,
177    ) -> (Vec<V>, bool)
178    where
179        P: Fn(&V) -> bool,
180    {
181        let mut results = Vec::with_capacity(result_limit.min(MAX_PREALLOC_SIZE));
182        let mut limit_reached = false;
183
184        for (scanned, entry) in self.data.iter().enumerate() {
185            // Check limit to ensure exactly scan_limit items are scanned
186            if scanned >= scan_limit {
187                limit_reached = true;
188                break;
189            }
190
191            if predicate(entry.value()) {
192                results.push(entry.value().clone());
193                if results.len() >= result_limit {
194                    limit_reached = true;
195                    break;
196                }
197            }
198        }
199
200        (results, limit_reached)
201    }
202
203    /// Check if key exists
204    ///
205    /// # Consistency
206    ///
207    /// Single-key lookups are always consistent and provide the most recent
208    /// committed state.
209    pub fn contains_key(&self, key: &K) -> bool {
210        self.data.contains_key(key)
211    }
212
213    /// Check if store is empty
214    pub fn is_empty(&self) -> bool {
215        self.data.is_empty()
216    }
217
218    /// Atomic read-modify-write operation
219    ///
220    /// Applies function to mutable value reference if key exists.
221    /// Returns the result of the function or None if key not found.
222    ///
223    /// # Consistency
224    ///
225    /// This operation is atomic with respect to the specific key. The function
226    /// is executed while holding the shard lock for that key, ensuring no
227    /// concurrent modifications to the same key.
228    ///
229    /// # Example
230    /// ```ignore
231    /// store.update_with(&stream_id, |stream| {
232    ///     stream.complete()
233    /// });
234    /// ```
235    pub fn update_with<F, R>(&self, key: &K, f: F) -> Option<R>
236    where
237        F: FnOnce(&mut V) -> R,
238    {
239        self.data.get_mut(key).map(|mut entry| f(entry.value_mut()))
240    }
241
242    /// Iterate over all entries
243    ///
244    /// Returns an iterator that yields references to each entry.
245    /// Useful for manual iteration with early abort.
246    ///
247    /// # Consistency
248    ///
249    /// Iteration is weakly consistent. Items added or removed during iteration
250    /// may or may not be included. This is a fundamental property of DashMap
251    /// that enables lock-free concurrent access.
252    pub fn iter(&self) -> impl Iterator<Item = dashmap::mapref::multiple::RefMulti<'_, K, V>> {
253        self.data.iter()
254    }
255}
256
257impl<K, V> Default for InMemoryStore<K, V>
258where
259    K: Eq + Hash + Clone + Send + Sync,
260    V: Clone + Send + Sync,
261{
262    fn default() -> Self {
263        Self::new()
264    }
265}
266
267impl<K, V> Clone for InMemoryStore<K, V>
268where
269    K: Eq + Hash + Clone + Send + Sync,
270    V: Clone + Send + Sync,
271{
272    fn clone(&self) -> Self {
273        Self {
274            data: Arc::clone(&self.data),
275        }
276    }
277}
278
279// Type aliases for domain-specific stores
280use crate::domain::{
281    aggregates::StreamSession,
282    entities::Stream,
283    value_objects::{SessionId, StreamId},
284};
285
286/// Session store type alias
287pub type SessionStore = InMemoryStore<SessionId, StreamSession>;
288
289/// Stream store type alias
290pub type StreamStore = InMemoryStore<StreamId, Stream>;
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn test_basic_operations() {
298        let store: InMemoryStore<String, i32> = InMemoryStore::new();
299
300        assert!(store.is_empty());
301        assert_eq!(store.count(), 0);
302
303        store.insert("a".to_string(), 1);
304        store.insert("b".to_string(), 2);
305
306        assert_eq!(store.count(), 2);
307        assert_eq!(store.get(&"a".to_string()), Some(1));
308        assert_eq!(store.get(&"c".to_string()), None);
309
310        let keys = store.all_keys();
311        assert_eq!(keys.len(), 2);
312
313        store.remove(&"a".to_string());
314        assert_eq!(store.count(), 1);
315
316        store.clear();
317        assert!(store.is_empty());
318    }
319
320    #[test]
321    fn test_filter() {
322        let store: InMemoryStore<String, i32> = InMemoryStore::new();
323
324        store.insert("a".to_string(), 1);
325        store.insert("b".to_string(), 2);
326        store.insert("c".to_string(), 3);
327
328        let evens = store.filter(|v| v % 2 == 0);
329        assert_eq!(evens, vec![2]);
330    }
331
332    #[test]
333    fn test_filter_limited_returns_at_most_limit_items() {
334        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
335
336        for i in 0..100 {
337            store.insert(i, i);
338        }
339
340        let (results, limit_reached) = store.filter_limited(|_| true, 10, 1000);
341
342        assert_eq!(results.len(), 10);
343        assert!(limit_reached);
344    }
345
346    #[test]
347    fn test_filter_limited_sets_limit_reached_when_scan_exceeded() {
348        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
349
350        for i in 0..100 {
351            store.insert(i, i);
352        }
353
354        let (results, limit_reached) = store.filter_limited(|v| *v > 1000, 100, 50);
355
356        assert!(results.is_empty());
357        assert!(limit_reached);
358    }
359
360    #[test]
361    fn test_filter_limited_sets_limit_reached_when_results_exceeded() {
362        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
363
364        for i in 0..100 {
365            store.insert(i, i);
366        }
367
368        let (results, limit_reached) = store.filter_limited(|_| true, 5, 1000);
369
370        assert_eq!(results.len(), 5);
371        assert!(limit_reached);
372    }
373
374    #[test]
375    fn test_filter_limited_empty_store() {
376        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
377
378        let (results, limit_reached) = store.filter_limited(|_| true, 10, 100);
379
380        assert!(results.is_empty());
381        assert!(!limit_reached);
382    }
383
384    #[test]
385    fn test_filter_limited_no_matches() {
386        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
387
388        for i in 0..10 {
389            store.insert(i, i);
390        }
391
392        let (results, limit_reached) = store.filter_limited(|v| *v > 100, 10, 100);
393
394        assert!(results.is_empty());
395        assert!(!limit_reached);
396    }
397
398    #[test]
399    fn test_filter_limited_partial_match_within_limits() {
400        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
401
402        for i in 0..10 {
403            store.insert(i, i);
404        }
405
406        let (results, limit_reached) = store.filter_limited(|v| v % 2 == 0, 100, 100);
407
408        assert_eq!(results.len(), 5);
409        assert!(!limit_reached);
410    }
411
412    #[test]
413    fn test_clone_shares_data() {
414        let store1: InMemoryStore<String, i32> = InMemoryStore::new();
415        store1.insert("key".to_string(), 42);
416
417        let store2 = store1.clone();
418        assert_eq!(store2.get(&"key".to_string()), Some(42));
419
420        store2.insert("another".to_string(), 100);
421        assert_eq!(store1.get(&"another".to_string()), Some(100));
422    }
423
424    #[test]
425    fn test_contains_key() {
426        let store: InMemoryStore<String, i32> = InMemoryStore::new();
427
428        assert!(!store.contains_key(&"key".to_string()));
429        store.insert("key".to_string(), 42);
430        assert!(store.contains_key(&"key".to_string()));
431    }
432
433    /// Test concurrent access from multiple threads
434    ///
435    /// Verifies DashMap's lock-free behavior under contention
436    #[test]
437    fn test_concurrent_access() {
438        use std::thread;
439
440        let store: InMemoryStore<i32, String> = InMemoryStore::new();
441        let store_clone = store.clone();
442
443        // Spawn thread to write concurrently
444        let write_handle = thread::spawn(move || {
445            for i in 0..100 {
446                store_clone.insert(i, format!("thread1-{}", i));
447            }
448        });
449
450        // Write from main thread concurrently
451        for i in 100..200 {
452            store.insert(i, format!("thread2-{}", i));
453        }
454
455        write_handle.join().unwrap();
456
457        // Verify all writes succeeded (DashMap handles concurrent writes)
458        assert_eq!(store.count(), 200);
459        assert_eq!(store.get(&50), Some("thread1-50".to_string()));
460        assert_eq!(store.get(&150), Some("thread2-150".to_string()));
461
462        // Test concurrent reads
463        let read_store = store.clone();
464        let read_handle = thread::spawn(move || {
465            for i in 0..200 {
466                read_store.get(&i); // Lock-free reads
467            }
468        });
469
470        // Read from main thread while other thread reads
471        for i in 0..200 {
472            store.get(&i);
473        }
474
475        read_handle.join().unwrap();
476    }
477
478    #[test]
479    fn test_iter() {
480        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
481
482        store.insert(1, 10);
483        store.insert(2, 20);
484        store.insert(3, 30);
485
486        let mut count = 0;
487        for entry in store.iter() {
488            assert!(entry.value() == &10 || entry.value() == &20 || entry.value() == &30);
489            count += 1;
490        }
491
492        assert_eq!(count, 3);
493    }
494
495    #[test]
496    fn test_max_prealloc_size_limits_allocation() {
497        // Verify that preallocation is bounded even with very large result_limit
498        let store: InMemoryStore<i32, i32> = InMemoryStore::new();
499        store.insert(1, 1);
500
501        // Even with huge result_limit, we only preallocate MAX_PREALLOC_SIZE
502        let (results, _) = store.filter_limited(|_| true, 1_000_000, 1_000_000);
503
504        // Should still work correctly
505        assert_eq!(results.len(), 1);
506    }
507}