summaryrefslogtreecommitdiffstats
path: root/vendor/gopkg.in/throttled/throttled.v1/store
diff options
context:
space:
mode:
Diffstat (limited to 'vendor/gopkg.in/throttled/throttled.v1/store')
-rw-r--r--vendor/gopkg.in/throttled/throttled.v1/store/doc.go2
-rw-r--r--vendor/gopkg.in/throttled/throttled.v1/store/mem.go90
-rw-r--r--vendor/gopkg.in/throttled/throttled.v1/store/redis.go85
3 files changed, 177 insertions, 0 deletions
diff --git a/vendor/gopkg.in/throttled/throttled.v1/store/doc.go b/vendor/gopkg.in/throttled/throttled.v1/store/doc.go
new file mode 100644
index 000000000..adb4618d3
--- /dev/null
+++ b/vendor/gopkg.in/throttled/throttled.v1/store/doc.go
@@ -0,0 +1,2 @@
+// Package store offers a memory-based and a Redis-based throttled.Store implementation.
+package store
diff --git a/vendor/gopkg.in/throttled/throttled.v1/store/mem.go b/vendor/gopkg.in/throttled/throttled.v1/store/mem.go
new file mode 100644
index 000000000..22d200e8d
--- /dev/null
+++ b/vendor/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
+}
diff --git a/vendor/gopkg.in/throttled/throttled.v1/store/redis.go b/vendor/gopkg.in/throttled/throttled.v1/store/redis.go
new file mode 100644
index 000000000..b089f9f4e
--- /dev/null
+++ b/vendor/gopkg.in/throttled/throttled.v1/store/redis.go
@@ -0,0 +1,85 @@
+package store
+
+import (
+ "time"
+
+ "github.com/garyburd/redigo/redis"
+ "gopkg.in/throttled/throttled.v1"
+)
+
+// redisStore implements a Redis-based store.
+type redisStore struct {
+ pool *redis.Pool
+ prefix string
+ db int
+}
+
+// NewRedisStore creates a new Redis-based store, using the provided pool to get its
+// connections. The keys will have the specified keyPrefix, which may be an empty string,
+// and the database index specified by db will be selected to store the keys.
+//
+func NewRedisStore(pool *redis.Pool, keyPrefix string, db int) throttled.Store {
+ return &redisStore{
+ pool: pool,
+ prefix: keyPrefix,
+ db: db,
+ }
+}
+
+// Incr increments the specified key. If the key did not exist, it sets it to 1
+// and sets it to expire after the number of seconds specified by window.
+//
+// It returns the new count value and the number of remaining seconds, or an error
+// if the operation fails.
+func (r *redisStore) Incr(key string, window time.Duration) (int, int, error) {
+ conn := r.pool.Get()
+ defer conn.Close()
+ if err := selectDB(r.db, conn); err != nil {
+ return 0, 0, err
+ }
+ // Atomically increment and read the TTL.
+ conn.Send("MULTI")
+ conn.Send("INCR", r.prefix+key)
+ conn.Send("TTL", r.prefix+key)
+ vals, err := redis.Values(conn.Do("EXEC"))
+ if err != nil {
+ conn.Do("DISCARD")
+ return 0, 0, err
+ }
+ var cnt, ttl int
+ if _, err = redis.Scan(vals, &cnt, &ttl); err != nil {
+ return 0, 0, err
+ }
+ // If there was no TTL set, then this is a newly created key (INCR creates the key
+ // if it didn't exist), so set it to expire.
+ if ttl == -1 {
+ ttl = int(window.Seconds())
+ _, err = conn.Do("EXPIRE", r.prefix+key, ttl)
+ if err != nil {
+ return 0, 0, err
+ }
+ }
+ return cnt, ttl, nil
+}
+
+// Reset sets the value of the key to 1, and resets its time window.
+func (r *redisStore) Reset(key string, window time.Duration) error {
+ conn := r.pool.Get()
+ defer conn.Close()
+ if err := selectDB(r.db, conn); err != nil {
+ return err
+ }
+ _, err := redis.String(conn.Do("SET", r.prefix+key, "1", "EX", int(window.Seconds()), "NX"))
+ return err
+}
+
+// Select the specified database index.
+func selectDB(db int, conn redis.Conn) error {
+ // Select the specified database
+ if db > 0 {
+ if _, err := redis.String(conn.Do("SELECT", db)); err != nil {
+ return err
+ }
+ }
+ return nil
+}