rskit_stateful/
manager.rs

1//! Multiplexed keyed accumulators.
2
3use std::collections::HashMap;
4use std::hash::Hash;
5use std::sync::Arc;
6
7use parking_lot::Mutex;
8use rskit_errors::AppResult;
9
10use crate::{Accumulator, AccumulatorConfig, Store};
11
12type StoreFactory<V> = dyn Fn() -> Box<dyn Store<V>> + Send + Sync;
13
14/// Manages per-key accumulators and TTL cleanup.
15pub struct Manager<K, V>
16where
17    K: Eq + Hash + Clone,
18    V: Clone + Send + Sync + 'static,
19{
20    accumulators: Mutex<HashMap<K, Arc<Accumulator<V>>>>,
21    config: AccumulatorConfig<V>,
22    store_factory: Arc<StoreFactory<V>>,
23}
24
25impl<K, V> Manager<K, V>
26where
27    K: Eq + Hash + Clone,
28    V: Clone + Send + Sync + 'static,
29{
30    /// Create a new keyed manager.
31    #[must_use]
32    pub fn new(
33        config: AccumulatorConfig<V>,
34        store_factory: impl Fn() -> Box<dyn Store<V>> + Send + Sync + 'static,
35    ) -> Self {
36        Self {
37            accumulators: Mutex::new(HashMap::new()),
38            config,
39            store_factory: Arc::new(store_factory),
40        }
41    }
42
43    /// Get the accumulator for `key`, creating it when missing.
44    pub fn get_or_create(&self, key: K) -> Arc<Accumulator<V>> {
45        let mut accumulators = self.accumulators.lock();
46        Arc::clone(accumulators.entry(key).or_insert_with(|| {
47            Arc::new(Accumulator::new(
48                (self.store_factory)(),
49                self.config.clone(),
50            ))
51        }))
52    }
53
54    /// Append a value to the keyed accumulator.
55    pub fn append(&self, key: K, value: V) -> AppResult<Option<Vec<V>>> {
56        self.get_or_create(key).append(value)
57    }
58
59    /// Flush a keyed accumulator.
60    pub fn flush(&self, key: &K) -> AppResult<Option<Vec<V>>> {
61        match self.accumulators.lock().get(key).cloned() {
62            Some(accumulator) => Ok(Some(accumulator.flush()?)),
63            None => Ok(None),
64        }
65    }
66
67    /// Return the buffered size for `key`.
68    pub fn size(&self, key: &K) -> AppResult<usize> {
69        match self.accumulators.lock().get(key).cloned() {
70            Some(accumulator) => accumulator.size(),
71            None => Ok(0),
72        }
73    }
74
75    /// List all currently active keys.
76    #[must_use]
77    pub fn keys(&self) -> Vec<K> {
78        self.accumulators.lock().keys().cloned().collect()
79    }
80
81    /// Remove expired accumulators and return the removed keys.
82    pub fn cleanup_expired(&self) -> AppResult<Vec<K>> {
83        let snapshot: Vec<(K, Arc<Accumulator<V>>)> = self
84            .accumulators
85            .lock()
86            .iter()
87            .map(|(key, accumulator)| (key.clone(), Arc::clone(accumulator)))
88            .collect();
89
90        let mut removed = Vec::new();
91        for (key, accumulator) in snapshot {
92            if accumulator.is_expired()? {
93                accumulator.close()?;
94                self.accumulators.lock().remove(&key);
95                removed.push(key);
96            }
97        }
98        Ok(removed)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use std::sync::atomic::{AtomicUsize, Ordering};
105    use std::sync::{Arc, Barrier};
106    use std::time::Duration;
107
108    use super::*;
109    use crate::{AccumulatorConfig, MemoryStore, SizeTrigger};
110
111    #[test]
112    fn manager_routes_values_by_key() {
113        let manager = Manager::new(AccumulatorConfig::new(), || {
114            Box::new(MemoryStore::<i32>::new())
115        });
116        manager.append("a", 1).unwrap();
117        manager.append("b", 2).unwrap();
118        assert_eq!(manager.size(&"a").unwrap(), 1);
119        assert_eq!(manager.size(&"b").unwrap(), 1);
120    }
121
122    #[tokio::test]
123    async fn manager_cleans_up_expired_accumulators() {
124        tokio::time::pause();
125        let manager = Manager::new(
126            AccumulatorConfig::new().with_ttl(Duration::from_secs(5)),
127            || Box::new(MemoryStore::<i32>::new()),
128        );
129        manager.append("a", 1).unwrap();
130        tokio::time::advance(Duration::from_secs(6)).await;
131        assert_eq!(manager.cleanup_expired().unwrap(), vec!["a"]);
132    }
133
134    #[test]
135    fn manager_propagates_flushes() {
136        let manager = Manager::new(
137            AccumulatorConfig::new().with_trigger(Arc::new(SizeTrigger::new(2))),
138            || Box::new(MemoryStore::<i32>::new()),
139        );
140        assert!(manager.append("a", 1).unwrap().is_none());
141        assert_eq!(manager.append("a", 2).unwrap(), Some(vec![1, 2]));
142    }
143
144    #[test]
145    fn get_or_create_invokes_factory_once_under_contention() {
146        const WORKERS: usize = 16;
147
148        let factory_calls = Arc::new(AtomicUsize::new(0));
149        let manager = Arc::new(Manager::new(AccumulatorConfig::new(), {
150            let factory_calls = Arc::clone(&factory_calls);
151            move || {
152                factory_calls.fetch_add(1, Ordering::SeqCst);
153                Box::new(MemoryStore::<i32>::new())
154            }
155        }));
156        let barrier = Arc::new(Barrier::new(WORKERS));
157
158        let threads: Vec<_> = (0..WORKERS)
159            .map(|_| {
160                let manager = Arc::clone(&manager);
161                let barrier = Arc::clone(&barrier);
162                std::thread::spawn(move || {
163                    barrier.wait();
164                    manager.get_or_create("shared")
165                })
166            })
167            .collect();
168
169        let accumulators: Vec<_> = threads
170            .into_iter()
171            .map(|thread| thread.join().unwrap())
172            .collect();
173
174        assert_eq!(factory_calls.load(Ordering::SeqCst), 1);
175        for accumulator in &accumulators[1..] {
176            assert!(Arc::ptr_eq(&accumulators[0], accumulator));
177        }
178    }
179}
rskit_stateful/manager.rs

rskit_stateful/
manager.rs
