mirror of
https://github.com/chenhg5/cc-connect.git
synced 2026-10-25 17:36:43 +08:00
Shawn
18eacddc44
RateLimiter.cleanupLoop() runs a background goroutine that never terminates. When SetRateLimitCfg() is called (e.g. during config reload), a new RateLimiter is created with a new goroutine, but the old goroutine keeps running forever — a goroutine leak. Add a stop channel and Stop() method to RateLimiter. Call Stop() on the previous limiter in SetRateLimitCfg() and in Engine.Stop() to ensure clean shutdown.
98 lines
2.0 KiB
Go
98 lines
2.0 KiB
Go
package core
|
|
|
|
import (
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRateLimiter_AllowWithinLimit(t *testing.T) {
|
|
rl := NewRateLimiter(5, time.Minute)
|
|
for i := 0; i < 5; i++ {
|
|
if !rl.Allow("user1") {
|
|
t.Errorf("request %d should be allowed", i+1)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_BlockExceedingLimit(t *testing.T) {
|
|
rl := NewRateLimiter(3, time.Minute)
|
|
for i := 0; i < 3; i++ {
|
|
rl.Allow("user1")
|
|
}
|
|
if rl.Allow("user1") {
|
|
t.Error("4th request should be blocked")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_DifferentKeys(t *testing.T) {
|
|
rl := NewRateLimiter(2, time.Minute)
|
|
rl.Allow("user1")
|
|
rl.Allow("user1")
|
|
|
|
if rl.Allow("user1") {
|
|
t.Error("user1 should be blocked")
|
|
}
|
|
if !rl.Allow("user2") {
|
|
t.Error("user2 should be allowed (independent bucket)")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_WindowExpiry(t *testing.T) {
|
|
rl := NewRateLimiter(2, 50*time.Millisecond)
|
|
rl.Allow("user1")
|
|
rl.Allow("user1")
|
|
|
|
if rl.Allow("user1") {
|
|
t.Error("should be blocked immediately")
|
|
}
|
|
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
if !rl.Allow("user1") {
|
|
t.Error("should be allowed after window expires")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Disabled(t *testing.T) {
|
|
rl := NewRateLimiter(0, time.Minute)
|
|
for i := 0; i < 100; i++ {
|
|
if !rl.Allow("user1") {
|
|
t.Error("should always allow when disabled")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Concurrent(t *testing.T) {
|
|
rl := NewRateLimiter(100, time.Minute)
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 200; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
rl.Allow("user1")
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
func TestRateLimiter_Stop(t *testing.T) {
|
|
rl := NewRateLimiter(5, time.Minute)
|
|
rl.Allow("user1")
|
|
|
|
// Stop should not panic and should be idempotent
|
|
rl.Stop()
|
|
rl.Stop() // second call should be safe
|
|
|
|
// Allow should still work after Stop (just no background cleanup)
|
|
if !rl.Allow("user2") {
|
|
t.Error("Allow should still work after Stop")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_StopDisabled(t *testing.T) {
|
|
// A disabled limiter (maxMessages=0) should also handle Stop gracefully
|
|
rl := NewRateLimiter(0, time.Minute)
|
|
rl.Stop()
|
|
}
|