Redis: expose KeyValue.WithPrefix (#64466)

This PR makes `WithPrefix` visible on the `KeyValue` interface.
Previously, we had other interface implementations that did not support
`WithPrefix`, but now `redisKeyValue` is the only implementation.
`WithPrefix` is currently just used in tests, but it's broadly useful
and will help clean up a bunch of places that wrap Redis and manually
add key prefixes.

internal/ratelimit/globallimiter.go

@@ -498,26 +498,26 @@ type TB interface {
 func SetupForTest(t TB) {
 	t.Helper()
 
-	kvMock = redispool.NewTestKeyValue()
-	tokenBucketGlobalPrefix = "__test__" + t.Name()
+	testStore = redispool.NewTestKeyValue()
 
 	// If we are not on CI, skip the test if our redis connection fails.
 	if os.Getenv("CI") == "" {
-		if err := kvMock.Ping(); err != nil {
+		if err := testStore.Ping(); err != nil {
 			t.Skip("could not connect to redis", err)
 		}
 	}
 
-	if err := redispool.DeleteAllKeysWithPrefix(kvMock, tokenBucketGlobalPrefix); err != nil {
+	tokenBucketGlobalPrefix = "__test__" + t.Name()
+	if err := redispool.DeleteAllKeysWithPrefix(testStore, tokenBucketGlobalPrefix); err != nil {
 		t.Fatalf("could not clear test prefix: &v", err)
 	}
 }
 
-var kvMock redispool.KeyValue
+var testStore redispool.KeyValue
 
 func kv() redispool.KeyValue {
-	if kvMock != nil {
-		return kvMock
+	if testStore != nil {
+		return testStore
 	}
 	return redispool.Store
 }

internal/rcache/fifo_list.go

@@ -129,8 +129,8 @@ func (l *FIFOList) globalPrefixKey() string {
 }
 
 func (l *FIFOList) kv() redispool.KeyValue {
-	if kvMock != nil {
-		return kvMock
+	if testStore != nil {
+		return testStore
 	}
 	return l._kv
 }

internal/rcache/rcache.go

@@ -243,37 +243,36 @@ const testAddr = "127.0.0.1:6379"
 func SetupForTest(t testing.TB) redispool.KeyValue {
 	t.Helper()
 
-	kvMock = redispool.NewTestKeyValue()
+	testStore = redispool.NewTestKeyValue()
 	t.Cleanup(func() {
-		kvMock.Pool().Close()
-		kvMock = nil
+		testStore.Pool().Close()
+		testStore = nil
 	})
 
-	globalPrefix = "__test__" + t.Name()
-
 	// If we are not on CI, skip the test if our redis connection fails.
 	if os.Getenv("CI") == "" {
-		if err := kvMock.Ping(); err != nil {
+		if err := testStore.Ping(); err != nil {
 			t.Skip("could not connect to redis", err)
 		}
 	}
 
-	if err := redispool.DeleteAllKeysWithPrefix(kvMock, globalPrefix); err != nil {
+	globalPrefix = "__test__" + t.Name()
+	if err := redispool.DeleteAllKeysWithPrefix(testStore, globalPrefix); err != nil {
 		log15.Error("Could not clear test prefix", "name", t.Name(), "globalPrefix", globalPrefix, "error", err)
 	}
 
-	return kvMock
+	return testStore
 }
 
-var kvMock redispool.KeyValue
+var testStore redispool.KeyValue
 
 func (r *Cache) kv() redispool.KeyValue {
 	// TODO: We should refactor the SetupForTest method to return a KV, not mock
 	// a global thing.
 	// That can only work when all tests pass the redis connection directly to the
 	// tested methods though.
-	if kvMock != nil {
-		return kvMock
+	if testStore != nil {
+		return testStore
 	}
 	return r._kv
 }

internal/redispool/keyvalue.go

@@ -56,6 +56,9 @@ type KeyValue interface {
 	WithContext(ctx context.Context) KeyValue
 	WithLatencyRecorder(r LatencyRecorder) KeyValue
 
+	// WithPrefix wraps r to return a RedisKeyValue that prefixes all keys with 'prefix:'
+	WithPrefix(prefix string) KeyValue
+
 	// Pool returns the underlying redis pool.
 	// The intention of this API is Pool is only for advanced use cases and the caller
 	// should consider if they need to use it. Pool is very hard to mock, while
@@ -159,14 +162,6 @@ func NewTestKeyValue() KeyValue {
 }
 
 // redisKeyValue is a KeyValue backed by pool
-//
-// Note: redisKeyValue additionally implements
-//
-//	interface {
-//	  // WithPrefix wraps r to return a RedisKeyValue that prefixes all keys with
-//	  // prefix + ":".
-//	  WithPrefix(prefix string) KeyValue
-//	}
 type redisKeyValue struct {
 	pool     *redis.Pool
 	ctx      context.Context
@@ -282,7 +277,6 @@ func (r *redisKeyValue) WithLatencyRecorder(rec LatencyRecorder) KeyValue {
 	}
 }
 
-// WithPrefix wraps r to return a RedisKeyValue that prefixes all keys with 'prefix:'
 func (r *redisKeyValue) WithPrefix(prefix string) KeyValue {
 	return &redisKeyValue{
 		pool:   r.pool,

internal/redispool/keyvalue_test.go

@@ -430,10 +430,7 @@ func redisKeyValueForTest(t *testing.T) redispool.KeyValue {
 		t.Logf("Could not clear test prefix name=%q prefix=%q error=%v", t.Name(), prefix, err)
 	}
 
-	redisKv := kv.(interface {
-		WithPrefix(string) redispool.KeyValue
-	})
-	return redisKv.WithPrefix(prefix)
+	return kv.WithPrefix(prefix)
 }
 
 func bytes(ss ...string) [][]byte {

internal/redispool/mocks.go

@@ -95,6 +95,9 @@ type MockKeyValue struct {
 	// WithLatencyRecorderFunc is an instance of a mock function object
 	// controlling the behavior of the method WithLatencyRecorder.
 	WithLatencyRecorderFunc *KeyValueWithLatencyRecorderFunc
+	// WithPrefixFunc is an instance of a mock function object controlling
+	// the behavior of the method WithPrefix.
+	WithPrefixFunc *KeyValueWithPrefixFunc
 }
 
 // NewMockKeyValue creates a new mock of the KeyValue interface. All methods
@@ -231,6 +234,11 @@ func NewMockKeyValue() *MockKeyValue {
 				return
 			},
 		},
+		WithPrefixFunc: &KeyValueWithPrefixFunc{
+			defaultHook: func(string) (r0 KeyValue) {
+				return
+			},
+		},
 	}
 }
 
@@ -368,6 +376,11 @@ func NewStrictMockKeyValue() *MockKeyValue {
 				panic("unexpected invocation of MockKeyValue.WithLatencyRecorder")
 			},
 		},
+		WithPrefixFunc: &KeyValueWithPrefixFunc{
+			defaultHook: func(string) KeyValue {
+				panic("unexpected invocation of MockKeyValue.WithPrefix")
+			},
+		},
 	}
 }
 
@@ -453,6 +466,9 @@ func NewMockKeyValueFrom(i KeyValue) *MockKeyValue {
 		WithLatencyRecorderFunc: &KeyValueWithLatencyRecorderFunc{
 			defaultHook: i.WithLatencyRecorder,
 		},
+		WithPrefixFunc: &KeyValueWithPrefixFunc{
+			defaultHook: i.WithPrefix,
+		},
 	}
 }
 
@@ -3158,3 +3174,104 @@ func (c KeyValueWithLatencyRecorderFuncCall) Args() []interface{} {
 func (c KeyValueWithLatencyRecorderFuncCall) Results() []interface{} {
 	return []interface{}{c.Result0}
 }
+
+// KeyValueWithPrefixFunc describes the behavior when the WithPrefix method
+// of the parent MockKeyValue instance is invoked.
+type KeyValueWithPrefixFunc struct {
+	defaultHook func(string) KeyValue
+	hooks       []func(string) KeyValue
+	history     []KeyValueWithPrefixFuncCall
+	mutex       sync.Mutex
+}
+
+// WithPrefix delegates to the next hook function in the queue and stores
+// the parameter and result values of this invocation.
+func (m *MockKeyValue) WithPrefix(v0 string) KeyValue {
+	r0 := m.WithPrefixFunc.nextHook()(v0)
+	m.WithPrefixFunc.appendCall(KeyValueWithPrefixFuncCall{v0, r0})
+	return r0
+}
+
+// SetDefaultHook sets function that is called when the WithPrefix method of
+// the parent MockKeyValue instance is invoked and the hook queue is empty.
+func (f *KeyValueWithPrefixFunc) SetDefaultHook(hook func(string) KeyValue) {
+	f.defaultHook = hook
+}
+
+// PushHook adds a function to the end of hook queue. Each invocation of the
+// WithPrefix method of the parent MockKeyValue instance invokes the hook at
+// the front of the queue and discards it. After the queue is empty, the
+// default hook function is invoked for any future action.
+func (f *KeyValueWithPrefixFunc) PushHook(hook func(string) KeyValue) {
+	f.mutex.Lock()
+	f.hooks = append(f.hooks, hook)
+	f.mutex.Unlock()
+}
+
+// SetDefaultReturn calls SetDefaultHook with a function that returns the
+// given values.
+func (f *KeyValueWithPrefixFunc) SetDefaultReturn(r0 KeyValue) {
+	f.SetDefaultHook(func(string) KeyValue {
+		return r0
+	})
+}
+
+// PushReturn calls PushHook with a function that returns the given values.
+func (f *KeyValueWithPrefixFunc) PushReturn(r0 KeyValue) {
+	f.PushHook(func(string) KeyValue {
+		return r0
+	})
+}
+
+func (f *KeyValueWithPrefixFunc) nextHook() func(string) KeyValue {
+	f.mutex.Lock()
+	defer f.mutex.Unlock()
+
+	if len(f.hooks) == 0 {
+		return f.defaultHook
+	}
+
+	hook := f.hooks[0]
+	f.hooks = f.hooks[1:]
+	return hook
+}
+
+func (f *KeyValueWithPrefixFunc) appendCall(r0 KeyValueWithPrefixFuncCall) {
+	f.mutex.Lock()
+	f.history = append(f.history, r0)
+	f.mutex.Unlock()
+}
+
+// History returns a sequence of KeyValueWithPrefixFuncCall objects
+// describing the invocations of this function.
+func (f *KeyValueWithPrefixFunc) History() []KeyValueWithPrefixFuncCall {
+	f.mutex.Lock()
+	history := make([]KeyValueWithPrefixFuncCall, len(f.history))
+	copy(history, f.history)
+	f.mutex.Unlock()
+
+	return history
+}
+
+// KeyValueWithPrefixFuncCall is an object that describes an invocation of
+// method WithPrefix on an instance of MockKeyValue.
+type KeyValueWithPrefixFuncCall struct {
+	// Arg0 is the value of the 1st argument passed to this method
+	// invocation.
+	Arg0 string
+	// Result0 is the value of the 1st result returned from this method
+	// invocation.
+	Result0 KeyValue
+}
+
+// Args returns an interface slice containing the arguments of this
+// invocation.
+func (c KeyValueWithPrefixFuncCall) Args() []interface{} {
+	return []interface{}{c.Arg0}
+}
+
+// Results returns an interface slice containing the results of this
+// invocation.
+func (c KeyValueWithPrefixFuncCall) Results() []interface{} {
+	return []interface{}{c.Result0}
+}