From d0f864c5700bdebf0825a504a4836b5ae02ef93f Mon Sep 17 00:00:00 2001 From: Sasha Varlamov Date: Wed, 29 Jul 2026 15:08:29 +0000 Subject: [PATCH] perf(daemon): throttle agent_usage emission with an in-memory limiter should_emit_agent_usage acquired the global metrics-DB mutex and ran a SELECT + upsert transaction on every AI checkpoint just to rate-limit agent_usage telemetry to one event per session per 150s. That put a sqlite round trip (and contention on the shared DB lock) on the checkpoint processing path. Replace it with a bounded in-memory limiter: a HashMap + expiry deque capped at 10,000 sessions, expiring entries after the emission interval and evicting the oldest when full. Sub-microsecond, no I/O, no shared lock with metrics delivery. A daemon restart resets the limiter and may re-emit one event per active session; the downstream pipeline already tolerates duplicate agent_usage events. The sqlite agent_usage_throttle table and its accessor remain for schema compatibility but are no longer on the checkpoint path. Ported from #1979 (closed) onto current main. Co-Authored-By: Claude Fable 5 --- src/daemon/checkpoint.rs | 181 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 168 insertions(+), 13 deletions(-) diff --git a/src/daemon/checkpoint.rs b/src/daemon/checkpoint.rs index 50d4e5d492..0cc8772284 100644 --- a/src/daemon/checkpoint.rs +++ b/src/daemon/checkpoint.rs @@ -16,8 +16,16 @@ use crate::git::repository::Repository; use futures::stream::{self, StreamExt}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; +#[cfg(any(test, not(feature = "test-support")))] +use std::collections::VecDeque; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +#[cfg(any(test, not(feature = "test-support")))] +use std::sync::Mutex; +#[cfg(not(any(test, feature = "test-support")))] +use std::sync::OnceLock; +#[cfg(any(test, not(feature = "test-support")))] +use std::time::Duration; use std::time::Instant; #[cfg(not(any(test, feature = "test-support")))] use std::time::{SystemTime, UNIX_EPOCH}; @@ -43,28 +51,105 @@ use crate::authorship::working_log::AgentId; #[cfg_attr(any(test, feature = "test-support"), allow(dead_code))] const AGENT_USAGE_MIN_INTERVAL_SECS: u64 = 150; +#[cfg(any(test, not(feature = "test-support")))] +const AGENT_USAGE_LIMITER_CAPACITY: usize = 10_000; #[cfg(not(any(test, feature = "test-support")))] const KNOWN_HUMAN_MIN_SECS_AFTER_AI: u64 = 1; +#[cfg(any(test, not(feature = "test-support")))] +struct AgentUsageLimiter { + interval: Duration, + capacity: usize, + last_emitted: HashMap, + expiry_order: VecDeque<(String, Instant)>, +} + +#[cfg(any(test, not(feature = "test-support")))] +impl AgentUsageLimiter { + fn new(interval: Duration, capacity: usize) -> Self { + Self { + interval, + capacity, + last_emitted: HashMap::with_capacity(capacity), + expiry_order: VecDeque::with_capacity(capacity), + } + } + + fn should_emit(&mut self, prompt_id: &str, now: Instant) -> bool { + self.expire_stale(now); + + if self.last_emitted.contains_key(prompt_id) { + return false; + } + if self.capacity == 0 { + return true; + } + + while self.last_emitted.len() >= self.capacity { + let Some((oldest_id, emitted_at)) = self.expiry_order.pop_front() else { + self.last_emitted.clear(); + break; + }; + if self.last_emitted.get(&oldest_id) == Some(&emitted_at) { + self.last_emitted.remove(&oldest_id); + } + } + + let prompt_id = prompt_id.to_string(); + self.last_emitted.insert(prompt_id.clone(), now); + self.expiry_order.push_back((prompt_id, now)); + true + } + + fn expire_stale(&mut self, now: Instant) { + while let Some((prompt_id, emitted_at)) = self.expiry_order.front() { + if now.saturating_duration_since(*emitted_at) < self.interval { + break; + } + if self.last_emitted.get(prompt_id) == Some(emitted_at) { + self.last_emitted.remove(prompt_id); + } + self.expiry_order.pop_front(); + } + } + + #[cfg(test)] + fn entry_count(&self) -> usize { + self.last_emitted.len() + } + + #[cfg(test)] + fn expiry_count(&self) -> usize { + self.expiry_order.len() + } +} + #[cfg(not(any(test, feature = "test-support")))] -fn should_emit_real_agent_usage(agent_id: &AgentId) -> bool { - let prompt_id = generate_short_hash(&agent_id.id, &agent_id.tool); - let now_ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); +static AGENT_USAGE_LIMITER: OnceLock> = OnceLock::new(); - let Ok(db) = crate::metrics::db::MetricsDatabase::global() else { - return true; - }; - let Ok(mut db_lock) = db.lock() else { +#[cfg(any(test, not(feature = "test-support")))] +fn should_emit_with_limiter( + limiter: &Mutex, + prompt_id: &str, + now: Instant, +) -> bool { + let Ok(mut limiter) = limiter.lock() else { return true; }; + limiter.should_emit(prompt_id, now) +} - db_lock - .should_emit_agent_usage(&prompt_id, now_ts, AGENT_USAGE_MIN_INTERVAL_SECS) - .unwrap_or(true) +#[cfg(not(any(test, feature = "test-support")))] +fn should_emit_real_agent_usage(agent_id: &AgentId) -> bool { + let prompt_id = generate_short_hash(&agent_id.id, &agent_id.tool); + let limiter = AGENT_USAGE_LIMITER.get_or_init(|| { + Mutex::new(AgentUsageLimiter::new( + Duration::from_secs(AGENT_USAGE_MIN_INTERVAL_SECS), + AGENT_USAGE_LIMITER_CAPACITY, + )) + }); + should_emit_with_limiter(limiter, &prompt_id, Instant::now()) } #[cfg(any(test, feature = "test-support"))] @@ -1144,3 +1229,73 @@ mod tests { assert!(should_emit_agent_usage_with_throttle(&agent_id, |_| true)); } } + +#[cfg(test)] +mod agent_usage_limiter_tests { + use super::*; + use std::time::Duration; + + #[test] + fn suppresses_until_the_interval_expires() { + let mut limiter = AgentUsageLimiter::new(Duration::from_secs(150), 10_000); + let start = Instant::now(); + + assert!(limiter.should_emit("session", start)); + assert!(!limiter.should_emit("session", start + Duration::from_secs(149))); + assert!(limiter.should_emit("session", start + Duration::from_secs(150))); + } + + #[test] + fn capacity_is_bounded_and_evicts_the_oldest_valid_entry() { + assert_eq!(AGENT_USAGE_LIMITER_CAPACITY, 10_000); + let mut limiter = AgentUsageLimiter::new(Duration::from_secs(150), 3); + let start = Instant::now(); + + assert!(limiter.should_emit("first", start)); + assert!(limiter.should_emit("second", start + Duration::from_secs(1))); + assert!(limiter.should_emit("third", start + Duration::from_secs(2))); + assert_eq!(limiter.entry_count(), 3); + assert_eq!(limiter.expiry_count(), 3); + + assert!(limiter.should_emit("fourth", start + Duration::from_secs(3))); + assert_eq!(limiter.entry_count(), 3); + assert_eq!(limiter.expiry_count(), 3); + assert!( + limiter.should_emit("first", start + Duration::from_secs(4)), + "the oldest valid entry should have been evicted" + ); + assert!( + !limiter.should_emit("third", start + Duration::from_secs(4)), + "newer entries should remain suppressed" + ); + } + + #[test] + fn stale_entries_are_expired_before_capacity_eviction() { + let mut limiter = AgentUsageLimiter::new(Duration::from_secs(10), 2); + let start = Instant::now(); + + assert!(limiter.should_emit("stale", start)); + assert!(limiter.should_emit("live", start + Duration::from_secs(9))); + assert!(limiter.should_emit("new", start + Duration::from_secs(10))); + + assert_eq!(limiter.entry_count(), 2); + assert_eq!(limiter.expiry_count(), 2); + assert!(!limiter.should_emit("live", start + Duration::from_secs(10))); + } + + #[test] + fn poisoned_lock_fails_open() { + let limiter = Mutex::new(AgentUsageLimiter::new(Duration::from_secs(150), 1)); + let _ = std::panic::catch_unwind(|| { + let _guard = limiter.lock().unwrap(); + panic!("poison limiter"); + }); + + assert!(should_emit_with_limiter( + &limiter, + "session", + Instant::now() + )); + } +}