diff options
author | Christopher Speller <crspeller@gmail.com> | 2016-05-12 15:08:58 -0400 |
---|---|---|
committer | Christopher Speller <crspeller@gmail.com> | 2016-05-12 16:37:29 -0400 |
commit | 84d2482ddbff9564c9ad75b2d30af66e3ddfd44d (patch) | |
tree | 8bfa567d2b6381f4a996ada2deff8a16aa85a3ac /vendor/gopkg.in/throttled/throttled.v1/store/redis.go | |
parent | d1efb66ad7b017f0fbfe6f0c20843b30f396e504 (diff) | |
download | chat-84d2482ddbff9564c9ad75b2d30af66e3ddfd44d.tar.gz chat-84d2482ddbff9564c9ad75b2d30af66e3ddfd44d.tar.bz2 chat-84d2482ddbff9564c9ad75b2d30af66e3ddfd44d.zip |
Updating go depencancies. Switching to go1.6 vendoring (#2949)
Diffstat (limited to 'vendor/gopkg.in/throttled/throttled.v1/store/redis.go')
-rw-r--r-- | vendor/gopkg.in/throttled/throttled.v1/store/redis.go | 85 |
1 files changed, 85 insertions, 0 deletions
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 +} |