50 lines
914 B
Go
50 lines
914 B
Go
package ratelimiter
|
|
|
|
// Copied from https://www.youtube.com/watch?v=m5oyY9fgZPs
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type FixedWindowRateLimiter struct {
|
|
sync.RWMutex
|
|
clients map[string]int
|
|
limit int
|
|
window time.Duration
|
|
}
|
|
|
|
func NewFixedWindowLimiter(limit int, window time.Duration) *FixedWindowRateLimiter {
|
|
return &FixedWindowRateLimiter{
|
|
clients: make(map[string]int),
|
|
limit: limit,
|
|
window: window,
|
|
}
|
|
}
|
|
|
|
func (rl *FixedWindowRateLimiter) Allow(ip string) (bool, time.Duration) {
|
|
rl.RLock()
|
|
count, exists := rl.clients[ip]
|
|
rl.RUnlock()
|
|
|
|
if !exists || count < rl.limit {
|
|
// begin the reset count for a new window
|
|
rl.Lock()
|
|
if !exists {
|
|
go rl.resetCount(ip)
|
|
}
|
|
rl.clients[ip]++
|
|
rl.Unlock()
|
|
return true, time.Duration(0)
|
|
}
|
|
return false, rl.window
|
|
|
|
}
|
|
|
|
func (rl *FixedWindowRateLimiter) resetCount(ip string) {
|
|
time.Sleep(rl.window)
|
|
rl.Lock()
|
|
delete(rl.clients, ip)
|
|
rl.Unlock()
|
|
}
|