From 0a0915a977c7acb087b270327ec5190b3b4428b5 Mon Sep 17 00:00:00 2001 From: Kyren223 Date: Mon, 14 Jul 2025 12:39:46 +0300 Subject: Added session based request rate limiting, uses token bucket strategy --- pkg/rate/rate.go | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 pkg/rate/rate.go (limited to 'pkg/rate') diff --git a/pkg/rate/rate.go b/pkg/rate/rate.go new file mode 100644 index 0000000..f1578e6 --- /dev/null +++ b/pkg/rate/rate.go @@ -0,0 +1,58 @@ +package rate + +import "time" + +type Limiter struct { + limit float64 + rate float64 + + lastRefill time.Time + tokens float64 +} + +// rate refills limiter rate tokens per second +func NewLimiter(rate float64, limit float64) Limiter { + return Limiter{ + limit: limit, + rate: rate, + lastRefill: time.Now().UTC(), + tokens: limit, + } +} + +func (rl *Limiter) Fill() { + rl.update() + rl.tokens = rl.limit +} + +func (rl *Limiter) SetRate(rate float64) { + rl.update() + rl.rate = rate +} + +func (rl *Limiter) SetLimit(limit float64) { + rl.update() + rl.limit = limit +} + +func (rl *Limiter) Take(tokens float64) bool { + rl.update() + has := rl.Has(tokens) + if has { + rl.tokens -= tokens + return true + } + return false +} + +func (rl *Limiter) Has(tokens float64) bool { + rl.update() + return rl.tokens >= tokens +} + +func (rl *Limiter) update() { + lastRefill := rl.lastRefill + rl.lastRefill = time.Now().UTC() + + rl.tokens += min(time.Since(lastRefill).Seconds()*rl.rate, rl.limit) +} -- cgit v1.3.1