Add Redis cluster support (#1045)

This should support Redis in both cluster and non-cluster modes.

Updates were made separately to github.com/throttled/throttled to support the slight changes in types.

Co-authored-by: Joseph Macaulay <joseph.macaulay@uber.com>
Co-authored-by: Zach Wasserman <zach@fleetdm.com>

cmd/fleet/serve.go

@@ -35,7 +35,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/spf13/cobra"
-	"github.com/throttled/throttled/store/redigostore"
+	"github.com/throttled/throttled/v2/store/redigostore"
 	"google.golang.org/grpc"
 )
 
@@ -215,7 +215,10 @@ the way that the Fleet server works.
 				}
 			}
 
-			redisPool := pubsub.NewRedisPool(config.Redis.Address, config.Redis.Password, config.Redis.Database, config.Redis.UseTLS)
+			redisPool, err := pubsub.NewRedisPool(config.Redis.Address, config.Redis.Password, config.Redis.Database, config.Redis.UseTLS)
+			if err != nil {
+				initFatal(err, "initialize Redis")
+			}
 			resultStore := pubsub.NewRedisQueryResults(redisPool, config.Redis.DuplicateResults)
 			liveQueryStore := live_query.NewRedisLiveQuery(redisPool)
 			ssoSessionStore := sso.NewSessionStore(redisPool)

go.mod

@@ -43,6 +43,7 @@ require (
 	github.com/mattn/go-isatty v0.0.12 // indirect
 	github.com/mattn/go-runewidth v0.0.8 // indirect
 	github.com/mitchellh/mapstructure v0.0.0-20180220230111-00c29f56e238 // indirect
+	github.com/mna/redisc v1.2.1
 	github.com/olekukonko/tablewriter v0.0.0-20180506121414-d4647c9c7a84
 	github.com/patrickmn/sortutil v0.0.0-20120526081524-abeda66eb583
 	github.com/pelletier/go-toml v1.1.0 // indirect
@@ -56,10 +57,9 @@ require (
 	github.com/spf13/jwalterweatherman v0.0.0-20180109140146-7c0cea34c8ec // indirect
 	github.com/spf13/pflag v1.0.1 // indirect
 	github.com/spf13/viper v1.0.2
-	github.com/stretchr/testify v1.6.1
+	github.com/stretchr/testify v1.7.0
 	github.com/theupdateframework/go-tuf v0.0.0-20201230183259-aee6270feb55
-	github.com/throttled/throttled v2.2.5+incompatible
-	github.com/throttled/throttled/v2 v2.7.2
+	github.com/throttled/throttled/v2 v2.8.0
 	github.com/urfave/cli/v2 v2.3.0
 	go.opencensus.io v0.20.2 // indirect
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
@@ -72,3 +72,4 @@ require (
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0-20170531160350-a96e63847dc3
 	gopkg.in/yaml.v2 v2.2.7
 )
+

go.sum

@@ -66,6 +66,8 @@ github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/fleetdm/goose v0.0.0-20210209032905-c3c01484bacb h1:p02npmJlTo+Px1s0VptKOJOJqH/rGlGBEVvLJRtzY3A=
 github.com/fleetdm/goose v0.0.0-20210209032905-c3c01484bacb/go.mod h1:d7Q+0eCENnKQUhkfAUVLfGnD4QcgJMF/uB9WRTN9TDI=
+github.com/fleetdm/throttled/v2 v2.7.3-0.20210611232555-e08fbdec1e43 h1:irhvqFd9Xx1vY7CpfYc0JQLiT6plyQOAXsL+wGfa7C0=
+github.com/fleetdm/throttled/v2 v2.7.3-0.20210611232555-e08fbdec1e43/go.mod h1:q1QyZVQXxb2NUfJ+Hjucmlrsrz9s/jt2ilMwSMo7a2I=
 github.com/flynn/go-docopt v0.0.0-20140912013429-f6dd2ebbb31e h1:Ss/B3/5wWRh8+emnK0++g5zQzwDTi30W10pKxKc4JXI=
 github.com/flynn/go-docopt v0.0.0-20140912013429-f6dd2ebbb31e/go.mod h1:HyVoz1Mz5Co8TFO8EupIdlcpwShBmY98dkT2xeHkvEI=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
@@ -195,6 +197,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0j
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mitchellh/mapstructure v0.0.0-20180220230111-00c29f56e238 h1:+MZW2uvHgN8kYvksEN3f7eFL2wpzk0GxmlFsMybWc7E=
 github.com/mitchellh/mapstructure v0.0.0-20180220230111-00c29f56e238/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
+github.com/mna/redisc v1.2.1 h1:7rI/qv2sa0OT8rsxDbKg7XPysr5AIDeXwL0T0vFOvlM=
+github.com/mna/redisc v1.2.1/go.mod h1:OxLEDNNDFOYJBo7MuSC+SEoP3k8bZY2dFW7T12TzX4c=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223 h1:F9x/1yl3T2AeKLr2AMdilSD8+f9bvMnNN8VS5iDtovc=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
@@ -266,6 +270,8 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/syndtr/goleveldb v1.0.0 h1:fBdIW9lB4Iz0n9khmH8w27SJ3QEJ7+IgjPEwGSZiFdE=
 github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
 github.com/tent/canonical-json-go v0.0.0-20130607151641-96e4ba3a7613 h1:iGnD/q9160NWqKZZ5vY4p0dMiYMRknzctfSkqA4nBDw=
@@ -276,6 +282,8 @@ github.com/throttled/throttled v2.2.5+incompatible h1:65UB52X0qNTYiT0Sohp8qLYVFw
 github.com/throttled/throttled v2.2.5+incompatible/go.mod h1:0BjlrEGQmvxps+HuXLsyRdqpSRvJpq0PNIsOtqP9Nos=
 github.com/throttled/throttled/v2 v2.7.2 h1:UHk0GmVH5Zk1YiY2Di/C8WdtJ9E0wjjfkEiuS/e0pns=
 github.com/throttled/throttled/v2 v2.7.2/go.mod h1:q1QyZVQXxb2NUfJ+Hjucmlrsrz9s/jt2ilMwSMo7a2I=
+github.com/throttled/throttled/v2 v2.8.0 h1:B5VfdM8BE+ClI2Ji238SbNOTWfYcocvuAhgT27lvwrE=
+github.com/throttled/throttled/v2 v2.8.0/go.mod h1:q1QyZVQXxb2NUfJ+Hjucmlrsrz9s/jt2ilMwSMo7a2I=
 github.com/urfave/cli/v2 v2.3.0 h1:qph92Y649prgesehzOrQjdWyxFOp/QVM+6imKHad91M=
 github.com/urfave/cli/v2 v2.3.0 h1:qph92Y649prgesehzOrQjdWyxFOp/QVM+6imKHad91M=
 github.com/urfave/cli/v2 v2.3.0/go.mod h1:LJmUH05zAU44vOAcrfzZQKsZbVcdbOG8rtL3/XcUArI=

server/live_query/redis_live_query.go

@@ -28,6 +28,7 @@ import (
 	"time"
 
 	"github.com/gomodule/redigo/redis"
+	"github.com/mna/redisc"
 	"github.com/pkg/errors"
 )
 
@@ -40,12 +41,12 @@ const (
 
 type redisLiveQuery struct {
 	// connection pool
-	pool *redis.Pool
+	pool *redisc.Cluster
 }
 
 // NewRedisQueryResults creats a new Redis implementation of the
 // QueryResultStore interface using the provided Redis connection pool.
-func NewRedisLiveQuery(pool *redis.Pool) *redisLiveQuery {
+func NewRedisLiveQuery(pool *redisc.Cluster) *redisLiveQuery {
 	return &redisLiveQuery{pool: pool}
 }
 

server/live_query/redis_live_query_test.go

@@ -32,9 +32,11 @@ func setupRedisLiveQuery(t *testing.T) (store *redisLiveQuery, teardown func())
 		useTLS   = false
 	)
 
-	store = NewRedisLiveQuery(pubsub.NewRedisPool(addr, password, database, useTLS))
+	pool, err := pubsub.NewRedisPool(addr, password, database, useTLS)
+	require.NoError(t, err)
+	store = NewRedisLiveQuery(pool)
 
-	_, err := store.pool.Get().Do("PING")
+	_, err = store.pool.Get().Do("PING")
 	require.NoError(t, err)
 
 	teardown = func() {

server/pubsub/query_results_test.go

@@ -68,9 +68,11 @@ func setupRedis(t *testing.T) (store *redisQueryResults, teardown func()) {
 		dupResults = false
 	)
 
-	store = NewRedisQueryResults(NewRedisPool(addr, password, database, useTLS), dupResults)
+	pool, err := NewRedisPool(addr, password, database, useTLS)
+	require.NoError(t, err)
+	store = NewRedisQueryResults(pool, dupResults)
 
-	_, err := store.pool.Get().Do("PING")
+	_, err = store.pool.Get().Do("PING")
 	require.Nil(t, err)
 
 	teardown = func() {

server/pubsub/redis_query_results.go

@@ -4,16 +4,18 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/fleetdm/fleet/server/kolide"
 	"github.com/gomodule/redigo/redis"
+	"github.com/mna/redisc"
 	"github.com/pkg/errors"
 )
 
 type redisQueryResults struct {
 	// connection pool
-	pool             *redis.Pool
+	pool             *redisc.Cluster
 	duplicateResults bool
 }
 
@@ -21,46 +23,63 @@ var _ kolide.QueryResultStore = &redisQueryResults{}
 
 // NewRedisPool creates a Redis connection pool using the provided server
 // address, password and database.
-func NewRedisPool(server, password string, database int, useTLS bool) *redis.Pool {
-	return &redis.Pool{
-		MaxIdle:     3,
-		IdleTimeout: 240 * time.Second,
-		Dial: func() (redis.Conn, error) {
-			c, err := redis.Dial(
-				"tcp",
-				server,
-				redis.DialDatabase(database),
-				redis.DialUseTLS(useTLS),
-				redis.DialConnectTimeout(5*time.Second),
-				redis.DialKeepAlive(10*time.Second),
-				// Read/Write timeouts not set here because we may see results
-				// only rarely on the pub/sub channel.
-			)
-
-			if err != nil {
-				return nil, err
-			}
-			if password != "" {
-				if _, err := c.Do("AUTH", password); err != nil {
-					c.Close()
-					return nil, err
-				}
-			}
-			return c, err
+func NewRedisPool(server, password string, database int, useTLS bool) (*redisc.Cluster, error) {
+	//Create the Cluster
+	cluster := &redisc.Cluster{
+		StartupNodes: []string{
+			fmt.Sprint(server),
 		},
-		TestOnBorrow: func(c redis.Conn, t time.Time) error {
-			if time.Since(t) < time.Minute {
-				return nil
-			}
-			_, err := c.Do("PING")
-			return err
+		CreatePool: func(server string, opts ...redis.DialOption) (*redis.Pool, error) {
+			return &redis.Pool{
+				MaxIdle:     3,
+				IdleTimeout: 240 * time.Second,
+				Dial: func() (redis.Conn, error) {
+					c, err := redis.Dial(
+						"tcp",
+						server,
+						redis.DialDatabase(database),
+						redis.DialUseTLS(useTLS),
+						redis.DialConnectTimeout(5*time.Second),
+						redis.DialKeepAlive(10*time.Second),
+						// Read/Write timeouts not set here because we may see results
+						// only rarely on the pub/sub channel.
+					)
+					if err != nil {
+						return nil, err
+					}
+					if password != "" {
+						if _, err := c.Do("AUTH", password); err != nil {
+							c.Close()
+							return nil, err
+						}
+					}
+					return c, err
+				},
+				TestOnBorrow: func(c redis.Conn, t time.Time) error {
+					if time.Since(t) < time.Minute {
+						return nil
+					}
+					_, err := c.Do("PING")
+					return err
+				},
+			}, nil
 		},
 	}
+
+	if err := cluster.Refresh(); err != nil && !isClusterDisabled(err) {
+		return nil, errors.Wrap(err, "refresh cluster")
+	}
+
+	return cluster, nil
+}
+
+func isClusterDisabled(err error) bool {
+	return strings.Contains(err.Error(), "ERR This instance has cluster support disabled")
 }
 
 // NewRedisQueryResults creats a new Redis implementation of the
 // QueryResultStore interface using the provided Redis connection pool.
-func NewRedisQueryResults(pool *redis.Pool, duplicateResults bool) *redisQueryResults {
+func NewRedisQueryResults(pool *redisc.Cluster, duplicateResults bool) *redisQueryResults {
 	return &redisQueryResults{pool: pool, duplicateResults: duplicateResults}
 }
 

server/service/endpoint_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/fleetdm/fleet/server/test"
 	kitlog "github.com/go-kit/kit/log"
 	"github.com/stretchr/testify/require"
-	"github.com/throttled/throttled/store/memstore"
+	"github.com/throttled/throttled/v2/store/memstore"
 )
 
 type testResource struct {

server/service/handler_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/go-kit/kit/log"
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
-	"github.com/throttled/throttled/store/memstore"
+	"github.com/throttled/throttled/v2/store/memstore"
 )
 
 func TestAPIRoutes(t *testing.T) {

server/service/http_auth_test.go

@@ -20,7 +20,7 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"github.com/throttled/throttled/store/memstore"
+	"github.com/throttled/throttled/v2/store/memstore"
 )
 
 func TestLogin(t *testing.T) {

server/service/middleware/ratelimit/ratelimit_test.go

@@ -5,8 +5,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/throttled/throttled/store/memstore"
 	"github.com/throttled/throttled/v2"
+	"github.com/throttled/throttled/v2/store/memstore"
 )
 
 // Intent is to test the middleware functionality. We rely on the tests within

server/sso/session_store.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gomodule/redigo/redis"
+	"github.com/mna/redisc"
 	"github.com/pkg/errors"
 )
 
@@ -33,12 +34,12 @@ type SessionStore interface {
 }
 
 // NewSessionStore creates a SessionStore
-func NewSessionStore(pool *redis.Pool) SessionStore {
+func NewSessionStore(pool *redisc.Cluster) SessionStore {
 	return &store{pool}
 }
 
 type store struct {
-	pool *redis.Pool
+	pool *redisc.Cluster
 }
 
 func (s *store) create(requestID, originalURL, metadata string, lifetimeSecs uint) error {

server/sso/session_store_test.go

@@ -6,12 +6,12 @@ import (
 	"time"
 
 	"github.com/fleetdm/fleet/server/pubsub"
-	"github.com/gomodule/redigo/redis"
+	"github.com/mna/redisc"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func newPool(t *testing.T) *redis.Pool {
+func newPool(t *testing.T) *redisc.Cluster {
 	if _, ok := os.LookupEnv("REDIS_TEST"); ok {
 		var (
 			addr     = "127.0.0.1:6379"
@@ -20,10 +20,11 @@ func newPool(t *testing.T) *redis.Pool {
 			useTLS   = false
 		)
 
-		p := pubsub.NewRedisPool(addr, password, database, useTLS)
-		_, err := p.Get().Do("PING")
+		pool, err := pubsub.NewRedisPool(addr, password, database, useTLS)
+		require.NoError(t, err)
+		_, err = pool.Get().Do("PING")
 		require.Nil(t, err)
-		return p
+		return pool
 	}
 	return nil
 }