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

vtcode_core/tools/
optimized_registry.rs
