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.

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)
+	}
+}

internal/redispool/test_utils.go

@@ -0,0 +1,17 @@
+package redispool
+
+// DeleteAllKeysWithPrefix retrieves all keys starting with 'prefix:' and sequentially deletes each one.
+//
+// NOTE: this should only be used for tests, as it is not atomic and cannot handle large keyspaces.
+func DeleteAllKeysWithPrefix(kv KeyValue, prefix string) error {
+	keys, err := kv.Keys(prefix + ":*")
+	if err != nil {
+		return err
+	}
+	for _, key := range keys {
+		if err := kv.Del(key); err != nil {
+			return err
+		}
+	}
+	return nil
+}

internal/redispool/utils.go

@@ -1,33 +0,0 @@
-package redispool
-
-import "github.com/gomodule/redigo/redis"
-
-// The number of keys to delete per batch.
-// The maximum number of keys that can be unpacked
-// is determined by the Lua config LUAI_MAXCSTACK
-// which is 8000 by default.
-// See https://www.lua.org/source/5.1/luaconf.h.html
-var deleteBatchSize = 5000
-
-func DeleteAllKeysWithPrefix(c redis.Conn, prefix string) error {
-	const script = `
-redis.replicate_commands()
-local cursor = '0'
-local prefix = ARGV[1]
-local batchSize = ARGV[2]
-local result = ''
-repeat
-	local keys = redis.call('SCAN', cursor, 'MATCH', prefix, 'COUNT', batchSize)
-	if #keys[2] > 0
-	then
-		result = redis.call('DEL', unpack(keys[2]))
-	end
-
-	cursor = keys[1]
-until cursor == '0'
-return result
-`
-
-	_, err := c.Do("EVAL", script, 0, prefix+":*", deleteBatchSize)
-	return err
-}