diff options
Diffstat (limited to 'vendor/gopkg.in/throttled/throttled.v2/store')
7 files changed, 618 insertions, 0 deletions
diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/deprecated.go b/vendor/gopkg.in/throttled/throttled.v2/store/deprecated.go new file mode 100644 index 000000000..5476e87ac --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/deprecated.go @@ -0,0 +1,32 @@ +// Package store contains deprecated aliases for subpackages +package store // import "gopkg.in/throttled/throttled.v2/store" + +import ( + "github.com/garyburd/redigo/redis" + + "gopkg.in/throttled/throttled.v2/store/memstore" + "gopkg.in/throttled/throttled.v2/store/redigostore" +) + +// DEPRECATED. NewMemStore is a compatible alias for mem.New +func NewMemStore(maxKeys int) *memstore.MemStore { + st, err := memstore.New(maxKeys) + if err != nil { + // As of this writing, `lru.New` can only return an error if you pass + // maxKeys <= 0 so this should never occur. + panic(err) + } + return st +} + +// DEPRECATED. NewMemStore is a compatible alias for redis.New +func NewRedisStore(pool *redis.Pool, keyPrefix string, db int) *redigostore.RedigoStore { + st, err := redigostore.New(pool, keyPrefix, db) + if err != nil { + // As of this writing, creating a Redis store never returns an error + // so this should be safe while providing some ability to return errors + // in the future. + panic(err) + } + return st +} diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore.go b/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore.go new file mode 100644 index 000000000..5d8fee8b5 --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore.go @@ -0,0 +1,127 @@ +// Package memstore offers an in-memory store implementation for throttled. +package memstore // import "gopkg.in/throttled/throttled.v2/store/memstore" + +import ( + "sync" + "sync/atomic" + "time" + + "github.com/hashicorp/golang-lru" +) + +// MemStore is an in-memory store implementation for throttled. It +// supports evicting the least recently used keys to control memory +// usage. It is stored in memory in the current process and thus +// doesn't share state with other rate limiters. +type MemStore struct { + sync.RWMutex + keys *lru.Cache + m map[string]*int64 +} + +// New initializes a Store. 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 maxKeys <= 0, there is no limit on the number of keys, +// which may use an unbounded amount of memory. +func New(maxKeys int) (*MemStore, error) { + var m *MemStore + + if maxKeys > 0 { + keys, err := lru.New(maxKeys) + if err != nil { + return nil, err + } + + m = &MemStore{ + keys: keys, + } + } else { + m = &MemStore{ + m: make(map[string]*int64), + } + } + return m, nil +} + +// GetWithTime returns the value of the key if it is in the store or +// -1 if it does not exist. It also returns the current local time on +// the machine. +func (ms *MemStore) GetWithTime(key string) (int64, time.Time, error) { + now := time.Now() + valP, ok := ms.get(key, false) + + if !ok { + return -1, now, nil + } + + return atomic.LoadInt64(valP), now, nil +} + +// SetIfNotExistsWithTTL sets the value of key only if it is not +// already set in the store it returns whether a new value was set. It +// ignores the ttl. +func (ms *MemStore) SetIfNotExistsWithTTL(key string, value int64, _ time.Duration) (bool, error) { + _, ok := ms.get(key, false) + + if ok { + return false, nil + } + + ms.Lock() + defer ms.Unlock() + + _, ok = ms.get(key, true) + + if ok { + return false, nil + } + + // Store a pointer to a new instance so that the caller + // can't mutate the value after setting + v := value + + if ms.keys != nil { + ms.keys.Add(key, &v) + } else { + ms.m[key] = &v + } + + return true, nil +} + +// CompareAndSwapWithTTL atomically compares the value at key to the +// old value. If it matches, it sets it to the new value and returns +// true. Otherwise, it returns false. If the key does not exist in the +// store, it returns false with no error. It ignores the ttl. +func (ms *MemStore) CompareAndSwapWithTTL(key string, old, new int64, _ time.Duration) (bool, error) { + valP, ok := ms.get(key, false) + + if !ok { + return false, nil + } + + return atomic.CompareAndSwapInt64(valP, old, new), nil +} + +func (ms *MemStore) get(key string, locked bool) (*int64, bool) { + var valP *int64 + var ok bool + + if ms.keys != nil { + var valI interface{} + + valI, ok = ms.keys.Get(key) + if ok { + valP = valI.(*int64) + } + } else { + if !locked { + ms.RLock() + defer ms.RUnlock() + } + valP, ok = ms.m[key] + } + + return valP, ok +} diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore_test.go b/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore_test.go new file mode 100644 index 000000000..ef003d3de --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/memstore/memstore_test.go @@ -0,0 +1,40 @@ +package memstore_test + +import ( + "testing" + + "gopkg.in/throttled/throttled.v2/store/memstore" + "gopkg.in/throttled/throttled.v2/store/storetest" +) + +func TestMemStoreLRU(t *testing.T) { + st, err := memstore.New(10) + if err != nil { + t.Fatal(err) + } + storetest.TestGCRAStore(t, st) +} + +func TestMemStoreUnlimited(t *testing.T) { + st, err := memstore.New(10) + if err != nil { + t.Fatal(err) + } + storetest.TestGCRAStore(t, st) +} + +func BenchmarkMemStoreLRU(b *testing.B) { + st, err := memstore.New(10) + if err != nil { + b.Fatal(err) + } + storetest.BenchmarkGCRAStore(b, st) +} + +func BenchmarkMemStoreUnlimited(b *testing.B) { + st, err := memstore.New(0) + if err != nil { + b.Fatal(err) + } + storetest.BenchmarkGCRAStore(b, st) +} diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redigostore.go b/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redigostore.go new file mode 100644 index 000000000..54208fa6d --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redigostore.go @@ -0,0 +1,156 @@ +// Package redigostore offers Redis-based store implementation for throttled using redigo. +package redigostore // import "gopkg.in/throttled/throttled.v2/store/redigostore" + +import ( + "strings" + "time" + + "github.com/garyburd/redigo/redis" +) + +const ( + redisCASMissingKey = "key does not exist" + redisCASScript = ` +local v = redis.call('get', KEYS[1]) +if v == false then + return redis.error_reply("key does not exist") +end +if v ~= ARGV[1] then + return 0 +end +if ARGV[3] ~= "0" then + redis.call('setex', KEYS[1], ARGV[3], ARGV[2]) +else + redis.call('set', KEYS[1], ARGV[2]) +end +return 1 +` +) + +// RedigoStore implements a Redis-based store using redigo. +type RedigoStore struct { + pool *redis.Pool + prefix string + db int +} + +// New 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. Any updating operations will reset +// the key TTL to the provided value rounded down to the nearest +// second. Depends on Redis 2.6+ for EVAL support. +func New(pool *redis.Pool, keyPrefix string, db int) (*RedigoStore, error) { + return &RedigoStore{ + pool: pool, + prefix: keyPrefix, + db: db, + }, nil +} + +// GetWithTime returns the value of the key if it is in the store +// or -1 if it does not exist. It also returns the current time at +// the redis server to microsecond precision. +func (r *RedigoStore) GetWithTime(key string) (int64, time.Time, error) { + var now time.Time + + key = r.prefix + key + + conn, err := r.getConn() + if err != nil { + return 0, now, err + } + defer conn.Close() + + conn.Send("TIME") + conn.Send("GET", key) + conn.Flush() + timeReply, err := redis.Values(conn.Receive()) + if err != nil { + return 0, now, err + } + + var s, us int64 + if _, err := redis.Scan(timeReply, &s, &us); err != nil { + return 0, now, err + } + now = time.Unix(s, us*int64(time.Microsecond)) + + v, err := redis.Int64(conn.Receive()) + if err == redis.ErrNil { + return -1, now, nil + } else if err != nil { + return 0, now, err + } + + return v, now, nil +} + +// SetIfNotExistsWithTTL sets the value of key only if it is not +// already set in the store it returns whether a new value was set. +// If a new value was set, the ttl in the key is also set, though this +// operation is not performed atomically. +func (r *RedigoStore) SetIfNotExistsWithTTL(key string, value int64, ttl time.Duration) (bool, error) { + key = r.prefix + key + + conn, err := r.getConn() + if err != nil { + return false, err + } + defer conn.Close() + + v, err := redis.Int64(conn.Do("SETNX", key, value)) + if err != nil { + return false, err + } + + updated := v == 1 + + if ttl >= time.Second { + if _, err := conn.Do("EXPIRE", key, int(ttl.Seconds())); err != nil { + return updated, err + } + } + + return updated, nil +} + +// CompareAndSwapWithTTL atomically compares the value at key to the +// old value. If it matches, it sets it to the new value and returns +// true. Otherwise, it returns false. If the key does not exist in the +// store, it returns false with no error. If the swap succeeds, the +// ttl for the key is updated atomically. +func (r *RedigoStore) CompareAndSwapWithTTL(key string, old, new int64, ttl time.Duration) (bool, error) { + key = r.prefix + key + conn, err := r.getConn() + if err != nil { + return false, err + } + defer conn.Close() + + swapped, err := redis.Bool(conn.Do("EVAL", redisCASScript, 1, key, old, new, int(ttl.Seconds()))) + if err != nil { + if strings.Contains(err.Error(), redisCASMissingKey) { + return false, nil + } + + return false, err + } + + return swapped, nil +} + +// Select the specified database index. +func (r *RedigoStore) getConn() (redis.Conn, error) { + conn := r.pool.Get() + + // Select the specified database + if r.db > 0 { + if _, err := redis.String(conn.Do("SELECT", r.db)); err != nil { + conn.Close() + return nil, err + } + } + + return conn, nil +} diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redisstore_test.go b/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redisstore_test.go new file mode 100644 index 000000000..d47b635d2 --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/redigostore/redisstore_test.go @@ -0,0 +1,85 @@ +package redigostore_test + +import ( + "testing" + "time" + + "github.com/garyburd/redigo/redis" + + "gopkg.in/throttled/throttled.v2/store/redigostore" + "gopkg.in/throttled/throttled.v2/store/storetest" +) + +const ( + redisTestDB = 1 + redisTestPrefix = "throttled:" +) + +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) { + c, st := setupRedis(t, 0) + defer c.Close() + defer clearRedis(c) + + clearRedis(c) + storetest.TestGCRAStore(t, st) + storetest.TestGCRAStoreTTL(t, st) +} + +func BenchmarkRedisStore(b *testing.B) { + c, st := setupRedis(b, 0) + defer c.Close() + defer clearRedis(c) + + storetest.BenchmarkGCRAStore(b, st) +} + +func clearRedis(c redis.Conn) error { + keys, err := redis.Values(c.Do("KEYS", redisTestPrefix+"*")) + if err != nil { + return err + } + + if _, err := redis.Int(c.Do("DEL", keys...)); err != nil { + return err + } + + return nil +} + +func setupRedis(tb testing.TB, ttl time.Duration) (redis.Conn, *redigostore.RedigoStore) { + pool := getPool() + c := pool.Get() + + if _, err := redis.String(c.Do("PING")); err != nil { + c.Close() + tb.Skip("redis server not available on localhost port 6379") + } + + if _, err := redis.String(c.Do("SELECT", redisTestDB)); err != nil { + c.Close() + tb.Fatal(err) + } + + st, err := redigostore.New(pool, redisTestPrefix, redisTestDB) + if err != nil { + c.Close() + tb.Fatal(err) + } + + return c, st +} diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/storetest/doc.go b/vendor/gopkg.in/throttled/throttled.v2/store/storetest/doc.go new file mode 100644 index 000000000..ecfee2638 --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/storetest/doc.go @@ -0,0 +1,2 @@ +// Package storetest provides a helper for testing throttled stores. +package storetest // import "gopkg.in/throttled/throttled.v2/store/storetest" diff --git a/vendor/gopkg.in/throttled/throttled.v2/store/storetest/storetest.go b/vendor/gopkg.in/throttled/throttled.v2/store/storetest/storetest.go new file mode 100644 index 000000000..191b40a4f --- /dev/null +++ b/vendor/gopkg.in/throttled/throttled.v2/store/storetest/storetest.go @@ -0,0 +1,176 @@ +// Package storetest provides a helper for testing throttled stores. +package storetest // import "gopkg.in/throttled/throttled.v2/store/storetest" + +import ( + "math/rand" + "strconv" + "sync/atomic" + "testing" + "time" + + "gopkg.in/throttled/throttled.v2" +) + +// TestGCRAStore tests the behavior of a GCRAStore implementation for +// compliance with the throttled API. It does not require support +// for TTLs. +func TestGCRAStore(t *testing.T, st throttled.GCRAStore) { + // GetWithTime a missing key + if have, _, err := st.GetWithTime("foo"); err != nil { + t.Fatal(err) + } else if have != -1 { + t.Errorf("expected GetWithTime to return -1 for a missing key but got %d", have) + } + + // SetIfNotExists on a new key + want := int64(1) + + if set, err := st.SetIfNotExistsWithTTL("foo", want, 0); err != nil { + t.Fatal(err) + } else if !set { + t.Errorf("expected SetIfNotExists on an empty key to succeed") + } + + before := time.Now() + + if have, now, err := st.GetWithTime("foo"); err != nil { + t.Fatal(err) + } else if have != want { + t.Errorf("expected GetWithTime to return %d but got %d", want, have) + } else if now.UnixNano() <= 0 { + t.Errorf("expected GetWithTime to return a time representable representable as a positive int64 of nanoseconds since the epoch") + } else if now.Before(before) || now.After(time.Now()) { + // Note that we make the assumption here that the store is running on + // the same machine as this test and thus shares a clock. This can be a + // little tricky in the case of Redis, which could be running + // elsewhere. The test assumes that it's running either locally on on + // Travis (where currently the Redis is available on localhost). If new + // test environments are procured, this may need to be revisited. + t.Errorf("expected GetWithTime to return a time between the time before the call and the time after the call") + } + + // SetIfNotExists on an existing key + if set, err := st.SetIfNotExistsWithTTL("foo", 123, 0); err != nil { + t.Fatal(err) + } else if set { + t.Errorf("expected SetIfNotExists on an existing key to fail") + } + + if have, _, err := st.GetWithTime("foo"); err != nil { + t.Fatal(err) + } else if have != want { + t.Errorf("expected GetWithTime to return %d but got %d", want, have) + } + + // SetIfNotExists on a different key + if set, err := st.SetIfNotExistsWithTTL("bar", 456, 0); err != nil { + t.Fatal(err) + } else if !set { + t.Errorf("expected SetIfNotExists on an empty key to succeed") + } + + // Returns the false on a missing key + if swapped, err := st.CompareAndSwapWithTTL("baz", 1, 2, 0); err != nil { + t.Fatal(err) + } else if swapped { + t.Errorf("expected CompareAndSwap to fail on a missing key") + } + + // Test a successful CAS + want = int64(2) + + if swapped, err := st.CompareAndSwapWithTTL("foo", 1, want, 0); err != nil { + t.Fatal(err) + } else if !swapped { + t.Errorf("expected CompareAndSwap to succeed") + } + + if have, _, err := st.GetWithTime("foo"); err != nil { + t.Fatal(err) + } else if have != want { + t.Errorf("expected GetWithTime to return %d but got %d", want, have) + } + + // Test an unsuccessful CAS + if swapped, err := st.CompareAndSwapWithTTL("foo", 1, 2, 0); err != nil { + t.Fatal(err) + } else if swapped { + t.Errorf("expected CompareAndSwap to fail") + } + + if have, _, err := st.GetWithTime("foo"); err != nil { + t.Fatal(err) + } else if have != want { + t.Errorf("expected GetWithTime to return %d but got %d", want, have) + } +} + +// TestGCRAStoreTTL tests the behavior of TTLs in a GCRAStore implementation. +func TestGCRAStoreTTL(t *testing.T, st throttled.GCRAStore) { + ttl := time.Second + want := int64(1) + key := "ttl" + + if _, err := st.SetIfNotExistsWithTTL(key, want, ttl); err != nil { + t.Fatal(err) + } + + if have, _, err := st.GetWithTime(key); err != nil { + t.Fatal(err) + } else if have != want { + t.Errorf("expected GetWithTime to return %d, got %d", want, have) + } + + // I can't think of a generic way to test expiration without a sleep + time.Sleep(ttl + time.Millisecond) + + if have, _, err := st.GetWithTime(key); err != nil { + t.Fatal(err) + } else if have != -1 { + t.Errorf("expected GetWithTime to fail on an expired key but got %d", have) + } +} + +// BenchmarkGCRAStore runs parallel benchmarks against a GCRAStore implementation. +// Aside from being useful for performance testing, this is useful for finding +// race conditions with the Go race detector. +func BenchmarkGCRAStore(b *testing.B, st throttled.GCRAStore) { + seed := int64(42) + var attempts, updates int64 + + b.RunParallel(func(pb *testing.PB) { + // We need atomic behavior around the RNG or go detects a race in the test + delta := int64(1) + seedValue := atomic.AddInt64(&seed, delta) - delta + gen := rand.New(rand.NewSource(seedValue)) + + for pb.Next() { + key := strconv.FormatInt(gen.Int63n(50), 10) + + var v int64 + var updated bool + + v, _, err := st.GetWithTime(key) + if v == -1 { + updated, err = st.SetIfNotExistsWithTTL(key, gen.Int63(), 0) + if err != nil { + b.Error(err) + } + } else if err != nil { + b.Error(err) + } else { + updated, err = st.CompareAndSwapWithTTL(key, v, gen.Int63(), 0) + if err != nil { + b.Error(err) + } + } + + atomic.AddInt64(&attempts, 1) + if updated { + atomic.AddInt64(&updates, 1) + } + } + }) + + b.Logf("%d/%d update operations succeeed", updates, attempts) +} |