vtcode_core/tools/
optimized_registry.rs

1//! Optimized tool registry with reduced lock contention and improved caching
2
3use anyhow::Result;
4use hashbrown::HashMap;
5use parking_lot::RwLock;
6use serde_json::Value;
7use std::sync::Arc;
8use tokio::sync::Semaphore;
9
10use crate::tools::registry::ToolExecutionRecord;
11
12/// Lock-free tool metadata cache
13#[derive(Clone)]
14pub struct CachedToolMetadata {
15    pub name: String,
16    pub description: String,
17    pub parameters: Value,
18    pub is_cached: bool,
19    pub avg_execution_time_ms: u64,
20}
21
22/// Optimized tool registry with minimal lock contention
23pub struct OptimizedToolRegistry {
24    /// Read-heavy tool metadata cache (rarely updated)
25    tool_metadata: Arc<RwLock<HashMap<String, Arc<CachedToolMetadata>>>>,
26
27    /// Execution semaphore to limit concurrent tool executions
28    execution_semaphore: Arc<Semaphore>,
29
30    /// Hot path LRU cache for frequently accessed tools.
31    /// Uses an LRU cache so that stale entries are evicted when the
32    /// capacity is reached, instead of silently refusing new entries.
33    hot_cache: Arc<RwLock<lru::LruCache<String, Arc<CachedToolMetadata>>>>,
34
35    /// Execution statistics (append-only for performance)
36    execution_stats: Arc<RwLock<Vec<ToolExecutionRecord>>>,
37}
38
39/// Maximum number of entries in the hot-path tool metadata cache.
40const HOT_CACHE_CAPACITY: usize = 16;
41
42impl OptimizedToolRegistry {
43    pub fn new(max_concurrent_tools: usize) -> Self {
44        Self {
45            tool_metadata: Arc::new(RwLock::new(HashMap::with_capacity(64))),
46            execution_semaphore: Arc::new(Semaphore::new(max_concurrent_tools)),
47            hot_cache: Arc::new(RwLock::new(lru::LruCache::new(
48                std::num::NonZeroUsize::new(HOT_CACHE_CAPACITY).unwrap_or(std::num::NonZeroUsize::MIN),
49            ))),
50            execution_stats: Arc::new(RwLock::new(Vec::with_capacity(1024))),
51        }
52    }
53
54    /// Fast tool lookup with hot cache optimization
55    pub fn get_tool_metadata(&self, tool_name: &str) -> Option<Arc<CachedToolMetadata>> {
56        // Try hot cache first (most frequently used tools).
57        // `get` on LRU promotes the entry to most-recently-used automatically.
58        if let Some(metadata) = self.hot_cache.write().get(tool_name) {
59            return Some(Arc::clone(metadata));
60        }
61
62        // Fallback to main cache
63        let metadata = self.tool_metadata.read().get(tool_name).cloned()?;
64
65        // Promote to hot cache
66        self.promote_to_hot_cache(tool_name, &metadata);
67
68        Some(metadata)
69    }
70
71    /// Register tool metadata with minimal locking
72    pub fn register_tool(&self, metadata: CachedToolMetadata) {
73        let tool_name = metadata.name.clone();
74        let metadata_arc = Arc::new(metadata);
75
76        self.tool_metadata.write().insert(tool_name, metadata_arc);
77    }
78
79    /// Execute tool with concurrency control and performance tracking
80    pub async fn execute_tool_optimized(&self, tool_name: &str, _args: Value) -> Result<Value> {
81        // Acquire execution permit
82        let _permit = self.execution_semaphore.acquire().await?;
83
84        let start_time = std::time::Instant::now();
85
86        // Simulate tool execution (replace with actual implementation)
87        let result = self.execute_tool_impl(tool_name, _args).await;
88
89        let execution_time = start_time.elapsed();
90
91        // Record execution statistics inline to keep behavior deterministic.
92        self.record_execution_stats(tool_name, execution_time, result.is_ok());
93
94        result
95    }
96
97    /// Promote frequently accessed tools to hot cache.
98    /// The LRU cache automatically evicts the least-recently-used entry
99    /// when capacity is reached.
100    fn promote_to_hot_cache(&self, tool_name: &str, metadata: &Arc<CachedToolMetadata>) {
101        self.hot_cache.write().push(tool_name.to_string(), Arc::clone(metadata));
102    }
103
104    /// Record execution statistics with minimal blocking
105    fn record_execution_stats(&self, tool_name: &str, execution_time: std::time::Duration, success: bool) {
106        let record = ToolExecutionRecord {
107            tool_name: tool_name.to_string(),
108            requested_name: String::new(), // Optimize: reuse from pool
109            is_mcp: false,
110            mcp_provider: None,
111            args: Value::Null,
112            result: if success {
113                Ok(Value::Null)
114            } else {
115                Err("Error".to_string())
116            },
117            timestamp: std::time::SystemTime::now(),
118            success,
119            context: crate::tools::registry::HarnessContextSnapshot::new("optimized".to_string(), None),
120            timeout_category: None,
121            base_timeout_ms: None,
122            adaptive_timeout_ms: None,
123            effective_timeout_ms: Some(execution_time.as_millis() as u64),
124            circuit_breaker: false,
125            attempt: 1,
126            retry_after_ms: None,
127            circuit_breaker_state: None,
128        };
129
130        let mut stats = self.execution_stats.write();
131        // Prune stats to prevent memory leaks (KISS/DRY/Perf)
132        if stats.len() >= 1024 {
133            stats.drain(..128); // Remove oldest 128 entries
134        }
135        stats.push(record);
136    }
137
138    /// Actual tool execution implementation
139    async fn execute_tool_impl(&self, _tool_name: &str, _args: Value) -> Result<Value> {
140        // Placeholder - replace with actual tool execution logic
141        Ok(Value::String("success".to_string()))
142    }
143
144    /// Get execution statistics without blocking
145    pub fn get_stats_snapshot(&self) -> Vec<ToolExecutionRecord> {
146        self.execution_stats.read().clone()
147    }
148
149    /// Clear hot cache periodically to prevent memory bloat
150    pub fn clear_hot_cache(&self) {
151        self.hot_cache.write().clear();
152    }
153}
vtcode_core/tools/optimized_registry.rs

vtcode_core/tools/
optimized_registry.rs
