Skip to main content

flag_rs/
completion_cache.rs

1//! Caching system for expensive completion operations
2//!
3//! This module provides a time-based cache for completion results to improve
4//! performance when users repeatedly request completions for the same context.
5
6use crate::completion::CompletionResult;
7use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11/// A cached completion entry with timestamp
12#[derive(Clone)]
13struct CacheEntry {
14    result: CompletionResult,
15    timestamp: Instant,
16}
17
18/// A thread-safe cache for completion results
19///
20/// The cache automatically expires entries after a configurable duration
21/// to ensure that completions remain fresh while still providing performance benefits.
22pub struct CompletionCache {
23    cache: Arc<Mutex<HashMap<String, CacheEntry>>>,
24    ttl: Duration,
25}
26
27impl CompletionCache {
28    /// Creates a new completion cache with the specified time-to-live
29    ///
30    /// # Arguments
31    ///
32    /// * `ttl` - How long cached entries should remain valid
33    pub fn new(ttl: Duration) -> Self {
34        Self {
35            cache: Arc::new(Mutex::new(HashMap::new())),
36            ttl,
37        }
38    }
39
40    /// Creates a new completion cache with a default TTL of 5 seconds
41    pub fn with_default_ttl() -> Self {
42        Self::new(Duration::from_secs(5))
43    }
44
45    /// Generates a cache key from completion context
46    ///
47    /// The key includes the command path and current prefix to ensure
48    /// we only return cached results for identical contexts.
49    pub fn make_key(
50        command_path: &[String],
51        prefix: &str,
52        flags: &HashMap<String, String>,
53    ) -> String {
54        let mut parts = vec![];
55
56        parts.extend(command_path.iter().cloned());
57
58        // Include the prefix being completed
59        parts.push(format!("__prefix:{prefix}"));
60
61        // Include relevant flags that might affect completion
62        let mut flag_parts: Vec<String> = flags.iter().map(|(k, v)| format!("{k}={v}")).collect();
63        flag_parts.sort(); // Ensure consistent ordering
64        parts.extend(flag_parts);
65
66        parts.join(":")
67    }
68
69    /// Attempts to get a cached completion result
70    ///
71    /// Returns `Some(CompletionResult)` if a valid cached entry exists,
72    /// or `None` if the entry doesn't exist or has expired.
73    pub fn get(&self, key: &str) -> Option<CompletionResult> {
74        let mut cache = self.cache.lock().ok()?;
75
76        if let Some(entry) = cache.get(key) {
77            if entry.timestamp.elapsed() < self.ttl {
78                return Some(entry.result.clone());
79            }
80            // Entry has expired, remove it
81            cache.remove(key);
82        }
83
84        None
85    }
86
87    /// Stores a completion result in the cache
88    ///
89    /// # Arguments
90    ///
91    /// * `key` - The cache key
92    /// * `result` - The completion result to cache
93    pub fn put(&self, key: String, result: CompletionResult) {
94        if let Ok(mut cache) = self.cache.lock() {
95            cache.insert(
96                key,
97                CacheEntry {
98                    result,
99                    timestamp: Instant::now(),
100                },
101            );
102
103            // Opportunistically clean up expired entries
104            self.cleanup_expired(&mut cache);
105        }
106    }
107
108    /// Removes expired entries from the cache
109    fn cleanup_expired(&self, cache: &mut HashMap<String, CacheEntry>) {
110        let now = Instant::now();
111        cache.retain(|_, entry| now.duration_since(entry.timestamp) < self.ttl);
112    }
113
114    /// Clears all cached entries
115    pub fn clear(&self) {
116        if let Ok(mut cache) = self.cache.lock() {
117            cache.clear();
118        }
119    }
120
121    /// Returns the number of cached entries
122    pub fn size(&self) -> usize {
123        self.cache.lock().map_or(0, |c| c.len())
124    }
125}
126
127impl Default for CompletionCache {
128    fn default() -> Self {
129        Self::with_default_ttl()
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use crate::completion::CompletionResult;
137
138    #[test]
139    fn test_cache_basic_operations() {
140        let cache = CompletionCache::new(Duration::from_secs(1));
141        let key = "test:key";
142        let result = CompletionResult::new().add("item1").add("item2");
143
144        // Test cache miss
145        assert!(cache.get(key).is_none());
146
147        // Test cache put and hit
148        cache.put(key.to_string(), result.clone());
149        let cached = cache.get(key).unwrap();
150        assert_eq!(cached.values, result.values);
151
152        // Test expiration
153        std::thread::sleep(Duration::from_millis(1100));
154        assert!(cache.get(key).is_none());
155    }
156
157    #[test]
158    fn test_cache_key_generation() {
159        let mut flags = HashMap::new();
160        flags.insert("namespace".to_string(), "default".to_string());
161        flags.insert("verbose".to_string(), "true".to_string());
162
163        let key1 =
164            CompletionCache::make_key(&["kubectl".to_string(), "get".to_string()], "po", &flags);
165        let key2 =
166            CompletionCache::make_key(&["kubectl".to_string(), "get".to_string()], "po", &flags);
167        assert_eq!(key1, key2);
168
169        // Different prefix should generate different key
170        let key3 =
171            CompletionCache::make_key(&["kubectl".to_string(), "get".to_string()], "pod", &flags);
172        assert_ne!(key1, key3);
173
174        // Different flags should generate different key
175        flags.insert("all-namespaces".to_string(), "true".to_string());
176        let key4 =
177            CompletionCache::make_key(&["kubectl".to_string(), "get".to_string()], "po", &flags);
178        assert_ne!(key1, key4);
179    }
180
181    #[test]
182    fn test_cache_cleanup() {
183        let cache = CompletionCache::new(Duration::from_millis(100));
184
185        // Add multiple entries
186        for i in 0..5 {
187            let key = format!("key{i}");
188            let result = CompletionResult::new().add(format!("item{i}"));
189            cache.put(key, result);
190        }
191
192        assert_eq!(cache.size(), 5);
193
194        // Wait for expiration
195        std::thread::sleep(Duration::from_millis(150));
196
197        // Trigger cleanup by adding a new entry
198        cache.put("new".to_string(), CompletionResult::new());
199
200        // Only the new entry should remain
201        assert_eq!(cache.size(), 1);
202    }
203}