diff --git a/cmd/worker/internal/ratelimit/handler_test.go b/cmd/worker/internal/ratelimit/handler_test.go index 010ff9bb7ac..6fc04f2fbc9 100644 --- a/cmd/worker/internal/ratelimit/handler_test.go +++ b/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{ diff --git a/internal/ratelimit/globallimiter.go b/internal/ratelimit/globallimiter.go index 41b0719332b..c7f681835a7 100644 --- a/internal/ratelimit/globallimiter.go +++ b/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) } } diff --git a/internal/ratelimit/globallimiter_test.go b/internal/ratelimit/globallimiter_test.go index ce1bf3dcbb3..2814be3cea3 100644 --- a/internal/ratelimit/globallimiter_test.go +++ b/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) } diff --git a/internal/rcache/BUILD.bazel b/internal/rcache/BUILD.bazel index 22e5f4630cc..956823db16e 100644 --- a/internal/rcache/BUILD.bazel +++ b/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", ], diff --git a/internal/rcache/fifo_list.go b/internal/rcache/fifo_list.go index 363ef0a655a..7a172a19ff1 100644 --- a/internal/rcache/fifo_list.go +++ b/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 +} diff --git a/internal/rcache/rcache.go b/internal/rcache/rcache.go index a8f80d12b64..10f0eb13974 100644 --- a/internal/rcache/rcache.go +++ b/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) } diff --git a/internal/rcache/rcache_test.go b/internal/rcache/rcache_test.go index dc3e973c290..11ff954c0f3 100644 --- a/internal/rcache/rcache_test.go +++ b/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 -} diff --git a/internal/redispool/BUILD.bazel b/internal/redispool/BUILD.bazel index 9c1f41483ff..2eaa47b9d34 100644 --- a/internal/redispool/BUILD.bazel +++ b/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__"], diff --git a/internal/redispool/keyvalue_test.go b/internal/redispool/keyvalue_test.go index 640e5d5638c..81e6ee31c82 100644 --- a/internal/redispool/keyvalue_test.go +++ b/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 { diff --git a/internal/redispool/redispool_test.go b/internal/redispool/redispool_test.go index 5400c1f1ce4..1cfdb4d3f27 100644 --- a/internal/redispool/redispool_test.go +++ b/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) + } +} diff --git a/internal/redispool/test_utils.go b/internal/redispool/test_utils.go new file mode 100644 index 00000000000..bb0a04e2f89 --- /dev/null +++ b/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 +} diff --git a/internal/redispool/utils.go b/internal/redispool/utils.go deleted file mode 100644 index e25256ea383..00000000000 --- a/internal/redispool/utils.go +++ /dev/null @@ -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 -}