Added session based request rate limiting, uses token bucket strategy

1 files changed, 58 insertions, 0 deletions

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)
+}