Skip to content

Commit 884fe78

Browse files
fix(tracker): add rate limiter to prevent event dispatch loops (#2515)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
1 parent db4973c commit 884fe78

3 files changed

Lines changed: 107 additions & 29 deletions

File tree

‎crates/forge_tracker/src/dispatch.rs‎

Lines changed: 45 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use super::Result;
1313
use crate::can_track::can_track;
1414
use crate::collect::{Collect, posthog};
1515
use crate::event::Identity;
16+
use crate::rate_limit::RateLimiter;
1617
use crate::{Event, EventKind, client_id};
1718

1819
const POSTHOG_API_SECRET: &str = match option_env!("POSTHOG_API_SECRET") {
@@ -47,6 +48,13 @@ static CACHED_PATH: LazyLock<Option<String>> = LazyLock::new(|| {
4748
});
4849
static CACHED_ARGS: LazyLock<Vec<String>> = LazyLock::new(|| std::env::args().skip(1).collect());
4950

51+
/// Maximum number of events that can be dispatched per minute.
52+
///
53+
/// This acts as a rate limiter to prevent runaway loops (e.g. when
54+
/// stdout/stderr is closed and every write error triggers another error event)
55+
/// while allowing normal tracking to continue for long-running sessions.
56+
const MAX_EVENTS_PER_MINUTE: usize = 1_000;
57+
5058
#[derive(Clone)]
5159
pub struct Tracker {
5260
collectors: Arc<Vec<Box<dyn Collect>>>,
@@ -56,6 +64,7 @@ pub struct Tracker {
5664
model: Arc<Mutex<Option<String>>>,
5765
conversation: Arc<Mutex<Option<Conversation>>>,
5866
is_logged_in: Arc<AtomicBool>,
67+
rate_limiter: Arc<RateLimiter>,
5968
}
6069

6170
impl Default for Tracker {
@@ -71,6 +80,7 @@ impl Default for Tracker {
7180
model: Arc::new(Mutex::new(None)),
7281
conversation: Arc::new(Mutex::new(None)),
7382
is_logged_in: Arc::new(AtomicBool::new(false)),
83+
rate_limiter: Arc::new(RateLimiter::new(MAX_EVENTS_PER_MINUTE)),
7484
}
7585
}
7686
}
@@ -93,35 +103,41 @@ impl Tracker {
93103
}
94104

95105
pub async fn dispatch(&self, event_kind: EventKind) -> Result<()> {
96-
if self.can_track {
97-
// Create a new event
98-
let email = self.email().await;
99-
let event = Event {
100-
event_name: event_kind.name(),
101-
event_value: event_kind.value(),
102-
start_time: self.start_time,
103-
cores: cores(),
104-
client_id: client_id(),
105-
os_name: os_name(),
106-
up_time: up_time(self.start_time),
107-
args: args(),
108-
path: path(),
109-
cwd: cwd(),
110-
user: user(),
111-
version: version(),
112-
email: email.clone(),
113-
model: self.model.lock().await.clone(),
114-
conversation: self.conversation().await,
115-
identity: match event_kind {
116-
EventKind::Login(id) => Some(id),
117-
_ => None,
118-
},
119-
};
120-
121-
// Dispatch the event to all collectors
122-
for collector in self.collectors.as_ref() {
123-
collector.collect(event.clone()).await?;
124-
}
106+
if !self.can_track {
107+
return Ok(());
108+
}
109+
110+
if !self.rate_limiter.check() {
111+
return Ok(()); // Drop event if rate limit exceeded
112+
}
113+
114+
// Create a new event
115+
let email = self.email().await;
116+
let event = Event {
117+
event_name: event_kind.name(),
118+
event_value: event_kind.value(),
119+
start_time: self.start_time,
120+
cores: cores(),
121+
client_id: client_id(),
122+
os_name: os_name(),
123+
up_time: up_time(self.start_time),
124+
args: args(),
125+
path: path(),
126+
cwd: cwd(),
127+
user: user(),
128+
version: version(),
129+
email: email.clone(),
130+
model: self.model.lock().await.clone(),
131+
conversation: self.conversation().await,
132+
identity: match event_kind {
133+
EventKind::Login(id) => Some(id),
134+
_ => None,
135+
},
136+
};
137+
138+
// Dispatch the event to all collectors
139+
for collector in self.collectors.as_ref() {
140+
collector.collect(event.clone()).await?;
125141
}
126142
Ok(())
127143
}

‎crates/forge_tracker/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ mod dispatch;
55
mod error;
66
mod event;
77
mod log;
8+
mod rate_limit;
89
pub use can_track::VERSION;
910
pub use dispatch::Tracker;
1011
use error::Result;
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
2+
3+
use chrono::Utc;
4+
5+
/// A simple thread-safe rate limiter that allows a maximum number of events per
6+
/// minute.
7+
#[derive(Debug)]
8+
pub struct RateLimiter {
9+
max_per_minute: usize,
10+
window_start: AtomicU64,
11+
count: AtomicUsize,
12+
}
13+
14+
impl RateLimiter {
15+
/// Creates a new rate limiter with the specified maximum events per minute.
16+
pub fn new(max_per_minute: usize) -> Self {
17+
Self {
18+
max_per_minute,
19+
window_start: AtomicU64::new(Utc::now().timestamp() as u64),
20+
count: AtomicUsize::new(0),
21+
}
22+
}
23+
24+
/// Checks if an event should be allowed based on the rate limit.
25+
/// Returns true if the event is allowed, false if it should be dropped.
26+
pub fn check(&self) -> bool {
27+
let now = Utc::now().timestamp() as u64;
28+
let window_start = self.window_start.load(Ordering::Relaxed);
29+
30+
// If a minute has passed, reset the window and counter
31+
if now.saturating_sub(window_start) >= 60 {
32+
// We use compare_exchange to avoid race conditions when multiple threads try to
33+
// reset
34+
if self
35+
.window_start
36+
.compare_exchange(window_start, now, Ordering::SeqCst, Ordering::Relaxed)
37+
.is_ok()
38+
{
39+
self.count.store(0, Ordering::SeqCst);
40+
}
41+
}
42+
43+
// Increment and check the rate limit
44+
self.count.fetch_add(1, Ordering::Relaxed) < self.max_per_minute
45+
}
46+
}
47+
48+
#[cfg(test)]
49+
mod tests {
50+
use super::*;
51+
52+
#[test]
53+
fn test_rate_limiter() {
54+
let limiter = RateLimiter::new(2);
55+
56+
assert!(limiter.check()); // 1
57+
assert!(limiter.check()); // 2
58+
assert!(!limiter.check()); // 3 - blocked
59+
assert!(!limiter.check()); // 4 - blocked
60+
}
61+
}

0 commit comments

Comments
 (0)