Discharge7443/fleet
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Use a redis cluster-friendly store for rate limit (#2577)

Browse Source
This commit is contained in:
Martin Angers
2021-10-19 16:24:13 -04:00
committed by GitHub
parent 1de7b8c4ca
commit a7420140ce
7 changed files with 299 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
changes/fix-rate-limiter-redis-store-in-cluster-mode
+1
View File
@@ -0,0 +1 @@
* Fix the rate limiter Redis store when using Redis Cluster.
cmd/fleet/serve.go
+5 -4
View File
@@ -46,7 +46,6 @@ import (
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/spf13/cobra"
"github.com/throttled/throttled/v2/store/redigostore"
"google.golang.org/grpc"
)
@@ -216,6 +215,8 @@ the way that the Fleet server works.
if err != nil {
initFatal(err, "initialize Redis")
}
level.Info(logger).Log("component", "redis", "mode", redisPool.Mode())
ds = cached_mysql.New(ds, redisPool)
resultStore := pubsub.NewRedisQueryResults(redisPool, config.Redis.DuplicateResults)
liveQueryStore := live_query.NewRedisLiveQuery(redisPool)
@@ -278,9 +279,9 @@ the way that the Fleet server works.
httpLogger := kitlog.With(logger, "component", "http")
limiterStore, err := redigostore.New(redisPool, "ratelimit::", 0)
if err != nil {
initFatal(err, "initialize rate limit store")
limiterStore := &redis.ThrottledStore{
Pool: redisPool,
KeyPrefix: "ratelimit::",
}
var apiHandler, frontendHandler http.Handler
server/datastore/redis/ratelimit_store.go
+122
View File
@@ -0,0 +1,122 @@
package redis
import (
"strings"
"time"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/gomodule/redigo/redis"
)
type ThrottledStore struct {
Pool fleet.RedisPool
KeyPrefix string
}
const (
getWithTimeScript = `
local tbl = redis.call('TIME')
local val = redis.call('GET', KEYS[1])
table.insert(tbl, val)
return tbl
`
compareAndSwapWithTTLScript = `
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
redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[3])
return 1
`
compareAndSwapNoKeyError = "key does not exist"
)
func (s *ThrottledStore) GetWithTime(key string) (int64, time.Time, error) {
var t time.Time
key = s.KeyPrefix + key
conn := s.Pool.Get()
defer conn.Close()
if err := BindConn(s.Pool, conn, key); err != nil {
return 0, t, err
}
// must come after BindConn due to redisc restrictions
conn = ConfigureDoer(s.Pool, conn)
script := redis.NewScript(1, getWithTimeScript)
res, err := redis.Values(script.Do(conn, key))
if err != nil {
return 0, t, err
}
if len(res) < 3 {
res = append(res, nil)
}
var secs, us, val int64
val = -1 // initialize val to -1, will stay untouched if res[2] is nil
if _, err := redis.Scan(res, &secs, &us, &val); err != nil {
return 0, t, err
}
t = time.Unix(secs, us*int64(time.Microsecond))
return val, t, nil
}
func (s *ThrottledStore) SetIfNotExistsWithTTL(key string, value int64, ttl time.Duration) (bool, error) {
key = s.KeyPrefix + key
conn := ConfigureDoer(s.Pool, s.Pool.Get())
defer conn.Close()
ttlSeconds := int(ttl.Seconds())
// An `EX 0` will fail, make sure that we set expiry for a minimum of one second
if ttlSeconds < 1 {
ttlSeconds = 1
}
_, err := redis.String(conn.Do("SET", key, value, "EX", ttlSeconds, "NX"))
if err != nil {
if err == redis.ErrNil {
// not set due to NX condition not met
return false, nil
}
return false, err
}
return true, nil
}
func (s *ThrottledStore) CompareAndSwapWithTTL(key string, old, new int64, ttl time.Duration) (bool, error) {
key = s.KeyPrefix + key
conn := s.Pool.Get()
defer conn.Close()
if err := BindConn(s.Pool, conn, key); err != nil {
return false, err
}
// must come after BindConn due to redisc restrictions
conn = ConfigureDoer(s.Pool, conn)
ttlSeconds := int(ttl.Seconds())
// An `EX 0` will fail, make sure that we set expiry for a minimum of one second
if ttlSeconds < 1 {
ttlSeconds = 1
}
script := redis.NewScript(1, compareAndSwapWithTTLScript)
swapped, err := redis.Bool(script.Do(conn, key, old, new, ttlSeconds))
if err != nil {
if strings.Contains(err.Error(), compareAndSwapNoKeyError) {
return false, nil
}
return false, err
}
return swapped, nil
}
server/datastore/redis/ratelimit_store_test.go
+127
View File
@@ -0,0 +1,127 @@
package redis_test
import (
"testing"
"time"
"github.com/fleetdm/fleet/v4/server/datastore/redis"
"github.com/fleetdm/fleet/v4/server/datastore/redis/redistest"
"github.com/fleetdm/fleet/v4/server/fleet"
redigo "github.com/gomodule/redigo/redis"
"github.com/stretchr/testify/require"
)
func TestThrottledStore(t *testing.T) {
const prefix = "TestThrottledStore:"
runTest := func(t *testing.T, pool fleet.RedisPool) {
store := redis.ThrottledStore{
Pool: pool,
KeyPrefix: prefix,
}
t.Run("GetWithTime", func(t *testing.T) {
conn := pool.Get()
defer conn.Close()
// key does not exist
v, ts, err := store.GetWithTime("a")
require.NoError(t, err)
require.Equal(t, v, int64(-1))
require.WithinDuration(t, time.Now(), ts, time.Second)
_, err = conn.Do("SET", store.KeyPrefix+"a", 1)
require.NoError(t, err)
// key exists
v, ts, err = store.GetWithTime("a")
require.NoError(t, err)
require.Equal(t, v, int64(1))
require.WithinDuration(t, time.Now(), ts, time.Second)
})
t.Run("SetIfNotExistsWithTTL", func(t *testing.T) {
conn := pool.Get()
defer conn.Close()
// key does not exist
ok, err := store.SetIfNotExistsWithTTL("{b}", 1, time.Second)
require.NoError(t, err)
require.True(t, ok)
v, err := redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}"))
require.NoError(t, err)
require.Equal(t, v, 1)
// key exists
ok, err = store.SetIfNotExistsWithTTL("{b}", 2, time.Second)
require.NoError(t, err)
require.False(t, ok)
// value is still 1
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}"))
require.NoError(t, err)
require.Equal(t, v, 1)
// key does not exist, but ttl less than a second
ok, err = store.SetIfNotExistsWithTTL("{b}2", 3, time.Millisecond)
require.NoError(t, err)
require.True(t, ok)
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}2"))
require.NoError(t, err)
require.Equal(t, v, 3)
})
t.Run("CompareAndSwapWithTTL", func(t *testing.T) {
conn := pool.Get()
defer conn.Close()
// key does not exist
ok, err := store.CompareAndSwapWithTTL("{c}", 1, 2, time.Second)
require.NoError(t, err)
require.False(t, ok)
_, err = conn.Do("SET", store.KeyPrefix+"{c}", 1)
require.NoError(t, err)
// key exists, but values do not match
ok, err = store.CompareAndSwapWithTTL("{c}", 2, 3, time.Second)
require.NoError(t, err)
require.False(t, ok)
// key exists, values match
ok, err = store.CompareAndSwapWithTTL("{c}", 1, 4, time.Second)
require.NoError(t, err)
require.True(t, ok)
v, err := redigo.Int(conn.Do("GET", store.KeyPrefix+"{c}"))
require.NoError(t, err)
require.Equal(t, v, 4)
// key exists, ttl less than a second
ok, err = store.CompareAndSwapWithTTL("{c}", 4, 5, time.Millisecond)
require.NoError(t, err)
require.True(t, ok)
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{c}"))
require.NoError(t, err)
require.Equal(t, v, 5)
})
}
t.Run("standalone", func(t *testing.T) {
pool := redistest.SetupRedis(t, false, false, false)
runTest(t, pool)
})
t.Run("cluster", func(t *testing.T) {
pool := redistest.SetupRedis(t, true, true, false)
runTest(t, pool)
})
t.Run("cluster_nofollow", func(t *testing.T) {
pool := redistest.SetupRedis(t, true, false, false)
runTest(t, pool)
})
}
server/datastore/redis/redis.go
+8
View File
@@ -25,12 +25,20 @@ func (p *standalonePool) Stats() map[string]redis.PoolStats {
}
}
func (p *standalonePool) Mode() fleet.RedisMode {
return fleet.RedisStandalone
}
type clusterPool struct {
*redisc.Cluster
followRedirs bool
readReplica bool
}
func (p *clusterPool) Mode() fleet.RedisMode {
return fleet.RedisCluster
}
// PoolConfig holds the redis pool configuration options.
type PoolConfig struct {
Server string
server/datastore/redis/redis_external_test.go
+12
View File
@@ -270,3 +270,15 @@ func TestReadOnlyConn(t *testing.T) {
require.Contains(t, err.Error(), "MOVED")
})
}
func TestRedisMode(t *testing.T) {
t.Run("standalone", func(t *testing.T) {
pool := redistest.SetupRedis(t, false, false, false)
require.Equal(t, pool.Mode(), fleet.RedisStandalone)
})
t.Run("cluster", func(t *testing.T) {
pool := redistest.SetupRedis(t, true, false, false)
require.Equal(t, pool.Mode(), fleet.RedisCluster)
})
}
server/fleet/redis_pool.go
+24
View File
@@ -13,4 +13,28 @@ type RedisPool interface {
// Stats returns a map of redis pool statistics for each server address.
Stats() map[string]redis.PoolStats
// Mode returns the mode in which Redis is running.
Mode() RedisMode
}
// RedisMode indicates the mode in which Redis is running.
type RedisMode byte
// List of supported Redis modes.
const (
RedisStandalone RedisMode = iota
RedisCluster
)
// String returns the string representation of the Redis mode.
func (m RedisMode) String() string {
switch m {
case RedisStandalone:
return "standalone"
case RedisCluster:
return "cluster"
default:
return "unknown"
}
}