diff options
Diffstat (limited to 'Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem.go')
-rw-r--r-- | Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem.go | 90 |
1 files changed, 90 insertions, 0 deletions
diff --git a/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem.go b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem.go new file mode 100644 index 000000000..22d200e8d --- /dev/null +++ b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem.go @@ -0,0 +1,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 +} |