1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
| package limiter
import ( "errors" "sync" "time" )
type SlidingWindowLimiter struct { limit int window int64 smallWindow int64 smallWindows int64 counters map[int64]int mutex sync.Mutex }
func NewSlidingWindowLimiter(limit int, window, smallWindow time.Duration) (*SlidingWindowLimiter, error) { if window%smallWindow != 0 { return nil, errors.New("window cannot be split by integers") }
return &SlidingWindowLimiter{ limit: limit, window: int64(window), smallWindow: int64(smallWindow), smallWindows: int64(window / smallWindow), counters: make(map[int64]int), }, nil }
func (l *SlidingWindowLimiter) TryAcquire() bool { l.mutex.Lock() defer l.mutex.Unlock()
currentSmallWindow := time.Now().UnixNano() / l.smallWindow * l.smallWindow startSmallWindow := currentSmallWindow - l.smallWindow*(l.smallWindows-1)
var count int for smallWindow, counter := range l.counters { if smallWindow < startSmallWindow { delete(l.counters, smallWindow) } else { count += counter } }
if count >= l.limit { return false } l.counters[currentSmallWindow]++ return true }
|