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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
|
package store
import (
"sync"
"time"
"github.com/golang/groupcache/lru"
"gopkg.in/throttled/throttled.v1"
)
// memStore implements an in-memory Store.
type memStore struct {
sync.Mutex
keys *lru.Cache
m map[string]*counter
}
// NewMemStore creates a new MemStore. If maxKeys > 0, the number of different keys
// is restricted to the specified amount. In this case, it uses an LRU algorithm to
// evict older keys to make room for newer ones. If a request is made for a key that
// has been evicted, it will be processed as if its count was 0, possibly allowing requests
// that should be denied.
//
// If maxKeys <= 0, there is no limit on the number of keys, which may use an unbounded amount of
// memory depending on the server's load.
//
// The MemStore is only for single-process rate-limiting. To share the rate limit state
// among multiple instances of the web server, use a database- or key-value-based
// store.
//
func NewMemStore(maxKeys int) throttled.Store {
var m *memStore
if maxKeys > 0 {
m = &memStore{
keys: lru.New(maxKeys),
}
} else {
m = &memStore{
m: make(map[string]*counter),
}
}
return m
}
// A counter represents a single entry in the MemStore.
type counter struct {
n int
ts time.Time
}
// Incr increments the counter for the specified key. It returns the new
// count value and the remaining number of seconds, or an error.
func (ms *memStore) Incr(key string, window time.Duration) (int, int, error) {
ms.Lock()
defer ms.Unlock()
var c *counter
if ms.keys != nil {
v, _ := ms.keys.Get(key)
if v != nil {
c = v.(*counter)
}
} else {
c = ms.m[key]
}
if c == nil {
c = &counter{0, time.Now().UTC()}
}
c.n++
if ms.keys != nil {
ms.keys.Add(key, c)
} else {
ms.m[key] = c
}
return c.n, throttled.RemainingSeconds(c.ts, window), nil
}
// Reset resets the counter for the specified key. It sets the count
// to 1 and initializes the timestamp with the current time, in UTC.
// It returns an error if the operation fails.
func (ms *memStore) Reset(key string, win time.Duration) error {
ms.Lock()
defer ms.Unlock()
c := &counter{1, time.Now().UTC()}
if ms.keys != nil {
ms.keys.Add(key, c)
} else {
ms.m[key] = c
}
return nil
}
|
