Redis: remove some direct pool usages (#64447)

We want to discourage direct usage of the Redis pool in favor of routing
all calls through the main `KeyValue` interface. This PR removes several
usages of `KeyValue.Pool`. To do so, it adds "PING" and "MGET" to the
`KeyValue` interface.

Author: jtibshirani

cmd/cody-gateway/internal/httpapi/BUILD.bazel

@@ -36,7 +36,6 @@ go_library(
         "//internal/version",
         "//lib/errors",
         "//lib/pointers",
-        "@com_github_gomodule_redigo//redis",
         "@com_github_gorilla_mux//:mux",
         "@com_github_khan_genqlient//graphql",
         "@com_github_sourcegraph_log//:log",

cmd/cody-gateway/internal/httpapi/diagnostics.go

@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/gomodule/redigo/redis"
 	"github.com/sourcegraph/log"
 	"github.com/sourcegraph/log/hook"
 	"github.com/sourcegraph/log/output"
@@ -17,16 +16,16 @@ import (
 	"github.com/sourcegraph/sourcegraph/cmd/cody-gateway/internal/response"
 	"github.com/sourcegraph/sourcegraph/internal/authbearer"
 	"github.com/sourcegraph/sourcegraph/internal/instrumentation"
+	"github.com/sourcegraph/sourcegraph/internal/redispool"
 	sgtrace "github.com/sourcegraph/sourcegraph/internal/trace"
 	"github.com/sourcegraph/sourcegraph/internal/version"
-	"github.com/sourcegraph/sourcegraph/lib/errors"
 )
 
 // NewDiagnosticsHandler creates a handler for diagnostic endpoints typically served
 // on "/-/..." paths. It should be placed before any authentication middleware, since
 // we do a simple auth on a static secret instead that is uniquely generated per
 // deployment.
-func NewDiagnosticsHandler(baseLogger log.Logger, next http.Handler, redisPool *redis.Pool, secret string, sources *actor.Sources) http.Handler {
+func NewDiagnosticsHandler(baseLogger log.Logger, next http.Handler, redisCache redispool.KeyValue, secret string, sources *actor.Sources) http.Handler {
 	baseLogger = baseLogger.Scoped("diagnostics")
 
 	hasValidSecret := func(l log.Logger, w http.ResponseWriter, r *http.Request) (yes bool) {
@@ -58,7 +57,7 @@ func NewDiagnosticsHandler(baseLogger log.Logger, next http.Handler, redisPool *
 				return
 			}
 
-			if err := healthz(r.Context(), redisPool); err != nil {
+			if err := healthz(r.Context(), redisCache); err != nil {
 				logger.Error("check failed", log.Error(err))
 
 				w.WriteHeader(http.StatusInternalServerError)
@@ -110,21 +109,6 @@ func NewDiagnosticsHandler(baseLogger log.Logger, next http.Handler, redisPool *
 	})
 }
 
-func healthz(ctx context.Context, rpool *redis.Pool) error {
-	// Check redis health
-	rconn, err := rpool.GetContext(ctx)
-	if err != nil {
-		return errors.Wrap(err, "redis: failed to get conn")
-	}
-	defer rconn.Close()
-
-	data, err := rconn.Do("PING")
-	if err != nil {
-		return errors.Wrap(err, "redis: failed to ping")
-	}
-	if data != "PONG" {
-		return errors.New("redis: failed to ping: no pong received")
-	}
-
-	return nil
+func healthz(ctx context.Context, cache redispool.KeyValue) error {
+	return cache.WithContext(ctx).Ping()
 }

cmd/cody-gateway/shared/main.go

@@ -241,7 +241,7 @@ func Main(ctx context.Context, obctx *observation.Context, ready service.ReadyFu
 		return errors.Wrap(err, "httpapi.NewHandler")
 	}
 	// Diagnostic and Maintenance layers, exposing additional APIs and endpoints.
-	handler = httpapi.NewDiagnosticsHandler(obctx.Logger, handler, redisCache.Pool(), cfg.DiagnosticsSecret, sources)
+	handler = httpapi.NewDiagnosticsHandler(obctx.Logger, handler, redisCache, cfg.DiagnosticsSecret, sources)
 	handler = httpapi.NewMaintenanceHandler(obctx.Logger, handler, cfg, redisCache)
 
 	// Collect request client for downstream handlers. Outside of dev, we always set up

cmd/frontend/internal/cli/serve_cmd.go

@@ -450,19 +450,6 @@ func GetInternalAddr() string {
 	return httpAddrInternal
 }
 
-func pingRedis(kv redispool.KeyValue) error {
-	conn := kv.Pool().Get()
-	defer conn.Close()
-	data, err := conn.Do("PING")
-	if err != nil {
-		return err
-	}
-	if data != "PONG" {
-		return errors.New("no pong received")
-	}
-	return nil
-}
-
 // waitForRedis waits up to a certain timeout for Redis to become reachable, to reduce the
 // likelihood of the HTTP handlers starting to serve requests while Redis (and therefore session
 // data) is still unavailable. After the timeout has elapsed, if Redis is still unreachable, it
@@ -473,7 +460,7 @@ func waitForRedis(logger sglog.Logger, kv redispool.KeyValue) {
 	var err error
 	for {
 		time.Sleep(150 * time.Millisecond)
-		err = pingRedis(kv)
+		err = kv.Ping()
 		if err == nil {
 			return
 		}

internal/metrics/store/BUILD.bazel

@@ -12,7 +12,6 @@ go_library(
     deps = [
         "//internal/redispool",
         "//lib/errors",
-        "@com_github_gomodule_redigo//redis",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_prometheus_client_model//go",
         "@com_github_prometheus_common//expfmt",

internal/metrics/store/store.go

@@ -5,7 +5,6 @@ import (
 	"io"
 	"strings"
 
-	"github.com/gomodule/redigo/redis"
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/prometheus/common/expfmt"
@@ -20,16 +19,6 @@ type Store interface {
 	prometheus.Gatherer
 }
 
-func NewDefaultStore() Store {
-	return &defaultStore{}
-}
-
-type defaultStore struct{}
-
-func (*defaultStore) Gather() ([]*dto.MetricFamily, error) {
-	return prometheus.DefaultGatherer.Gather()
-}
-
 type DistributedStore interface {
 	Store
 	Ingest(instance string, mfs []*dto.MetricFamily) error
@@ -48,13 +37,10 @@ type distributedStore struct {
 }
 
 func (d *distributedStore) Gather() ([]*dto.MetricFamily, error) {
-	pool := redispool.Cache.Pool()
-
-	reConn := pool.Get()
-	defer reConn.Close()
+	cache := redispool.Cache
 
 	// First, list all the keys for which we hold metrics.
-	keys, err := redis.Values(reConn.Do("KEYS", d.prefix+"*"))
+	keys, err := cache.Keys(d.prefix + "*")
 	if err != nil {
 		return nil, errors.Wrap(err, "listing entries from redis")
 	}
@@ -64,7 +50,7 @@ func (d *distributedStore) Gather() ([]*dto.MetricFamily, error) {
 	}
 
 	// Then bulk retrieve all the metrics blobs for all the instances.
-	encodedMetrics, err := redis.Strings(reConn.Do("MGET", keys...))
+	encodedMetrics, err := cache.MGet(keys).Strings()
 	if err != nil {
 		return nil, errors.Wrap(err, "retrieving blobs from redis")
 	}
@@ -92,7 +78,7 @@ func (d *distributedStore) Gather() ([]*dto.MetricFamily, error) {
 }
 
 func (d *distributedStore) Ingest(instance string, mfs []*dto.MetricFamily) error {
-	pool := redispool.Cache.Pool()
+	cache := redispool.Cache
 
 	// First, encode the metrics to text format so we can store them.
 	var enc bytes.Buffer
@@ -106,13 +92,10 @@ func (d *distributedStore) Ingest(instance string, mfs []*dto.MetricFamily) erro
 
 	encodedMetrics := enc.String()
 
-	reConn := pool.Get()
-	defer reConn.Close()
-
 	// Store the metrics and set an expiry on the key, if we haven't retrieved
 	// an updated set of metric data, we consider the host down and prune it
 	// from the gatherer.
-	err := reConn.Send("SETEX", d.prefix+instance, d.expiry, encodedMetrics)
+	err := cache.SetEx(d.prefix+instance, d.expiry, encodedMetrics)
 	if err != nil {
 		return errors.Wrap(err, "writing metrics blob to redis")
 	}

internal/ratelimit/globallimiter.go

@@ -499,15 +499,11 @@ func SetupForTest(t TB) {
 	t.Helper()
 
 	kvMock = redispool.NewTestKeyValue()
-
 	tokenBucketGlobalPrefix = "__test__" + t.Name()
-	c := kvMock.Pool().Get()
-	defer c.Close()
 
 	// If we are not on CI, skip the test if our redis connection fails.
 	if os.Getenv("CI") == "" {
-		_, err := c.Do("PING")
-		if err != nil {
+		if err := kvMock.Ping(); err != nil {
 			t.Skip("could not connect to redis", err)
 		}
 	}

internal/rcache/rcache.go

@@ -250,13 +250,10 @@ func SetupForTest(t testing.TB) redispool.KeyValue {
 	})
 
 	globalPrefix = "__test__" + t.Name()
-	c := kvMock.Pool().Get()
-	defer c.Close()
 
 	// If we are not on CI, skip the test if our redis connection fails.
 	if os.Getenv("CI") == "" {
-		_, err := c.Do("PING")
-		if err != nil {
+		if err := kvMock.Ping(); err != nil {
 			t.Skip("could not connect to redis", err)
 		}
 	}