Skip to main content

dbkit_rs/
cache.rs

1use dashmap::DashMap;
2use std::sync::Arc;
3
4/// Concurrent key-value cache with named buckets.
5///
6/// Uses DashMap internally for lock-free concurrent reads. Create named
7/// buckets for different entity types (products, listings, etc.).
8#[derive(Debug, Clone)]
9pub struct Cache {
10    buckets: Arc<DashMap<String, Arc<DashMap<String, String>>>>,
11}
12
13impl Cache {
14    pub fn new() -> Self {
15        Self {
16            buckets: Arc::new(DashMap::new()),
17        }
18    }
19
20    /// Create a set of named buckets upfront.
21    pub fn with_buckets(names: &[&str]) -> Self {
22        let cache = Self::new();
23        for name in names {
24            cache.buckets.insert(name.to_string(), Arc::new(DashMap::new()));
25        }
26        cache
27    }
28
29    /// Get a value from a named bucket.
30    pub fn get(&self, bucket: &str, key: &str) -> Option<String> {
31        self.buckets
32            .get(bucket)
33            .and_then(|b| b.get(key).map(|v| v.clone()))
34    }
35
36    /// Set a value in a named bucket (creates the bucket if it doesn't exist).
37    pub fn set(&self, bucket: &str, key: String, value: String) {
38        self.buckets
39            .entry(bucket.to_string())
40            .or_insert_with(|| Arc::new(DashMap::new()))
41            .insert(key, value);
42    }
43
44    /// Remove a value from a named bucket.
45    pub fn remove(&self, bucket: &str, key: &str) -> Option<String> {
46        self.buckets
47            .get(bucket)
48            .and_then(|b| b.remove(key).map(|(_, v)| v))
49    }
50
51    /// Clear all entries in a specific bucket.
52    pub fn clear_bucket(&self, bucket: &str) {
53        if let Some(b) = self.buckets.get(bucket) {
54            b.clear();
55        }
56    }
57
58    /// Clear all buckets.
59    pub fn clear_all(&self) {
60        for entry in self.buckets.iter() {
61            entry.value().clear();
62        }
63    }
64}
65
66impl Default for Cache {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_cache_basic() {
78        let cache = Cache::with_buckets(&["products", "listings"]);
79
80        cache.set("products", "upc:123".into(), "uuid-1".into());
81        assert_eq!(cache.get("products", "upc:123"), Some("uuid-1".into()));
82        assert_eq!(cache.get("products", "upc:999"), None);
83        assert_eq!(cache.get("listings", "upc:123"), None);
84    }
85
86    #[test]
87    fn test_cache_auto_create_bucket() {
88        let cache = Cache::new();
89        cache.set("new_bucket", "key".into(), "value".into());
90        assert_eq!(cache.get("new_bucket", "key"), Some("value".into()));
91    }
92
93    #[test]
94    fn test_cache_clear() {
95        let cache = Cache::with_buckets(&["a", "b"]);
96        cache.set("a", "k1".into(), "v1".into());
97        cache.set("b", "k2".into(), "v2".into());
98
99        cache.clear_bucket("a");
100        assert_eq!(cache.get("a", "k1"), None);
101        assert_eq!(cache.get("b", "k2"), Some("v2".into()));
102
103        cache.clear_all();
104        assert_eq!(cache.get("b", "k2"), None);
105    }
106}