fix(kvblock): avoid Redis full-scan clear

Signed-off-by: Kay Yan <kay.yan@daocloud.io>

Author: yankay

pkg/kvcache/kvblock/redis.go

@@ -272,6 +272,7 @@ func (r *RedisIndex) Add(ctx context.Context, engineKeys, requestKeys []BlockHas
 				return err
 			}
 			pipe.HSet(ctx, redisKey, field, "")
+			pipe.SAdd(ctx, redisPodEntriesKey(entry.PodIdentifier), redisPodEntryMember(redisKey, field))
 		}
 	}
 
@@ -330,6 +331,7 @@ func (r *RedisIndex) evictPodsFromRequestKey(ctx context.Context, requestKey Blo
 			return err
 		}
 		pipe.HDel(ctx, redisKey, field)
+		pipe.SRem(ctx, redisPodEntriesKey(entry.PodIdentifier), redisPodEntryMember(redisKey, field))
 	}
 
 	if _, err := pipe.Exec(ctx); err != nil {
@@ -408,60 +410,79 @@ func redisEngineKey(engineKey BlockHash) string {
 	return "engine:" + engineKey.String()
 }
 
-// Clear removes every hash field for the pod across all request-key hashes and
-// device tiers. Each field is a JSON-encoded PodEntry, so matching decodes the
-// field and compares PodIdentifier — catching every tier, group, and speculative
-// variant. It pages the keyspace with SCAN (skipping engine: keys), HDELs the
-// pod's fields, and prunes now-empty hashes. O(keyspace), but Clear is rare and
-// off the Lookup/Add hot path. Because it deletes from the shared store, it is
-// correct for multi-replica deployments with no cross-process coordination.
+func redisPodEntriesKey(podIdentifier string) string {
+	return "kvblock:pod:" + podIdentifier + ":entries"
+}
+
+func redisPodEntryMember(requestKey, field string) string {
+	return requestKey + "\x00" + field
+}
+
+type redisPodEntryRef struct {
+	requestKey string
+	field      string
+}
+
+func parseRedisPodEntryMember(member string) (redisPodEntryRef, bool) {
+	requestKey, field, ok := strings.Cut(member, "\x00")
+	if !ok || requestKey == "" || field == "" {
+		return redisPodEntryRef{}, false
+	}
+	return redisPodEntryRef{requestKey: requestKey, field: field}, true
+}
+
+// Clear removes every reverse-indexed field for the pod across all request-key
+// hashes and device tiers. Add maintains an exact per-pod set of
+// "<requestKey>\x00<encodedPodField>" entries, so Clear only touches keys this
+// index recorded for that pod instead of scanning the whole Redis keyspace.
 func (r *RedisIndex) Clear(ctx context.Context, podIdentifier string) error {
 	logger := log.FromContext(ctx).WithName("kvblock.RedisIndex.Clear")
+	podEntriesKey := redisPodEntriesKey(podIdentifier)
 
 	const scanBatch int64 = 1024
-	removed := 0
+	processed := 0
 	var cursor uint64
 	for {
-		keys, next, err := r.RedisClient.Scan(ctx, cursor, "*", scanBatch).Result()
+		members, next, err := r.RedisClient.SScan(ctx, podEntriesKey, cursor, "", scanBatch).Result()
 		if err != nil {
-			return fmt.Errorf("clear scan failed: %w", err)
+			return fmt.Errorf("clear reverse-index scan failed: %w", err)
 		}
-		for _, key := range keys {
-			if strings.HasPrefix(key, "engine:") {
-				continue // engine:<hash> ZSETs hold no pod fields
-			}
-
-			fields, err := r.RedisClient.HKeys(ctx, key).Result()
-			if err != nil {
-				return fmt.Errorf("clear hkeys failed for %s: %w", key, err)
-			}
-
-			var stale []string
-			for _, field := range fields {
-				if entry, ok := decodeRedisPodField(field); ok && entry.PodIdentifier == podIdentifier {
-					stale = append(stale, field)
-				}
-			}
-			if len(stale) == 0 {
-				continue
-			}
-
-			if err := r.RedisClient.HDel(ctx, key, stale...).Err(); err != nil {
-				return fmt.Errorf("clear hdel failed for %s: %w", key, err)
-			}
-			removed += len(stale)
-
-			if err := pruneRequestKeyScript.Run(ctx, r.RedisClient, []string{key}).Err(); err != nil &&
-				!errors.Is(err, redis.Nil) {
-				return fmt.Errorf("clear prune failed for %s: %w", key, err)
-			}
+		if err := r.clearRedisPodEntryMembers(ctx, podEntriesKey, members); err != nil {
+			return err
 		}
+		processed += len(members)
 
-		if cursor = next; cursor == 0 {
+		if cursor = next; cursor != 0 {
+			continue
+		}
+		remaining, err := r.RedisClient.SCard(ctx, podEntriesKey).Result()
+		if err != nil {
+			return fmt.Errorf("clear reverse-index cardinality failed: %w", err)
+		}
+		if remaining == 0 {
 			break
 		}
 	}
 
-	logger.Info("cleared pod from index", "pod", podIdentifier, "removed", removed)
+	logger.Info("cleared pod from index", "pod", podIdentifier, "members", processed)
+	return nil
+}
+
+func (r *RedisIndex) clearRedisPodEntryMembers(ctx context.Context, podEntriesKey string, members []string) error {
+	if len(members) == 0 {
+		return nil
+	}
+
+	pipe := r.RedisClient.Pipeline()
+	for _, member := range members {
+		ref, ok := parseRedisPodEntryMember(member)
+		if ok {
+			pipe.HDel(ctx, ref.requestKey, ref.field)
+		}
+		pipe.SRem(ctx, podEntriesKey, member)
+	}
+	if _, err := pipe.Exec(ctx); err != nil {
+		return fmt.Errorf("clear reverse-index pipeline failed: %w", err)
+	}
 	return nil
 }

pkg/kvcache/kvblock/redis_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package kvblock_test
 
 import (
+	"encoding/json"
 	"testing"
 
 	"github.com/alicebob/miniredis/v2"
@@ -27,6 +28,11 @@ import (
 
 // createRedisIndexForTesting creates a new RedisIndex with a mock Redis server for testing.
 func createRedisIndexForTesting(t *testing.T) Index {
+	t.Helper()
+	return createRedisIndexForInspection(t)
+}
+
+func createRedisIndexForInspection(t *testing.T) *RedisIndex {
 	t.Helper()
 	server, err := miniredis.Run()
 	require.NoError(t, err)
@@ -41,10 +47,81 @@ func createRedisIndexForTesting(t *testing.T) Index {
 	}
 	index, err := NewRedisIndex(redisConfig)
 	require.NoError(t, err)
-	return index
+	redisIndex, ok := index.(*RedisIndex)
+	require.True(t, ok)
+	return redisIndex
 }
 
 // TestRedisIndexBehavior tests the Redis index implementation using common test behaviors.
 func TestRedisIndexBehavior(t *testing.T) {
 	testCommonIndexBehavior(t, createRedisIndexForTesting)
 }
+
+func TestRedisClearPreservesUnrelatedKeys(t *testing.T) {
+	index := createRedisIndexForInspection(t)
+	ctx := t.Context()
+	pod := PodEntry{PodIdentifier: "pod-clear", DeviceTier: "gpu"}
+	requestKey := BlockHash(0xC1EA1001)
+	unrelatedField := redisTestPodField(t, pod)
+
+	require.NoError(t, index.Add(ctx, nil, []BlockHash{requestKey}, []PodEntry{pod}))
+	require.NoError(t, index.RedisClient.HSet(ctx, "unrelated-hash", unrelatedField, "keep").Err())
+	require.NoError(t, index.RedisClient.Set(ctx, "unrelated-string", "keep", 0).Err())
+
+	require.NoError(t, index.Clear(ctx, pod.PodIdentifier))
+
+	gotHash, err := index.RedisClient.HGet(ctx, "unrelated-hash", unrelatedField).Result()
+	require.NoError(t, err)
+	require.Equal(t, "keep", gotHash)
+
+	gotString, err := index.RedisClient.Get(ctx, "unrelated-string").Result()
+	require.NoError(t, err)
+	require.Equal(t, "keep", gotString)
+
+	hits, err := index.Lookup(ctx, []BlockHash{requestKey}, nil)
+	require.NoError(t, err)
+	require.Empty(t, hits[requestKey])
+
+	exists, err := index.RedisClient.Exists(ctx, requestKey.String()).Result()
+	require.NoError(t, err)
+	require.Zero(t, exists)
+}
+
+func TestRedisEvictPrunesReverseIndex(t *testing.T) {
+	index := createRedisIndexForInspection(t)
+	ctx := t.Context()
+	pod := PodEntry{PodIdentifier: "pod-evict", DeviceTier: "gpu"}
+	requestKey := BlockHash(0xC1EA1002)
+	podEntriesKey := "kvblock:pod:" + pod.PodIdentifier + ":entries"
+
+	require.NoError(t, index.Add(ctx, nil, []BlockHash{requestKey}, []PodEntry{pod}))
+	count, err := index.RedisClient.SCard(ctx, podEntriesKey).Result()
+	require.NoError(t, err)
+	require.EqualValues(t, 1, count)
+
+	require.NoError(t, index.Evict(ctx, requestKey, RequestKey, []PodEntry{pod}))
+
+	count, err = index.RedisClient.SCard(ctx, podEntriesKey).Result()
+	require.NoError(t, err)
+	require.Zero(t, count)
+}
+
+func TestRedisClearDropsMalformedReverseEntry(t *testing.T) {
+	index := createRedisIndexForInspection(t)
+	ctx := t.Context()
+	podEntriesKey := "kvblock:pod:pod-malformed:entries"
+
+	require.NoError(t, index.RedisClient.SAdd(ctx, podEntriesKey, "malformed").Err())
+	require.NoError(t, index.Clear(ctx, "pod-malformed"))
+
+	count, err := index.RedisClient.SCard(ctx, podEntriesKey).Result()
+	require.NoError(t, err)
+	require.Zero(t, count)
+}
+
+func redisTestPodField(t *testing.T, entry PodEntry) string {
+	t.Helper()
+	value, err := json.Marshal(entry)
+	require.NoError(t, err)
+	return string(value)
+}

tests/profiling/kv_cache_index/README.md

@@ -35,6 +35,44 @@ use the -bench option to filter
 go test -bench=Redis -benchmem
 ```
 
+### Redis Clear benchmarks
+
+The Redis Clear benchmark has two modes:
+
+* Default mode uses `miniredis`, so it needs no external Redis server and is suitable for quick local checks.
+* Real Redis mode is guarded by the `redis_real_bench` build tag and requires `KV_CACHE_BENCH_REDIS_ADDR`.
+
+Run the default `miniredis` benchmark:
+
+```bash
+go test ./tests/profiling/kv_cache_index \
+  -run '^$' \
+  -bench 'BenchmarkRedisClear(SharedKeyspace|PodFanout|Scale)' \
+  -benchmem \
+  -benchtime=3x \
+  -count=1
+```
+
+Run the same Clear scenarios against a real Redis server:
+
+```bash
+docker run --rm --name llmd-redis-bench \
+  -p 127.0.0.1:6380:6379 \
+  redis:7.4.1 redis-server --save '' --appendonly no
+```
+
+In another shell:
+
+```bash
+KV_CACHE_BENCH_REDIS_ADDR='redis://127.0.0.1:6380' \
+go test -tags redis_real_bench ./tests/profiling/kv_cache_index \
+  -run '^$' \
+  -bench 'BenchmarkRedisClearReal(Scale|SharedKeyspace|PodFanout)' \
+  -benchmem \
+  -benchtime=3x \
+  -count=1
+```
+
 ### Understanding the Output
 
 `BenchmarkInMemory_Add-12      192   6086106 ns/op    500 B/op      5 allocs/op`