435c0d9fc8
This PR switches the Nomad repository from using govendor to Go modules for managing dependencies. Aspects of the Nomad workflow remain pretty much the same. The usual Makefile targets should continue to work as they always did. The API submodule simply defers to the parent Nomad version on the repository, keeping the semantics of API versioning that currently exists.
61 lines
1 KiB
Go
61 lines
1 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
|
|
}
|
|
}
|