Skip to main content

smcp_computer/mcp_clients/
resource_cache.rs

1/**
2* 文件名: resource_cache
3* 作者: Claude Code
4* 创建日期: 2025-12-26
5* 最后修改日期: 2025-12-26
6* 版权: 2023 JQQ. All rights reserved.
7* 依赖: tokio, smcp-computer
8* 描述: 资源缓存管理器 / Resource cache manager
9*
10* ================================================================================
11* 功能说明 / Functionality
12* ================================================================================
13*
14* 本模块实现了资源缓存的管理,包括:
15* - 缓存资源数据
16* - TTL(Time To Live)管理
17* - 缓存失效和刷新
18* - 线程安全的缓存访问
19*
20* ================================================================================
21*/
22use serde_json::Value;
23use std::collections::HashMap;
24use std::sync::Arc;
25use std::time::{Duration, Instant};
26use tokio::sync::RwLock;
27
28/// 缓存的资源条目
29#[derive(Debug, Clone)]
30pub struct CachedResource {
31    /// 资源数据
32    pub data: Value,
33    /// 缓存时间戳
34    pub cached_at: Instant,
35    /// TTL(生存时间)
36    pub ttl: Duration,
37    /// 版本号(每次更新递增)
38    pub version: u64,
39}
40
41impl CachedResource {
42    /// 创建新的缓存条目
43    pub fn new(data: Value, ttl: Duration) -> Self {
44        Self {
45            data,
46            cached_at: Instant::now(),
47            ttl,
48            version: 1,
49        }
50    }
51
52    /// 检查缓存是否已过期
53    pub fn is_expired(&self) -> bool {
54        self.cached_at.elapsed() > self.ttl
55    }
56
57    /// 获取剩余 TTL
58    pub fn remaining_ttl(&self) -> Duration {
59        if self.is_expired() {
60            Duration::ZERO
61        } else {
62            self.ttl.saturating_sub(self.cached_at.elapsed())
63        }
64    }
65
66    /// 更新缓存数据
67    pub fn refresh(&mut self, new_data: Value) {
68        self.data = new_data;
69        self.cached_at = Instant::now();
70        self.version += 1;
71    }
72}
73
74/// 资源缓存管理器
75///
76/// 负责管理资源数据的本地缓存,提供线程安全的缓存操作接口。
77#[derive(Debug, Clone)]
78pub struct ResourceCache {
79    /// 缓存数据(URI -> CachedResource)
80    cache: Arc<RwLock<HashMap<String, CachedResource>>>,
81    /// 默认 TTL
82    default_ttl: Duration,
83}
84
85impl ResourceCache {
86    /// 创建新的资源缓存管理器
87    ///
88    /// # 参数
89    /// - `default_ttl`: 默认的缓存生存时间
90    pub fn new(default_ttl: Duration) -> Self {
91        Self {
92            cache: Arc::new(RwLock::new(HashMap::new())),
93            default_ttl,
94        }
95    }
96
97    /// 设置缓存
98    ///
99    /// # 参数
100    /// - `uri`: 资源 URI
101    /// - `data`: 资源数据
102    /// - `ttl`: 可选的 TTL,如果为 None 则使用默认 TTL
103    pub async fn set(&self, uri: String, data: Value, ttl: Option<Duration>) {
104        let ttl = ttl.unwrap_or(self.default_ttl);
105        let mut cache = self.cache.write().await;
106        cache.insert(uri, CachedResource::new(data, ttl));
107    }
108
109    /// 获取缓存
110    ///
111    /// # 参数
112    /// - `uri`: 资源 URI
113    ///
114    /// # 返回
115    /// - `Some(Value)`: 缓存存在且未过期
116    /// - `None`: 缓存不存在或已过期
117    pub async fn get(&self, uri: &str) -> Option<Value> {
118        let mut cache = self.cache.write().await;
119
120        if let Some(cached) = cache.get(uri) {
121            if !cached.is_expired() {
122                return Some(cached.data.clone());
123            } else {
124                // 缓存已过期,移除
125                cache.remove(uri);
126            }
127        }
128        None
129    }
130
131    /// 刷新缓存(保留版本号)
132    ///
133    /// # 参数
134    /// - `uri`: 资源 URI
135    /// - `new_data`: 新的资源数据
136    ///
137    /// # 返回
138    /// - `Ok(version)`: 刷新成功,返回新版本号
139    /// - `Err(_)`: 缓存不存在
140    pub async fn refresh(&self, uri: &str, new_data: Value) -> Result<u64, String> {
141        let mut cache = self.cache.write().await;
142
143        if let Some(cached) = cache.get_mut(uri) {
144            cached.refresh(new_data);
145            Ok(cached.version)
146        } else {
147            Err(format!("Resource not cached: {}", uri))
148        }
149    }
150
151    /// 移除缓存
152    ///
153    /// # 参数
154    /// - `uri`: 资源 URI
155    ///
156    /// # 返回
157    /// - `true`: 找到并移除
158    /// - `false`: 未找到
159    pub async fn remove(&self, uri: &str) -> bool {
160        let mut cache = self.cache.write().await;
161        cache.remove(uri).is_some()
162    }
163
164    /// 清空所有缓存
165    pub async fn clear(&self) {
166        let mut cache = self.cache.write().await;
167        cache.clear();
168    }
169
170    /// 检查资源是否已缓存且未过期
171    ///
172    /// # 参数
173    /// - `uri`: 资源 URI
174    ///
175    /// # 返回
176    /// - `true`: 已缓存且未过期
177    /// - `false`: 未缓存或已过期
178    pub async fn contains(&self, uri: &str) -> bool {
179        let cache = self.cache.read().await;
180        if let Some(cached) = cache.get(uri) {
181            !cached.is_expired()
182        } else {
183            false
184        }
185    }
186
187    /// 获取缓存条目的详细信息
188    ///
189    /// # 参数
190    /// - `uri`: 资源 URI
191    ///
192    /// # 返回
193    /// - `Some(CachedResource)`: 缓存条目
194    /// - `None`: 不存在
195    pub async fn get_entry(&self, uri: &str) -> Option<CachedResource> {
196        let cache = self.cache.read().await;
197        cache.get(uri).cloned()
198    }
199
200    /// 获取缓存大小
201    pub async fn size(&self) -> usize {
202        let cache = self.cache.read().await;
203        cache.len()
204    }
205
206    /// 获取所有缓存的 URI 列表
207    pub async fn keys(&self) -> Vec<String> {
208        let cache = self.cache.read().await;
209        cache.keys().cloned().collect()
210    }
211
212    /// 清理过期的缓存条目
213    pub async fn cleanup_expired(&self) -> usize {
214        let mut cache = self.cache.write().await;
215        let initial_size = cache.len();
216
217        cache.retain(|_, cached| !cached.is_expired());
218
219        initial_size - cache.len()
220    }
221}
222
223impl Default for ResourceCache {
224    fn default() -> Self {
225        Self::new(Duration::from_secs(60)) // 默认 60 秒 TTL
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    #[tokio::test]
234    async fn test_set_and_get_cache() {
235        let cache = ResourceCache::new(Duration::from_secs(60));
236
237        let data = Value::String("test data".to_string());
238        cache
239            .set("window://test".to_string(), data.clone(), None)
240            .await;
241
242        let retrieved = cache.get("window://test").await;
243        assert!(retrieved.is_some());
244        assert_eq!(retrieved.unwrap(), data);
245    }
246
247    #[tokio::test]
248    async fn test_cache_expiration() {
249        let cache = ResourceCache::new(Duration::from_millis(100)); // 100ms TTL
250
251        let data = Value::String("test data".to_string());
252        cache.set("window://test".to_string(), data, None).await;
253
254        // 立即获取应该成功
255        assert!(cache.get("window://test").await.is_some());
256
257        // 等待过期
258        tokio::time::sleep(Duration::from_millis(150)).await;
259
260        // 过期后应该返回 None
261        assert!(cache.get("window://test").await.is_none());
262    }
263
264    #[tokio::test]
265    async fn test_refresh_cache() {
266        let cache = ResourceCache::new(Duration::from_secs(60));
267
268        let data1 = Value::String("version 1".to_string());
269        cache
270            .set("window://test".to_string(), data1.clone(), None)
271            .await;
272
273        let entry = cache.get_entry("window://test").await.unwrap();
274        assert_eq!(entry.version, 1);
275
276        let data2 = Value::String("version 2".to_string());
277        cache.refresh("window://test", data2.clone()).await.unwrap();
278
279        let entry = cache.get_entry("window://test").await.unwrap();
280        assert_eq!(entry.version, 2);
281        assert_eq!(entry.data, data2);
282    }
283
284    #[tokio::test]
285    async fn test_remove_cache() {
286        let cache = ResourceCache::new(Duration::from_secs(60));
287
288        cache
289            .set(
290                "window://test".to_string(),
291                Value::String("test".to_string()),
292                None,
293            )
294            .await;
295
296        assert!(cache.contains("window://test").await);
297
298        cache.remove("window://test").await;
299        assert!(!cache.contains("window://test").await);
300    }
301
302    #[tokio::test]
303    async fn test_cleanup_expired() {
304        let cache = ResourceCache::new(Duration::from_millis(100));
305
306        cache
307            .set(
308                "window://test1".to_string(),
309                Value::String("test1".to_string()),
310                None,
311            )
312            .await;
313        cache
314            .set(
315                "window://test2".to_string(),
316                Value::String("test2".to_string()),
317                Some(Duration::from_secs(60)), // 较长的 TTL
318            )
319            .await;
320
321        tokio::time::sleep(Duration::from_millis(150)).await;
322
323        let cleaned = cache.cleanup_expired().await;
324        assert_eq!(cleaned, 1);
325
326        assert!(!cache.contains("window://test1").await);
327        assert!(cache.contains("window://test2").await);
328    }
329}