43 lines
949 B
Go
43 lines
949 B
Go
package ratelimiter
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type SlidingWindowRateLimiter struct {
|
|
sync.RWMutex
|
|
clients map[string][]time.Time
|
|
limit int
|
|
window time.Duration
|
|
}
|
|
|
|
func NewSlidingWindowLimiter(limit int, window time.Duration) *SlidingWindowRateLimiter {
|
|
return &SlidingWindowRateLimiter{
|
|
clients: make(map[string][]time.Time),
|
|
limit: limit,
|
|
window: window,
|
|
}
|
|
}
|
|
|
|
func (rl *SlidingWindowRateLimiter) Allow(ip string) (bool, time.Duration) {
|
|
rl.RLock()
|
|
defer rl.Unlock()
|
|
|
|
// add new request attempt
|
|
rl.clients[ip] = append(rl.clients[ip], time.Now())
|
|
|
|
// remove ones outside the window
|
|
for len(rl.clients[ip]) > 0 && time.Since(rl.clients[ip][0]) > rl.window {
|
|
rl.clients[ip] = rl.clients[ip][1:]
|
|
}
|
|
|
|
// do actual check now
|
|
if len(rl.clients[ip]) > rl.limit {
|
|
// calc retry wait time
|
|
retryAfter := rl.window - time.Since(rl.clients[ip][0])
|
|
return false, time.Duration(retryAfter)
|
|
}
|
|
return true, time.Duration(0)
|
|
}
|