summaryrefslogtreecommitdiff
path: root/vendor/github.com/nxadm/tail/ratelimiter/memory.go
blob: bf3c2131b1e9a8d0a53db0b9260a047978f8eae9 (plain)
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
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
	}
}