61 lines
1.0 KiB
Go
61 lines
1.0 KiB
Go
|
package ratelimiter
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
GC_SIZE int = 100
|
||
|
|
GC_PERIOD time.Duration = 60 * time.Second
|
||
|
|
)
|
||
|
|
|
||
|
|
type Memory struct {
|
||
|
|
store map[string]LeakyBucket
|
||
|
|
lastGCCollected time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewMemory() *Memory {
|
||
|
|
m := new(Memory)
|
||
|
|
m.store = make(map[string]LeakyBucket)
|
||
|
|
m.lastGCCollected = time.Now()
|
||
|
|
return m
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *Memory) GetBucketFor(key string) (*LeakyBucket, error) {
|
||
|
|
|
||
|
|
bucket, ok := m.store[key]
|
||
|
|
if !ok {
|
||
|
|
return nil, errors.New("miss")
|
||
|
|
}
|
||
|
|
|
||
|
|
return &bucket, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *Memory) SetBucketFor(key string, bucket LeakyBucket) error {
|
||
|
|
|
||
|
|
if len(m.store) > GC_SIZE {
|
||
|
|
m.GarbageCollect()
|
||
|
|
}
|
||
|
|
|
||
|
|
m.store[key] = bucket
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *Memory) GarbageCollect() {
|
||
|
|
now := time.Now()
|
||
|
|
|
||
|
|
// rate limit GC to once per minute
|
||
|
|
if now.Unix() >= m.lastGCCollected.Add(GC_PERIOD).Unix() {
|
||
|
|
for key, bucket := range m.store {
|
||
|
|
// if the bucket is drained, then GC
|
||
|
|
if bucket.DrainedAt().Unix() < now.Unix() {
|
||
|
|
delete(m.store, key)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
m.lastGCCollected = now
|
||
|
|
}
|
||
|
|
}
|