watchdog/internal/ratelimit/limiter.go
NotAShelf 7ecc03ac19
various: cleanup
Fixes a status code conflict in `LimitedResponseWriter`, and a clock
skew bug that I probably introduced last time I dealt with time. I hate
computers.

We now use `tie.Since()`, which employs a monotonic clock that is immune
to system wall clock changes. 

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: Iec3147c21c5a295170f48cbf1a4620596a6a6964
2026-05-25 10:19:21 +03:00

55 lines
1.3 KiB
Go

package ratelimit
import (
"sync"
"time"
)
// Implements a simple token bucket rate limiter
type TokenBucket struct {
mu sync.Mutex
tokens int
capacity int
refill int
interval time.Duration
lastFill time.Time
}
// Creates a rate limiter with specified capacity and refill rate capacity
func NewTokenBucket(capacity, refillPerInterval int, interval time.Duration) *TokenBucket {
return &TokenBucket{
tokens: capacity,
capacity: capacity,
refill: refillPerInterval,
interval: interval,
lastFill: time.Now(),
}
}
// Allow checks if a request should be allowed
// Uses monotonic time via time.Since() to prevent clock skew issues
func (tb *TokenBucket) Allow() bool {
tb.mu.Lock()
defer tb.mu.Unlock()
// Refill tokens based on elapsed time using monotonic clock
// time.Since() uses monotonic readings when available, unaffected by wall clock changes
elapsed := time.Since(tb.lastFill)
if elapsed >= tb.interval {
periods := int(elapsed / tb.interval)
tb.tokens += periods * tb.refill
if tb.tokens > tb.capacity {
tb.tokens = tb.capacity
}
// Advance lastFill by exact periods to prevent drift
tb.lastFill = tb.lastFill.Add(time.Duration(periods) * tb.interval)
}
// Check if we have tokens available
if tb.tokens > 0 {
tb.tokens--
return true
}
return false
}