Skip to main content

teaql_runtime/data_service/
cache.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3use std::time::{Duration, Instant};
4
5use teaql_core::Record;
6
7#[derive(Debug, Default)]
8pub struct InMemoryAggregationCache {
9    namespace: String,
10    entries: Mutex<HashMap<String, AggregationCacheEntry>>,
11}
12
13pub trait AggregationCacheBackend: Send + Sync {
14    fn namespace(&self) -> &str;
15    fn get(&self, key: &str, max_age_millis: u64) -> Option<Vec<Record>>;
16    fn put(&self, key: String, rows: Vec<Record>);
17    fn invalidate_namespace(&self, namespace: &str);
18}
19
20#[derive(Debug, Clone)]
21struct AggregationCacheEntry {
22    stored_at: Instant,
23    rows: Vec<Record>,
24}
25
26impl InMemoryAggregationCache {
27    pub fn with_namespace(namespace: impl Into<String>) -> Self {
28        Self {
29            namespace: namespace.into(),
30            entries: Mutex::new(HashMap::new()),
31        }
32    }
33
34    pub fn namespace(&self) -> &str {
35        &self.namespace
36    }
37}
38
39impl AggregationCacheBackend for InMemoryAggregationCache {
40    fn namespace(&self) -> &str {
41        &self.namespace
42    }
43
44    fn get(&self, key: &str, max_age_millis: u64) -> Option<Vec<Record>> {
45        let entries = self.entries.lock().ok()?;
46        let entry = entries.get(key)?;
47        (max_age_millis == 0 || entry.stored_at.elapsed() <= Duration::from_millis(max_age_millis))
48            .then(|| entry.rows.clone())
49    }
50
51    fn put(&self, key: String, rows: Vec<Record>) {
52        if let Ok(mut entries) = self.entries.lock() {
53            entries.insert(
54                key,
55                AggregationCacheEntry {
56                    stored_at: Instant::now(),
57                    rows,
58                },
59            );
60        }
61    }
62
63    fn invalidate_namespace(&self, namespace: &str) {
64        if let Ok(mut entries) = self.entries.lock() {
65            let prefix = format!("{namespace}::");
66            entries.retain(|key, _| !key.starts_with(&prefix));
67        }
68    }
69}
70
71impl InMemoryAggregationCache {
72    pub fn get(&self, key: &str, max_age_millis: u64) -> Option<Vec<Record>> {
73        AggregationCacheBackend::get(self, key, max_age_millis)
74    }
75
76    pub fn put(&self, key: String, rows: Vec<Record>) {
77        AggregationCacheBackend::put(self, key, rows);
78    }
79
80    pub fn clear(&self) {
81        if let Ok(mut entries) = self.entries.lock() {
82            entries.clear();
83        }
84    }
85
86    pub fn invalidate_namespace(&self, namespace: &str) {
87        AggregationCacheBackend::invalidate_namespace(self, namespace);
88    }
89}