Skip to content

Commit d63d61f

Browse files
committed
fix: Fix goroutine leak in hosted agent rate limiter
- Added stop channel and stopped flag to RateLimiter struct - Modified replenishTokens to listen for stop signal and exit cleanly - Added Stop() method to gracefully shutdown rate limiter - Added Stop() method to HostedAgent to cleanup rate limiter on shutdown Fixes cursor bot issue: Rate Limiter Goroutine Leak
1 parent d8ec5cc commit d63d61f

1 file changed

Lines changed: 30 additions & 5 deletions

File tree

pkg/analyze/agents/hosted/hosted_agent.go

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ type RateLimiter struct {
5050
tokens chan struct{}
5151
interval time.Duration
5252
lastReset time.Time
53+
stopCh chan struct{}
54+
stopped bool
5355
}
5456

5557
// RetryConfig defines retry behavior
@@ -188,6 +190,8 @@ func NewRateLimiter(requestsPerMinute int) *RateLimiter {
188190
tokens: tokens,
189191
interval: interval,
190192
lastReset: time.Now(),
193+
stopCh: make(chan struct{}),
194+
stopped: false,
191195
}
192196

193197
// Start token replenishment goroutine
@@ -201,12 +205,18 @@ func (rl *RateLimiter) replenishTokens() {
201205
ticker := time.NewTicker(rl.interval)
202206
defer ticker.Stop()
203207

204-
for range ticker.C {
208+
for {
205209
select {
206-
case rl.tokens <- struct{}{}:
207-
// Token added successfully
208-
default:
209-
// Bucket is full, skip
210+
case <-rl.stopCh:
211+
// Stop signal received, exit goroutine
212+
return
213+
case <-ticker.C:
214+
select {
215+
case rl.tokens <- struct{}{}:
216+
// Token added successfully
217+
default:
218+
// Bucket is full, skip
219+
}
210220
}
211221
}
212222
}
@@ -221,6 +231,14 @@ func (rl *RateLimiter) waitForToken(ctx context.Context) error {
221231
}
222232
}
223233

234+
// Stop cleanly shuts down the rate limiter and stops the replenishment goroutine
235+
func (rl *RateLimiter) Stop() {
236+
if !rl.stopped {
237+
rl.stopped = true
238+
close(rl.stopCh)
239+
}
240+
}
241+
224242
// Name returns the agent name
225243
func (a *HostedAgent) Name() string {
226244
return a.name
@@ -461,6 +479,13 @@ func (a *HostedAgent) SetEnabled(enabled bool) {
461479
a.enabled = enabled
462480
}
463481

482+
// Stop cleanly shuts down the hosted agent and stops background goroutines
483+
func (a *HostedAgent) Stop() {
484+
if a.rateLimiter != nil {
485+
a.rateLimiter.Stop()
486+
}
487+
}
488+
464489
// UpdateCredentials updates the API key for authentication
465490
func (a *HostedAgent) UpdateCredentials(apiKey string) error {
466491
if apiKey == "" {

0 commit comments

Comments
 (0)