fix: update existing limiter rates & adapt flush test for dedup

This commit is contained in:
xboard
2026-03-25 18:31:37 +08:00
parent b51cecfe17
commit 691ca93c82
2 changed files with 19 additions and 5 deletions
+16 -2
View File
@@ -43,8 +43,7 @@ func (t *SpeedTracker) SetLogCallback(f SpeedTrackerLogCallback) {
t.logFunc = f t.logFunc = f
} }
// UpdateBuckets updates the UUID→userID mapping and cleans up removed users. // UpdateBuckets updates the UUID→userID mapping and syncs existing limiters.
// Rate limiters are created on-demand in GetLimiter.
func (t *SpeedTracker) UpdateBuckets() { func (t *SpeedTracker) UpdateBuckets() {
currentUsers := make([]panel.User, 0, 32) currentUsers := make([]panel.User, 0, 32)
t.limiter.mu.RLock() t.limiter.mu.RLock()
@@ -65,6 +64,21 @@ func (t *SpeedTracker) UpdateBuckets() {
if user.UUID != "" { if user.UUID != "" {
newUUIDMap[user.UUID] = user.ID newUUIDMap[user.UUID] = user.ID
} }
// Update existing limiter if speed changed
if lim, ok := t.buckets[user.ID]; ok {
if user.SpeedLimit > 0 {
bytesPerSec := int(user.SpeedLimit) * 1_000_000 / 8
burst := bytesPerSec
if burst < 64*1024 {
burst = 64 * 1024
}
lim.SetLimit(rate.Limit(bytesPerSec))
lim.SetBurst(burst)
} else {
delete(t.buckets, user.ID)
}
}
} }
// Clean up buckets for removed users // Clean up buckets for removed users
+3 -3
View File
@@ -164,10 +164,10 @@ func TestFlushAliveIPs(t *testing.T) {
t.Errorf("user 2 IPs: got %d, want 1", len(flushed[2])) t.Errorf("user 2 IPs: got %d, want 1", len(flushed[2]))
} }
// Without another Process call, aliveIPs still reflect last known state // Without another Process call, hash is same → returns nil (skip duplicate)
flushed2 := tr.FlushAliveIPs() flushed2 := tr.FlushAliveIPs()
if len(flushed2[1]) != 2 { if flushed2 != nil {
t.Errorf("expected same IPs without new Process, got %v", flushed2) t.Errorf("expected nil on duplicate flush, got %v", flushed2)
} }
} }