45 lines
747 B
Go
45 lines
747 B
Go
package logic
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type burst struct {
|
|
mu sync.Mutex
|
|
hits map[string][]time.Time
|
|
}
|
|
|
|
func newBurst() *burst {
|
|
return &burst{hits: make(map[string][]time.Time)}
|
|
}
|
|
|
|
func (b *burst) allow(key string, window time.Duration, n int) bool {
|
|
if b == nil || n < 1 {
|
|
return true
|
|
}
|
|
now := time.Now()
|
|
cut := now.Add(-window)
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
kept := b.hits[key][:0]
|
|
for _, t := range b.hits[key] {
|
|
if t.After(cut) {
|
|
kept = append(kept, t)
|
|
}
|
|
}
|
|
if len(kept) >= n {
|
|
b.hits[key] = kept
|
|
return false
|
|
}
|
|
b.hits[key] = append(kept, now)
|
|
if len(b.hits) > 8192 {
|
|
for k, ts := range b.hits {
|
|
if len(ts) == 0 || now.Sub(ts[len(ts)-1]) > time.Hour {
|
|
delete(b.hits, k)
|
|
}
|
|
}
|
|
}
|
|
return true
|
|
}
|