Spaces:
Paused
Paused
File size: 1,559 Bytes
a48ca26 |
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 62 63 64 65 66 67 68 69 70 71 |
package common
import (
"sync"
"time"
)
type InMemoryRateLimiter struct {
store map[string]*[]int64
mutex sync.Mutex
expirationDuration time.Duration
}
func (l *InMemoryRateLimiter) Init(expirationDuration time.Duration) {
if l.store == nil {
l.mutex.Lock()
if l.store == nil {
l.store = make(map[string]*[]int64)
l.expirationDuration = expirationDuration
if expirationDuration > 0 {
go l.clearExpiredItems()
}
}
l.mutex.Unlock()
}
}
func (l *InMemoryRateLimiter) clearExpiredItems() {
for {
time.Sleep(l.expirationDuration)
l.mutex.Lock()
now := time.Now().Unix()
for key := range l.store {
queue := l.store[key]
size := len(*queue)
if size == 0 || now-(*queue)[size-1] > int64(l.expirationDuration.Seconds()) {
delete(l.store, key)
}
}
l.mutex.Unlock()
}
}
// Request parameter duration's unit is seconds
func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool {
l.mutex.Lock()
defer l.mutex.Unlock()
// [old <-- new]
queue, ok := l.store[key]
now := time.Now().Unix()
if ok {
if len(*queue) < maxRequestNum {
*queue = append(*queue, now)
return true
} else {
if now-(*queue)[0] >= duration {
*queue = (*queue)[1:]
*queue = append(*queue, now)
return true
} else {
return false
}
}
} else {
s := make([]int64, 0, maxRequestNum)
l.store[key] = &s
*(l.store[key]) = append(*(l.store[key]), now)
}
return true
}
|