Skip to main content

ai_lib_rust/cache/
backend.rs

1//! Cache backend implementations.
2
3use super::key::CacheKey;
4use crate::Result;
5use async_trait::async_trait;
6use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8use std::time::{Duration, Instant};
9
10#[derive(Clone)]
11struct CacheEntry {
12    data: Vec<u8>,
13    created_at: Instant,
14    ttl: Duration,
15    last_accessed: Instant,
16}
17
18impl CacheEntry {
19    fn new(data: Vec<u8>, ttl: Duration) -> Self {
20        let now = Instant::now();
21        Self {
22            data,
23            created_at: now,
24            ttl,
25            last_accessed: now,
26        }
27    }
28    fn is_expired(&self) -> bool {
29        self.created_at.elapsed() > self.ttl
30    }
31}
32
33#[async_trait]
34pub trait CacheBackend: Send + Sync {
35    async fn get(&self, key: &CacheKey) -> Result<Option<Vec<u8>>>;
36    async fn set(&self, key: &CacheKey, value: &[u8], ttl: Duration) -> Result<()>;
37    async fn delete(&self, key: &CacheKey) -> Result<bool>;
38    async fn exists(&self, key: &CacheKey) -> Result<bool>;
39    async fn clear(&self) -> Result<()>;
40    async fn len(&self) -> Result<usize>;
41    fn name(&self) -> &'static str;
42}
43
44pub struct MemoryCache {
45    entries: Arc<RwLock<HashMap<String, CacheEntry>>>,
46    max_entries: usize,
47}
48
49impl MemoryCache {
50    pub fn new(max_entries: usize) -> Self {
51        Self {
52            entries: Arc::new(RwLock::new(HashMap::new())),
53            max_entries,
54        }
55    }
56    fn evict_if_needed(&self, entries: &mut HashMap<String, CacheEntry>) {
57        entries.retain(|_, e| !e.is_expired());
58        while entries.len() >= self.max_entries {
59            let oldest = entries
60                .iter()
61                .min_by_key(|(_, e)| e.last_accessed)
62                .map(|(k, _)| k.clone());
63            if let Some(k) = oldest {
64                entries.remove(&k);
65            } else {
66                break;
67            }
68        }
69    }
70}
71
72#[async_trait]
73impl CacheBackend for MemoryCache {
74    async fn get(&self, key: &CacheKey) -> Result<Option<Vec<u8>>> {
75        let mut entries = self.entries.write().unwrap();
76        if let Some(entry) = entries.get_mut(&key.hash) {
77            if entry.is_expired() {
78                entries.remove(&key.hash);
79                return Ok(None);
80            }
81            entry.last_accessed = Instant::now();
82            return Ok(Some(entry.data.clone()));
83        }
84        Ok(None)
85    }
86    async fn set(&self, key: &CacheKey, value: &[u8], ttl: Duration) -> Result<()> {
87        let mut entries = self.entries.write().unwrap();
88        self.evict_if_needed(&mut entries);
89        entries.insert(key.hash.clone(), CacheEntry::new(value.to_vec(), ttl));
90        Ok(())
91    }
92    async fn delete(&self, key: &CacheKey) -> Result<bool> {
93        Ok(self.entries.write().unwrap().remove(&key.hash).is_some())
94    }
95    async fn exists(&self, key: &CacheKey) -> Result<bool> {
96        let entries = self.entries.read().unwrap();
97        Ok(entries
98            .get(&key.hash)
99            .map(|e| !e.is_expired())
100            .unwrap_or(false))
101    }
102    async fn clear(&self) -> Result<()> {
103        self.entries.write().unwrap().clear();
104        Ok(())
105    }
106    async fn len(&self) -> Result<usize> {
107        Ok(self
108            .entries
109            .read()
110            .unwrap()
111            .values()
112            .filter(|e| !e.is_expired())
113            .count())
114    }
115    fn name(&self) -> &'static str {
116        "memory"
117    }
118}
119
120pub struct NullCache;
121impl NullCache {
122    pub fn new() -> Self {
123        Self
124    }
125}
126impl Default for NullCache {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132#[async_trait]
133impl CacheBackend for NullCache {
134    async fn get(&self, _: &CacheKey) -> Result<Option<Vec<u8>>> {
135        Ok(None)
136    }
137    async fn set(&self, _: &CacheKey, _: &[u8], _: Duration) -> Result<()> {
138        Ok(())
139    }
140    async fn delete(&self, _: &CacheKey) -> Result<bool> {
141        Ok(false)
142    }
143    async fn exists(&self, _: &CacheKey) -> Result<bool> {
144        Ok(false)
145    }
146    async fn clear(&self) -> Result<()> {
147        Ok(())
148    }
149    async fn len(&self) -> Result<usize> {
150        Ok(0)
151    }
152    fn name(&self) -> &'static str {
153        "null"
154    }
155}