| #![allow(dead_code)] |
| |
|
|
| use once_cell::sync::Lazy; |
| use serde_json::Value; |
| use std::collections::HashMap; |
| use std::sync::RwLock; |
| use std::time::Instant; |
|
|
| |
| #[derive(Clone)] |
| struct CacheEntry { |
| |
| schema: Value, |
| |
| last_used: Instant, |
| |
| hit_count: usize, |
| } |
|
|
| |
| struct SchemaCache { |
| |
| cache: HashMap<String, CacheEntry>, |
| |
| stats: CacheStats, |
| } |
|
|
| |
| #[derive(Default, Clone, Debug)] |
| pub struct CacheStats { |
| |
| pub total_requests: usize, |
| |
| pub cache_hits: usize, |
| |
| pub cache_misses: usize, |
| } |
|
|
| impl CacheStats { |
| |
| pub fn hit_rate(&self) -> f64 { |
| if self.total_requests == 0 { |
| 0.0 |
| } else { |
| self.cache_hits as f64 / self.total_requests as f64 |
| } |
| } |
| } |
|
|
| impl SchemaCache { |
| fn new() -> Self { |
| Self { |
| cache: HashMap::new(), |
| stats: CacheStats::default(), |
| } |
| } |
|
|
| |
| fn get(&mut self, key: &str) -> Option<Value> { |
| self.stats.total_requests += 1; |
|
|
| if let Some(entry) = self.cache.get_mut(key) { |
| |
| entry.last_used = Instant::now(); |
| entry.hit_count += 1; |
| self.stats.cache_hits += 1; |
| Some(entry.schema.clone()) |
| } else { |
| self.stats.cache_misses += 1; |
| None |
| } |
| } |
|
|
| |
| fn insert(&mut self, key: String, schema: Value) { |
| |
| const MAX_CACHE_SIZE: usize = 1000; |
| if self.cache.len() >= MAX_CACHE_SIZE { |
| self.evict_lru(); |
| } |
|
|
| let entry = CacheEntry { |
| schema, |
| last_used: Instant::now(), |
| hit_count: 0, |
| }; |
| self.cache.insert(key, entry); |
| } |
|
|
| |
| fn evict_lru(&mut self) { |
| if self.cache.is_empty() { |
| return; |
| } |
|
|
| |
| let oldest_key = self |
| .cache |
| .iter() |
| .min_by_key(|(_, entry)| entry.last_used) |
| .map(|(key, _)| key.clone()); |
|
|
| if let Some(key) = oldest_key { |
| self.cache.remove(&key); |
| } |
| } |
|
|
| |
| fn stats(&self) -> CacheStats { |
| self.stats.clone() |
| } |
|
|
| |
| fn clear(&mut self) { |
| self.cache.clear(); |
| self.stats = CacheStats::default(); |
| } |
| } |
|
|
| |
| static SCHEMA_CACHE: Lazy<RwLock<SchemaCache>> = Lazy::new(|| RwLock::new(SchemaCache::new())); |
|
|
| |
| |
| |
| fn compute_schema_hash(schema: &Value) -> String { |
| use sha2::{Digest, Sha256}; |
|
|
| let mut hasher = Sha256::new(); |
| |
| let schema_str = schema.to_string(); |
| hasher.update(schema_str.as_bytes()); |
|
|
| |
| format!("{:x}", hasher.finalize())[..16].to_string() |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| pub fn clean_json_schema_cached(schema: &mut Value, tool_name: &str) { |
| |
| let hash = compute_schema_hash(schema); |
| let cache_key = format!("{}:{}", tool_name, hash); |
|
|
| |
| { |
| if let Ok(mut cache) = SCHEMA_CACHE.write() { |
| if let Some(cached) = cache.get(&cache_key) { |
| *schema = cached; |
| return; |
| } |
| } |
| } |
|
|
| |
| super::json_schema::clean_json_schema_for_tool(schema, tool_name); |
|
|
| |
| if let Ok(mut cache) = SCHEMA_CACHE.write() { |
| cache.insert(cache_key, schema.clone()); |
| } |
| } |
|
|
| |
| pub fn get_cache_stats() -> CacheStats { |
| SCHEMA_CACHE |
| .read() |
| .map(|cache| cache.stats()) |
| .unwrap_or_default() |
| } |
|
|
| |
| pub fn clear_cache() { |
| if let Ok(mut cache) = SCHEMA_CACHE.write() { |
| cache.clear(); |
| } |
| } |
|
|
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use serde_json::json; |
|
|
| #[test] |
| fn test_compute_schema_hash() { |
| let schema1 = json!({"type": "string"}); |
| let schema2 = json!({"type": "string"}); |
| let schema3 = json!({"type": "number"}); |
|
|
| let hash1 = compute_schema_hash(&schema1); |
| let hash2 = compute_schema_hash(&schema2); |
| let hash3 = compute_schema_hash(&schema3); |
|
|
| |
| assert_eq!(hash1, hash2); |
| |
| assert_ne!(hash1, hash3); |
| } |
|
|
| #[test] |
| fn test_cache_hit() { |
| clear_cache(); |
|
|
| let mut schema = json!({"type": "string", "minLength": 5}); |
| let tool_name = "test_tool"; |
|
|
| |
| clean_json_schema_cached(&mut schema, tool_name); |
|
|
| |
| let mut schema2 = json!({"type": "string", "minLength": 5}); |
| clean_json_schema_cached(&mut schema2, tool_name); |
|
|
| let stats = get_cache_stats(); |
| |
| assert!( |
| stats.cache_hits > 0, |
| "Expected cache hits, got: {:?}", |
| stats |
| ); |
| assert!(stats.hit_rate() > 0.0); |
| } |
|
|
| #[test] |
| fn test_cache_eviction() { |
| clear_cache(); |
|
|
| |
| for i in 0..1100 { |
| let mut schema = json!({"type": "string", "index": i}); |
| let tool_name = format!("tool_{}", i); |
| clean_json_schema_cached(&mut schema, &tool_name); |
| } |
|
|
| |
| let stats = get_cache_stats(); |
| assert!(stats.total_requests > 0); |
| } |
| } |
|
|