Redis: simplify DeleteAllKeysWithPrefix (#64407)

The `DeleteAllKeysWithPrefix` method is now only used in tests to ensure
the test keyspace is clear. This PR makes it clear this method is only
used in tests, and simplifies the implementation so it no longer needs a
script + direct Redis connection.

Another tiny Redis refactor to prepare for multi-tenancy work.

Author: jtibshirani

cmd/worker/internal/ratelimit/handler_test.go

@@ -41,13 +41,11 @@ func TestHandler_Handle(t *testing.T) {
 			return err
 		},
 	}
+
 	t.Cleanup(func() {
-		c := pool.Get()
-		err := redispool.DeleteAllKeysWithPrefix(c, prefix)
-		if err != nil {
+		if err := redispool.DeleteAllKeysWithPrefix(redispool.RedisKeyValue(pool), prefix); err != nil {
 			t.Logf("Failed to clear redis: %+v\n", err)
 		}
-		c.Close()
 	})
 
 	conf.Mock(&conf.Unified{

internal/ratelimit/globallimiter.go

@@ -540,9 +540,8 @@ func SetupForTest(t TB) {
 		}
 	}
 
-	err := redispool.DeleteAllKeysWithPrefix(c, tokenBucketGlobalPrefix)
-	if err != nil {
-		t.Fatalf("cold not clear test prefix: &v", err)
+	if err := redispool.DeleteAllKeysWithPrefix(redispool.RedisKeyValue(pool), tokenBucketGlobalPrefix); err != nil {
+		t.Fatalf("could not clear test prefix: &v", err)
 	}
 }
 

internal/ratelimit/globallimiter_test.go

@@ -318,12 +318,7 @@ func redisPoolForTest(t *testing.T, prefix string) *redis.Pool {
 		},
 	}
 
-	c := pool.Get()
-	t.Cleanup(func() {
-		_ = c.Close()
-	})
-
-	if err := redispool.DeleteAllKeysWithPrefix(c, prefix); err != nil {
+	if err := redispool.DeleteAllKeysWithPrefix(redispool.RedisKeyValue(pool), prefix); err != nil {
 		t.Logf("Could not clear test prefix name=%q prefix=%q error=%v", t.Name(), prefix, err)
 	}
 

internal/rcache/BUILD.bazel

@@ -30,7 +30,6 @@ go_test(
         "requires-network",
     ],
     deps = [
-        "//internal/redispool",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],

internal/rcache/fifo_list.go

@@ -134,3 +134,11 @@ func (l *FIFOList) kv() redispool.KeyValue {
 	}
 	return l._kv
 }
+
+func bytes(s ...string) [][]byte {
+	t := make([][]byte, len(s))
+	for i, v := range s {
+		t[i] = []byte(v)
+	}
+	return t
+}

internal/rcache/rcache.go

@@ -272,8 +272,7 @@ func SetupForTest(t testing.TB) redispool.KeyValue {
 		}
 	}
 
-	err := redispool.DeleteAllKeysWithPrefix(c, globalPrefix)
-	if err != nil {
+	if err := redispool.DeleteAllKeysWithPrefix(kvMock, globalPrefix); err != nil {
 		log15.Error("Could not clear test prefix", "name", t.Name(), "globalPrefix", globalPrefix, "error", err)
 	}
 

internal/rcache/rcache_test.go

@@ -1,14 +1,10 @@
 package rcache
 
 import (
-	"reflect"
-	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
-
-	"github.com/sourcegraph/sourcegraph/internal/redispool"
 )
 
 func TestCache_namespace(t *testing.T) {
@@ -95,55 +91,6 @@ func TestCache_simple(t *testing.T) {
 	}
 }
 
-func TestCache_deleteAllKeysWithPrefix(t *testing.T) {
-	kv := SetupForTest(t)
-
-	c := New(kv, "some_prefix")
-	var aKeys, bKeys []string
-	var key string
-	for i := range 10 {
-		if i%2 == 0 {
-			key = "a:" + strconv.Itoa(i)
-			aKeys = append(aKeys, key)
-		} else {
-			key = "b:" + strconv.Itoa(i)
-			bKeys = append(bKeys, key)
-		}
-
-		c.Set(key, []byte(strconv.Itoa(i)))
-	}
-
-	pool := kv.Pool()
-
-	conn := pool.Get()
-	defer conn.Close()
-
-	err := redispool.DeleteAllKeysWithPrefix(conn, c.rkeyPrefix()+"a")
-	if err != nil {
-		t.Error(err)
-	}
-
-	getMulti := func(keys ...string) [][]byte {
-		t.Helper()
-		var vals [][]byte
-		for _, k := range keys {
-			v, _ := c.Get(k)
-			vals = append(vals, v)
-		}
-		return vals
-	}
-
-	vals := getMulti(aKeys...)
-	if got, exp := vals, [][]byte{nil, nil, nil, nil, nil}; !reflect.DeepEqual(exp, got) {
-		t.Errorf("Expected %v, but got %v", exp, got)
-	}
-
-	vals = getMulti(bKeys...)
-	if got, exp := vals, bytes("1", "3", "5", "7", "9"); !reflect.DeepEqual(exp, got) {
-		t.Errorf("Expected %v, but got %v", exp, got)
-	}
-}
-
 func TestCache_Increase(t *testing.T) {
 	kv := SetupForTest(t)
 
@@ -274,11 +221,3 @@ func TestCache_Hashes(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, 0, del4)
 }
-
-func bytes(s ...string) [][]byte {
-	t := make([][]byte, len(s))
-	for i, v := range s {
-		t[i] = []byte(v)
-	}
-	return t
-}

internal/redispool/BUILD.bazel

@@ -9,7 +9,7 @@ go_library(
         "mocks.go",
         "redispool.go",
         "sysreq.go",
-        "utils.go",
+        "test_utils.go",
     ],
     importpath = "github.com/sourcegraph/sourcegraph/internal/redispool",
     visibility = ["//:__subpackages__"],

internal/redispool/keyvalue_test.go

@@ -389,14 +389,15 @@ func redisKeyValueForTest(t *testing.T) redispool.KeyValue {
 		}
 	}
 
-	if err := redispool.DeleteAllKeysWithPrefix(c, prefix); err != nil {
+	kv := redispool.RedisKeyValue(pool)
+	if err := redispool.DeleteAllKeysWithPrefix(kv, prefix); err != nil {
 		t.Logf("Could not clear test prefix name=%q prefix=%q error=%v", t.Name(), prefix, err)
 	}
 
-	kv := redispool.RedisKeyValue(pool).(interface {
+	redisKv := kv.(interface {
 		WithPrefix(string) redispool.KeyValue
 	})
-	return kv.WithPrefix(prefix)
+	return redisKv.WithPrefix(prefix)
 }
 
 func bytes(ss ...string) [][]byte {

internal/redispool/redispool_test.go

@@ -3,8 +3,12 @@ package redispool
 import (
 	"flag"
 	"os"
+	"reflect"
+	"strconv"
 	"testing"
+	"time"
 
+	"github.com/gomodule/redigo/redis"
 	"github.com/sourcegraph/log/logtest"
 )
 
@@ -32,3 +36,71 @@ func TestMain(m *testing.M) {
 	logtest.Init(m)
 	os.Exit(m.Run())
 }
+
+func TestDeleteAllKeysWithPrefix(t *testing.T) {
+	t.Helper()
+
+	pool := &redis.Pool{
+		MaxIdle:     3,
+		IdleTimeout: 240 * time.Second,
+		Dial: func() (redis.Conn, error) {
+			return redis.Dial("tcp", "127.0.0.1:6379")
+		},
+		TestOnBorrow: func(c redis.Conn, t time.Time) error {
+			_, err := c.Do("PING")
+			return err
+		},
+	}
+
+	c := 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 {
+			t.Skip("could not connect to redis", err)
+		}
+	}
+
+	kv := RedisKeyValue(pool)
+	var aKeys, bKeys []string
+	var key string
+	for i := range 10 {
+		if i%2 == 0 {
+			key = "a:" + strconv.Itoa(i)
+			aKeys = append(aKeys, key)
+		} else {
+			key = "b:" + strconv.Itoa(i)
+			bKeys = append(bKeys, key)
+		}
+
+		if err := kv.Set(key, []byte(strconv.Itoa(i))); err != nil {
+			t.Fatalf("could not set key %s: %v", key, err)
+		}
+	}
+
+	if err := DeleteAllKeysWithPrefix(kv, "a"); err != nil {
+		t.Fatal(err)
+	}
+
+	getMulti := func(keys ...string) []string {
+		t.Helper()
+		var vals []string
+		for _, k := range keys {
+			v, _ := kv.Get(k).String()
+			vals = append(vals, v)
+		}
+		return vals
+	}
+
+	vals := getMulti(aKeys...)
+	if got, exp := vals, []string{"", "", "", "", ""}; !reflect.DeepEqual(exp, got) {
+		t.Errorf("Expected %v, but got %v", exp, got)
+	}
+
+	vals = getMulti(bKeys...)
+	if got, exp := vals, []string{"1", "3", "5", "7", "9"}; !reflect.DeepEqual(exp, got) {
+		t.Errorf("Expected %v, but got %v", exp, got)
+	}
+}