use lru::LruCache;
use rch_common::Classification;
use std::num::NonZeroUsize;
use std::sync::Mutex;
use std::time::{Duration, Instant};
const DEFAULT_CACHE_CAPACITY: usize = 1000;
const DEFAULT_TTL_SECS: u64 = 30;
#[derive(Debug, Clone)]
struct CachedClassification {
classification: Classification,
inserted_at: Instant,
ttl: Duration,
}
impl CachedClassification {
fn is_expired(&self) -> bool {
self.inserted_at.elapsed() > self.ttl
}
}
#[derive(Debug, Clone)]
pub struct CacheConfig {
pub max_entries: usize,
pub ttl: Duration,
pub enabled: bool,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
max_entries: DEFAULT_CACHE_CAPACITY,
ttl: Duration::from_secs(DEFAULT_TTL_SECS),
enabled: true,
}
}
}
pub struct ClassificationCache {
cache: Mutex<LruCache<String, CachedClassification>>,
config: CacheConfig,
hits: Mutex<u64>,
misses: Mutex<u64>,
}
impl ClassificationCache {
pub fn new(config: CacheConfig) -> Self {
let capacity = NonZeroUsize::new(config.max_entries.max(1))
.unwrap_or(NonZeroUsize::new(DEFAULT_CACHE_CAPACITY).unwrap());
Self {
cache: Mutex::new(LruCache::new(capacity)),
config,
hits: Mutex::new(0),
misses: Mutex::new(0),
}
}
pub fn with_defaults() -> Self {
Self::new(CacheConfig::default())
}
pub fn get(&self, command: &str) -> Option<Classification> {
if !self.config.enabled {
return None;
}
let key = command;
let mut cache = self.cache.lock().ok()?;
if let Some(entry) = cache.get(key) {
if entry.is_expired() {
cache.pop(key);
if let Ok(mut misses) = self.misses.lock() {
*misses += 1;
}
tracing::trace!(command = %command, "Classification cache miss (expired)");
None
} else {
if let Ok(mut hits) = self.hits.lock() {
*hits += 1;
}
tracing::trace!(command = %command, "Classification cache hit");
Some(entry.classification.clone())
}
} else {
if let Ok(mut misses) = self.misses.lock() {
*misses += 1;
}
tracing::trace!(command = %command, "Classification cache miss");
None
}
}
pub fn put(&self, command: &str, classification: Classification) {
if !self.config.enabled {
return;
}
let key = command.to_string();
let entry = CachedClassification {
classification,
inserted_at: Instant::now(),
ttl: self.config.ttl,
};
if let Ok(mut cache) = self.cache.lock() {
cache.put(key, entry);
tracing::trace!(command = %command, "Classification cached");
}
}
#[allow(dead_code)] pub fn clear(&self) {
if let Ok(mut cache) = self.cache.lock() {
cache.clear();
}
}
#[allow(dead_code)] pub fn stats(&self) -> CacheStats {
let hits = self.hits.lock().map(|g| *g).unwrap_or(0);
let misses = self.misses.lock().map(|g| *g).unwrap_or(0);
let len = self.cache.lock().map(|c| c.len()).unwrap_or(0);
CacheStats { hits, misses, len }
}
}
impl Default for ClassificationCache {
fn default() -> Self {
Self::with_defaults()
}
}
#[allow(dead_code)] #[derive(Debug, Clone, Default)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub len: usize,
}
impl CacheStats {
#[allow(dead_code)] pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
(self.hits as f64 / total as f64) * 100.0
}
}
}
static GLOBAL_CACHE: std::sync::OnceLock<ClassificationCache> = std::sync::OnceLock::new();
pub fn global_cache() -> &'static ClassificationCache {
GLOBAL_CACHE.get_or_init(ClassificationCache::with_defaults)
}
#[allow(dead_code)] pub fn classify_cached<F>(command: &str, classify_fn: F) -> Classification
where
F: FnOnce(&str) -> Classification,
{
classify_cached_inner(command, classify_fn, is_hook_mode())
}
pub(crate) fn classify_hook_command<F>(command: &str, classify_fn: F) -> Classification
where
F: FnOnce(&str) -> Classification,
{
classify_cached_inner(command, classify_fn, true)
}
fn classify_cached_inner<F>(command: &str, classify_fn: F, hook_mode: bool) -> Classification
where
F: FnOnce(&str) -> Classification,
{
let trimmed = command.trim();
if hook_mode {
return classify_fn(trimmed);
}
let cache = global_cache();
if let Some(cached) = cache.get(trimmed) {
return cached;
}
let classification = classify_fn(trimmed);
cache.put(trimmed, classification.clone());
classification
}
#[allow(dead_code)] fn is_hook_mode() -> bool {
std::env::var_os("RCH_HOOK_MODE").is_some_and(|v| {
let s = v.to_string_lossy().to_lowercase();
s == "1" || s == "true"
})
}
#[cfg(test)]
mod tests {
use super::*;
use rch_common::{CompilationKind, classify_command};
use std::borrow::Cow;
fn make_classification(is_compilation: bool) -> Classification {
Classification {
is_compilation,
confidence: if is_compilation { 0.95 } else { 0.0 },
kind: if is_compilation {
Some(CompilationKind::CargoBuild)
} else {
None
},
reason: Cow::Borrowed("test"),
command_prefix: None,
extracted_command: None,
}
}
#[test]
fn test_cache_basic_operations() {
let cache = ClassificationCache::with_defaults();
assert!(cache.get("cargo build").is_none());
let class = make_classification(true);
cache.put("cargo build", class.clone());
let cached = cache.get("cargo build");
assert!(cached.is_some());
assert!(cached.unwrap().is_compilation);
assert!(cache.get("cargo test").is_none());
}
#[test]
fn test_cache_ttl_expiration() {
let config = CacheConfig {
max_entries: 100,
ttl: Duration::from_millis(50), enabled: true,
};
let cache = ClassificationCache::new(config);
let class = make_classification(true);
cache.put("cargo build", class);
assert!(cache.get("cargo build").is_some());
std::thread::sleep(Duration::from_millis(60));
assert!(cache.get("cargo build").is_none());
}
#[test]
fn test_cache_lru_eviction() {
let config = CacheConfig {
max_entries: 2, ttl: Duration::from_secs(30),
enabled: true,
};
let cache = ClassificationCache::new(config);
cache.put("cmd1", make_classification(true));
cache.put("cmd2", make_classification(true));
assert!(cache.get("cmd1").is_some());
assert!(cache.get("cmd2").is_some());
cache.get("cmd1");
cache.put("cmd3", make_classification(true));
assert!(cache.get("cmd1").is_some());
assert!(cache.get("cmd3").is_some());
assert!(cache.get("cmd2").is_none());
}
#[test]
fn test_cache_disabled() {
let config = CacheConfig {
max_entries: 100,
ttl: Duration::from_secs(30),
enabled: false,
};
let cache = ClassificationCache::new(config);
let class = make_classification(true);
cache.put("cargo build", class);
assert!(cache.get("cargo build").is_none());
}
#[test]
fn test_cache_stats() {
let cache = ClassificationCache::with_defaults();
let stats = cache.stats();
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
assert_eq!(stats.len, 0);
cache.get("cargo build");
let stats = cache.stats();
assert_eq!(stats.misses, 1);
cache.put("cargo build", make_classification(true));
cache.get("cargo build");
let stats = cache.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.len, 1);
}
#[test]
fn test_cache_assumes_pretrimmed_keys() {
let cache = ClassificationCache::with_defaults();
cache.put("cargo build", make_classification(true));
assert!(cache.get("cargo build").is_some());
assert!(cache.get(" cargo build ").is_none());
}
#[test]
fn test_classify_cached_normalizes_whitespace_via_single_trim() {
let trimmed_result = classify_cached("cargo build", classify_command);
let padded_result = classify_cached(" cargo build ", classify_command);
assert_eq!(trimmed_result.is_compilation, padded_result.is_compilation);
assert_eq!(trimmed_result.kind, padded_result.kind);
}
#[test]
fn test_classify_cached_hook_mode_bypasses_cache() {
let unique_cmd = "cargo build --t16-hook-bypass-marker";
for _ in 0..50 {
let _ = classify_cached_inner(unique_cmd, classify_command, true);
}
assert!(global_cache().get(unique_cmd).is_none());
}
#[test]
fn test_classify_hook_command_bypasses_cache_without_env() {
let unique_cmd = "cargo build --hook-entry-bypass-without-env-marker";
for _ in 0..50 {
let _ = classify_hook_command(unique_cmd, classify_command);
}
assert!(global_cache().get(unique_cmd).is_none());
}
#[test]
fn test_classify_cached_default_mode_uses_cache() {
let unique_cmd = "cargo test --t16-default-mode-marker";
let r1 = classify_cached_inner(unique_cmd, classify_command, false);
let r2 = classify_cached_inner(unique_cmd, classify_command, false);
assert_eq!(r1.is_compilation, r2.is_compilation);
assert_eq!(r1.kind, r2.kind);
assert!(
global_cache().get(unique_cmd).is_some(),
"default mode should populate the cache"
);
}
#[test]
fn test_is_hook_mode_truthy_env_values() {
let truthy = ["1", "true", "TRUE", "True"];
let falsy = ["0", "false", "no", "yes", ""];
for val in truthy {
let lc = val.to_lowercase();
assert!(lc == "1" || lc == "true", "expected {val:?} to be truthy");
}
for val in falsy {
let lc = val.to_lowercase();
assert!(
!(lc == "1" || lc == "true"),
"expected {val:?} to NOT be truthy"
);
}
}
#[test]
fn test_classify_cached() {
let call_count = std::sync::atomic::AtomicUsize::new(0);
let classify = |_cmd: &str| {
call_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
make_classification(true)
};
let result = classify_cached("cargo build", classify);
assert!(result.is_compilation);
}
#[test]
fn test_hit_rate_calculation() {
let stats = CacheStats {
hits: 80,
misses: 20,
len: 10,
};
assert!((stats.hit_rate() - 80.0).abs() < 0.001);
let empty_stats = CacheStats::default();
assert_eq!(empty_stats.hit_rate(), 0.0);
}
#[test]
fn test_cache_config_default() {
let config = CacheConfig::default();
assert_eq!(config.max_entries, DEFAULT_CACHE_CAPACITY);
assert_eq!(config.ttl, Duration::from_secs(DEFAULT_TTL_SECS));
assert!(config.enabled);
}
#[test]
fn test_cache_config_clone() {
let config = CacheConfig {
max_entries: 500,
ttl: Duration::from_secs(60),
enabled: false,
};
let cloned = config.clone();
assert_eq!(cloned.max_entries, 500);
assert_eq!(cloned.ttl, Duration::from_secs(60));
assert!(!cloned.enabled);
}
#[test]
fn test_cache_config_debug() {
let config = CacheConfig {
max_entries: 100,
ttl: Duration::from_secs(10),
enabled: true,
};
let debug = format!("{:?}", config);
assert!(debug.contains("max_entries"));
assert!(debug.contains("100"));
assert!(debug.contains("enabled"));
}
#[test]
fn test_cache_stats_clone() {
let stats = CacheStats {
hits: 100,
misses: 50,
len: 25,
};
let cloned = stats.clone();
assert_eq!(cloned.hits, 100);
assert_eq!(cloned.misses, 50);
assert_eq!(cloned.len, 25);
}
#[test]
fn test_cache_stats_debug() {
let stats = CacheStats {
hits: 42,
misses: 13,
len: 10,
};
let debug = format!("{:?}", stats);
assert!(debug.contains("hits"));
assert!(debug.contains("42"));
assert!(debug.contains("misses"));
assert!(debug.contains("13"));
}
#[test]
fn test_cache_stats_default() {
let stats = CacheStats::default();
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
assert_eq!(stats.len, 0);
assert_eq!(stats.hit_rate(), 0.0);
}
#[test]
fn test_classification_cache_default() {
let cache = ClassificationCache::default();
let stats = cache.stats();
assert_eq!(stats.len, 0);
cache.put("test cmd", make_classification(true));
assert!(cache.get("test cmd").is_some());
}
#[test]
fn test_cache_clear() {
let cache = ClassificationCache::with_defaults();
cache.put("cmd1", make_classification(true));
cache.put("cmd2", make_classification(true));
cache.put("cmd3", make_classification(false));
let stats = cache.stats();
assert_eq!(stats.len, 3);
cache.clear();
let stats = cache.stats();
assert_eq!(stats.len, 0);
assert!(cache.get("cmd1").is_none());
assert!(cache.get("cmd2").is_none());
assert!(cache.get("cmd3").is_none());
}
#[test]
fn test_cache_zero_capacity() {
let config = CacheConfig {
max_entries: 0,
ttl: Duration::from_secs(30),
enabled: true,
};
let cache = ClassificationCache::new(config);
cache.put("cmd1", make_classification(true));
assert!(cache.get("cmd1").is_some());
cache.put("cmd2", make_classification(true));
assert!(cache.get("cmd2").is_some());
assert!(cache.get("cmd1").is_none());
}
#[test]
fn test_global_cache_returns_same_instance() {
let cache1 = global_cache();
let cache2 = global_cache();
assert!(std::ptr::eq(cache1, cache2));
}
#[test]
fn test_cache_hit_rate_various_scenarios() {
let stats = CacheStats {
hits: 50,
misses: 50,
len: 10,
};
assert!((stats.hit_rate() - 50.0).abs() < 0.001);
let stats = CacheStats {
hits: 100,
misses: 0,
len: 5,
};
assert!((stats.hit_rate() - 100.0).abs() < 0.001);
let stats = CacheStats {
hits: 0,
misses: 100,
len: 0,
};
assert!((stats.hit_rate() - 0.0).abs() < 0.001);
}
#[test]
fn test_cache_put_overwrites_existing() {
let cache = ClassificationCache::with_defaults();
let class1 = make_classification(true);
cache.put("cmd", class1);
let cached = cache.get("cmd").unwrap();
assert!(cached.is_compilation);
let class2 = make_classification(false);
cache.put("cmd", class2);
let cached = cache.get("cmd").unwrap();
assert!(!cached.is_compilation);
}
#[test]
fn test_cache_direct_keys_are_byte_exact() {
let cache = ClassificationCache::with_defaults();
cache.put(" command ", make_classification(true));
assert!(cache.get(" command ").is_some());
assert!(cache.get("command").is_none());
assert!(cache.get(" command").is_none());
assert!(cache.get("command ").is_none());
let stats = cache.stats();
assert_eq!(stats.len, 1);
}
#[test]
fn test_cache_multiple_misses_tracked() {
let cache = ClassificationCache::with_defaults();
cache.get("missing1");
cache.get("missing2");
cache.get("missing3");
cache.get("missing4");
cache.get("missing5");
let stats = cache.stats();
assert_eq!(stats.misses, 5);
assert_eq!(stats.hits, 0);
}
#[test]
fn test_cache_multiple_hits_tracked() {
let cache = ClassificationCache::with_defaults();
cache.put("cmd", make_classification(true));
cache.get("cmd");
cache.get("cmd");
cache.get("cmd");
let stats = cache.stats();
assert_eq!(stats.hits, 3);
}
#[test]
fn test_cached_classification_debug() {
let entry = CachedClassification {
classification: make_classification(true),
inserted_at: Instant::now(),
ttl: Duration::from_secs(30),
};
let debug = format!("{:?}", entry);
assert!(debug.contains("classification"));
assert!(debug.contains("inserted_at"));
assert!(debug.contains("ttl"));
}
#[test]
fn test_cached_classification_clone() {
let entry = CachedClassification {
classification: make_classification(true),
inserted_at: Instant::now(),
ttl: Duration::from_secs(60),
};
let cloned = entry.clone();
assert_eq!(
cloned.classification.is_compilation,
entry.classification.is_compilation
);
assert_eq!(cloned.ttl, entry.ttl);
}
#[test]
fn test_cached_classification_is_expired() {
let entry = CachedClassification {
classification: make_classification(true),
inserted_at: Instant::now(),
ttl: Duration::from_millis(10),
};
assert!(!entry.is_expired());
std::thread::sleep(Duration::from_millis(15));
assert!(entry.is_expired());
}
#[test]
fn test_cache_non_compilation_classification() {
let cache = ClassificationCache::with_defaults();
let class = make_classification(false);
cache.put("echo hello", class);
let cached = cache.get("echo hello").unwrap();
assert!(!cached.is_compilation);
assert!(cached.kind.is_none());
}
}