fix(kvblock): simplify Redis clear cleanup

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

Author: yankay

pkg/kvcache/kvblock/redis.go

@@ -154,29 +154,6 @@ var pruneRequestKeyScript = redis.NewScript(`
 	return 0
 `)
 
-// clearPodEntriesScript removes one batch of reverse-index members, deletes
-// matching request-key hash fields, and prunes request-key hashes that become
-// empty. Malformed members are only removed from the reverse-index set.
-var clearPodEntriesScript = redis.NewScript(`
-	local removed = 0
-	for i = 1, #ARGV do
-		local member = ARGV[i]
-		redis.call('SREM', KEYS[1], member)
-		local sep = string.find(member, string.char(0), 1, true)
-		if sep ~= nil then
-			local requestKey = string.sub(member, 1, sep - 1)
-			local field = string.sub(member, sep + 1)
-			if requestKey ~= '' and field ~= '' then
-				removed = removed + redis.call('HDEL', requestKey, field)
-				if redis.call('HLEN', requestKey) == 0 then
-					redis.call('DEL', requestKey)
-				end
-			end
-		end
-	end
-	return removed
-`)
-
 // pruneEngineKeyScript atomically deletes an engine key mapping only if all
 // associated request key hashes are empty. This prevents a TOCTOU race where a
 // concurrent Add could insert into a request key between checking and deleting.
@@ -441,6 +418,19 @@ 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
@@ -450,18 +440,17 @@ func (r *RedisIndex) Clear(ctx context.Context, podIdentifier string) error {
 	podEntriesKey := redisPodEntriesKey(podIdentifier)
 
 	const scanBatch int64 = 1024
-	removed := 0
+	processed := 0
 	var cursor uint64
 	for {
 		members, next, err := r.RedisClient.SScan(ctx, podEntriesKey, cursor, "", scanBatch).Result()
 		if err != nil {
 			return fmt.Errorf("clear reverse-index scan failed: %w", err)
 		}
-		count, err := r.clearRedisPodEntryMembers(ctx, podEntriesKey, members)
-		if err != nil {
+		if err := r.clearRedisPodEntryMembers(ctx, podEntriesKey, members); err != nil {
 			return err
 		}
-		removed += count
+		processed += len(members)
 
 		if cursor = next; cursor != 0 {
 			continue
@@ -475,22 +464,25 @@ func (r *RedisIndex) Clear(ctx context.Context, podIdentifier string) error {
 		}
 	}
 
-	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) (int, error) {
+func (r *RedisIndex) clearRedisPodEntryMembers(ctx context.Context, podEntriesKey string, members []string) error {
 	if len(members) == 0 {
-		return 0, nil
+		return nil
 	}
 
-	args := make([]interface{}, len(members))
-	for i, member := range members {
-		args[i] = member
+	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)
 	}
-	removed, err := clearPodEntriesScript.Run(ctx, r.RedisClient, []string{podEntriesKey}, args...).Int()
-	if err != nil && !errors.Is(err, redis.Nil) {
-		return 0, fmt.Errorf("clear reverse-index batch failed: %w", err)
+	if _, err := pipe.Exec(ctx); err != nil {
+		return fmt.Errorf("clear reverse-index pipeline failed: %w", err)
 	}
-	return removed, nil
+	return nil
 }

pkg/kvcache/kvblock/redis_test.go

@@ -106,6 +106,19 @@ func TestRedisEvictPrunesReverseIndex(t *testing.T) {
 	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)

tests/profiling/kv_cache_index/README.md

@@ -67,7 +67,7 @@ In another shell:
 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(SharedKeyspace|PodFanout)' \
+  -bench 'BenchmarkRedisClearReal(Scale|SharedKeyspace|PodFanout)' \
   -benchmem \
   -benchtime=3x \
   -count=1

tests/profiling/kv_cache_index/redis_clear_benchmark_test.go

@@ -30,7 +30,10 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 )
 
-const redisClearTargetPod = "pod-0"
+const (
+	redisClearTargetPod     = "pod-0"
+	redisClearSeedBatchSize = 1024
+)
 
 type redisClearBenchmarkCase struct {
 	name          string
@@ -166,8 +169,11 @@ func seedRedisClearReverseIndexData(
 	}
 	for pod := 0; pod < tc.podCount; pod++ {
 		entry := kvblock.PodEntry{PodIdentifier: fmt.Sprintf("pod-%d", pod), DeviceTier: "gpu"}
-		if err := index.Add(ctx, nil, requestKeys, []kvblock.PodEntry{entry}); err != nil {
-			b.Fatalf("failed to seed redis data through Add: %v", err)
+		for start := 0; start < len(requestKeys); start += redisClearSeedBatchSize {
+			end := min(start+redisClearSeedBatchSize, len(requestKeys))
+			if err := index.Add(ctx, nil, requestKeys[start:end], []kvblock.PodEntry{entry}); err != nil {
+				b.Fatalf("failed to seed redis data through Add: %v", err)
+			}
 		}
 	}
 }
@@ -180,16 +186,23 @@ func seedRedisClearLegacyData(
 ) {
 	b.Helper()
 	pipe := index.RedisClient.Pipeline()
+	pending := 0
 	for key := 0; key < tc.requestKeys; key++ {
 		redisKey := redisClearBlockHash(b, key+1).String()
 		for pod := 0; pod < tc.podCount; pod++ {
 			entry := kvblock.PodEntry{PodIdentifier: fmt.Sprintf("pod-%d", pod), DeviceTier: "gpu"}
 			field := redisClearPodField(b, entry)
 			pipe.HSet(ctx, redisKey, field, "")
+			pending++
+			if pending == redisClearSeedBatchSize {
+				execRedisClearSeedPipeline(b, ctx, pipe)
+				pipe = index.RedisClient.Pipeline()
+				pending = 0
+			}
 		}
 	}
-	if _, err := pipe.Exec(ctx); err != nil {
-		b.Fatalf("failed to seed redis data: %v", err)
+	if pending > 0 {
+		execRedisClearSeedPipeline(b, ctx, pipe)
 	}
 }
 
@@ -199,10 +212,24 @@ func seedUnrelatedRedisData(b *testing.B, ctx context.Context, index *kvblock.Re
 		return
 	}
 	pipe := index.RedisClient.Pipeline()
+	pending := 0
 	for key := 0; key < count; key++ {
 		field := redisClearPodField(b, kvblock.PodEntry{PodIdentifier: "other", DeviceTier: "gpu"})
 		pipe.HSet(ctx, fmt.Sprintf("unrelated:%d", key), field, "")
+		pending++
+		if pending == redisClearSeedBatchSize {
+			execRedisClearSeedPipeline(b, ctx, pipe)
+			pipe = index.RedisClient.Pipeline()
+			pending = 0
+		}
+	}
+	if pending > 0 {
+		execRedisClearSeedPipeline(b, ctx, pipe)
 	}
+}
+
+func execRedisClearSeedPipeline(b *testing.B, ctx context.Context, pipe redis.Pipeliner) {
+	b.Helper()
 	if _, err := pipe.Exec(ctx); err != nil {
 		b.Fatalf("failed to seed redis data: %v", err)
 	}

tests/profiling/kv_cache_index/redis_clear_real_benchmark_test.go

@@ -25,6 +25,14 @@ import (
 
 const redisClearRealAddressEnv = "KV_CACHE_BENCH_REDIS_ADDR"
 
+func BenchmarkRedisClearRealScaleLegacyFullScan(b *testing.B) {
+	benchmarkRedisClearRealCases(b, false, redisClearScaleCases)
+}
+
+func BenchmarkRedisClearRealScaleReverseIndex(b *testing.B) {
+	benchmarkRedisClearRealCases(b, true, redisClearScaleCases)
+}
+
 func BenchmarkRedisClearRealSharedKeyspaceLegacyFullScan(b *testing.B) {
 	benchmarkRedisClearRealCases(b, false, redisClearSharedKeyspaceCases)
 }