diff options
Diffstat (limited to 'Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store')
5 files changed, 286 insertions, 0 deletions
diff --git a/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/doc.go b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/doc.go new file mode 100644 index 000000000..adb4618d3 --- /dev/null +++ b/Godeps/_workspace/src/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/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..74d48a73e --- /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/mattermost/platform/Godeps/_workspace/src/github.com/golang/groupcache/lru" + "github.com/mattermost/platform/Godeps/_workspace/src/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/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem_test.go b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem_test.go new file mode 100644 index 000000000..e8ef8d0da --- /dev/null +++ b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/mem_test.go @@ -0,0 +1,43 @@ +package store + +import ( + "testing" + "time" +) + +func TestMemStore(t *testing.T) { + st := NewMemStore(0) + win := time.Second + + // Reset stores a key with count of 1, current timestamp + err := st.Reset("k", time.Second) + if err != nil { + t.Errorf("expected reset to return nil, got %s", err) + } + cnt, sec1, _ := st.Incr("k", win) + if cnt != 2 { + t.Errorf("expected reset+incr to set count to 2, got %d", cnt) + } + + // Incr increments the key, keeps same timestamp + cnt, sec2, err := st.Incr("k", win) + if err != nil { + t.Errorf("expected 2nd incr to return nil error, got %s", err) + } + if cnt != 3 { + t.Errorf("expected 2nd incr to return 3, got %d", cnt) + } + if sec1 != sec2 { + t.Errorf("expected 2nd incr to return %d secs, got %d", sec1, sec2) + } + + // Reset on existing key brings it back to 1, new timestamp + err = st.Reset("k", win) + if err != nil { + t.Errorf("expected reset on existing key to return nil, got %s", err) + } + cnt, _, _ = st.Incr("k", win) + if cnt != 2 { + t.Errorf("expected last reset+incr to return 2, got %d", cnt) + } +} diff --git a/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis.go b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis.go new file mode 100644 index 000000000..44ef43122 --- /dev/null +++ b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis.go @@ -0,0 +1,85 @@ +package store + +import ( + "time" + + "github.com/mattermost/platform/Godeps/_workspace/src/github.com/garyburd/redigo/redis" + "github.com/mattermost/platform/Godeps/_workspace/src/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 +} diff --git a/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis_test.go b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis_test.go new file mode 100644 index 000000000..d7d1d1743 --- /dev/null +++ b/Godeps/_workspace/src/gopkg.in/throttled/throttled.v1/store/redis_test.go @@ -0,0 +1,66 @@ +package store + +import ( + "testing" + "time" + + "github.com/mattermost/platform/Godeps/_workspace/src/github.com/garyburd/redigo/redis" +) + +func getPool() *redis.Pool { + pool := &redis.Pool{ + MaxIdle: 3, + IdleTimeout: 30 * time.Second, + Dial: func() (redis.Conn, error) { + return redis.Dial("tcp", ":6379") + }, + TestOnBorrow: func(c redis.Conn, t time.Time) error { + _, err := c.Do("PING") + return err + }, + } + return pool +} + +func TestRedisStore(t *testing.T) { + pool := getPool() + c := pool.Get() + if _, err := redis.String(c.Do("PING")); err != nil { + c.Close() + t.Skip("redis server not available on localhost port 6379") + } + st := NewRedisStore(pool, "throttled:", 1) + win := 2 * time.Second + + // Incr increments the key, even if it does not exist + cnt, secs, err := st.Incr("k", win) + if err != nil { + t.Errorf("expected initial incr to return nil error, got %s", err) + } + if cnt != 1 { + t.Errorf("expected initial incr to return 1, got %d", cnt) + } + if secs != int(win.Seconds()) { + t.Errorf("expected initial incr to return %d secs, got %d", int(win.Seconds()), secs) + } + + // Waiting a second diminishes the remaining seconds + time.Sleep(time.Second) + _, sec2, _ := st.Incr("k", win) + if sec2 != secs-1 { + t.Errorf("expected 2nd incr after a 1s sleep to return %d secs, got %d", secs-1, sec2) + } + + // Waiting a second so the key expires, Incr should set back to 1, initial secs + time.Sleep(1100 * time.Millisecond) + cnt, sec3, err := st.Incr("k", win) + if err != nil { + t.Errorf("expected last incr to return nil error, got %s", err) + } + if cnt != 1 { + t.Errorf("expected last incr to return 1, got %d", cnt) + } + if sec3 != int(win.Seconds()) { + t.Errorf("expected last incr to return %d secs, got %d", int(win.Seconds()), sec3) + } +} |