1package ratelimiter
 2
 3import (
 4	"errors"
 5	"time"
 6)
 7
 8const (
 9	GC_SIZE   int           = 100
10	GC_PERIOD time.Duration = 60 * time.Second
11)
12
13type Memory struct {
14	store           map[string]LeakyBucket
15	lastGCCollected time.Time
16}
17
18func NewMemory() *Memory {
19	m := new(Memory)
20	m.store = make(map[string]LeakyBucket)
21	m.lastGCCollected = time.Now()
22	return m
23}
24
25func (m *Memory) GetBucketFor(key string) (*LeakyBucket, error) {
26
27	bucket, ok := m.store[key]
28	if !ok {
29		return nil, errors.New("miss")
30	}
31
32	return &bucket, nil
33}
34
35func (m *Memory) SetBucketFor(key string, bucket LeakyBucket) error {
36
37	if len(m.store) > GC_SIZE {
38		m.GarbageCollect()
39	}
40
41	m.store[key] = bucket
42
43	return nil
44}
45
46func (m *Memory) GarbageCollect() {
47	now := time.Now()
48
49	// rate limit GC to once per minute
50	if now.Unix() >= m.lastGCCollected.Add(GC_PERIOD).Unix() {
51		for key, bucket := range m.store {
52			// if the bucket is drained, then GC
53			if bucket.DrainedAt().Unix() < now.Unix() {
54				delete(m.store, key)
55			}
56		}
57
58		m.lastGCCollected = now
59	}
60}