forked from ulule/limiter
/
store_memory.go
73 lines (56 loc) · 1.53 KB
/
store_memory.go
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
package limiter
import (
"fmt"
"time"
cache "github.com/pmylund/go-cache"
)
// MemoryStore is the in-memory store.
type MemoryStore struct {
Cache *cache.Cache
Prefix string
}
// NewMemoryStore creates a new instance of memory store with defaults.
func NewMemoryStore() Store {
return NewMemoryStoreWithOptions(StoreOptions{
Prefix: DefaultPrefix,
CleanUpInterval: DefaultCleanUpInterval,
})
}
// NewMemoryStoreWithOptions creates a new instance of memory store with options.
func NewMemoryStoreWithOptions(options StoreOptions) Store {
return &MemoryStore{
Prefix: options.Prefix,
Cache: cache.New(cache.NoExpiration, options.CleanUpInterval),
}
}
// Get implement Store.Get() method.
func (s *MemoryStore) Get(key string, rate Rate) (Context, error) {
ctx := Context{}
key = fmt.Sprintf("%s:%s", s.Prefix, key)
item, found := s.Cache.Items()[key]
ms := int64(time.Millisecond)
if !found || item.Expired() {
s.Cache.Set(key, int64(1), rate.Period)
return Context{
Limit: rate.Limit,
Remaining: rate.Limit - 1,
Reset: (time.Now().UnixNano()/ms + int64(rate.Period)/ms) / 1000,
Reached: false,
}, nil
}
err := s.Cache.Increment(key, int64(1))
if err != nil {
return ctx, nil
}
remaining := int64(0)
count := item.Object.(int64)
if count < rate.Limit {
remaining = rate.Limit - count
}
return Context{
Limit: rate.Limit,
Remaining: remaining,
Reset: time.Now().Add(time.Duration(item.Expiration) * time.Second).Unix(),
Reached: count > rate.Limit,
}, nil
}