diff --git a/internal/engine/kvrocks/delete_integration_test.go b/internal/engine/kvrocks/delete_integration_test.go new file mode 100644 index 0000000000..dcdda8a605 --- /dev/null +++ b/internal/engine/kvrocks/delete_integration_test.go @@ -0,0 +1,128 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +//go:build integration + +package kvrocks + +import ( + "context" + "os" + "testing" + + "github.com/alicebob/miniredis/v2" +) + +// TestKvrocksDeleteMany exercises the real go-redis DEL path behind +// DeleteMany: a single call removes every named key, a key that no longer +// exists is not an error, and an empty key list is a safe no-op (go-redis +// rejects a zero-argument DEL). +func TestKvrocksDeleteMany(t *testing.T) { + resetKvrocksState() + addr := os.Getenv("KVROCKS_TEST_ADDR") + if addr == "" { + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("miniredis: %v", err) + } + t.Cleanup(mr.Close) + addr = mr.Addr() + } + loadConfigForTest(t, addr) + if err := Init(context.Background()); err != nil { + t.Fatalf("kvrocks.Init: %v", err) + } + defer Close() + + ctx := context.Background() + c := Get() + if c == nil { + t.Fatal("kvrocks client should be initialized") + } + + // Batch remove: set three keys, drop them all in one call. + for _, k := range []string{"dm:a", "dm:b", "dm:c"} { + if !c.Set(ctx, k, "v", 0) { + t.Fatalf("seed %s failed", k) + } + } + if !c.DeleteMany(ctx, "dm:a", "dm:b", "dm:c") { + t.Fatal("DeleteMany returned false for a successful batch") + } + for _, k := range []string{"dm:a", "dm:b", "dm:c"} { + if v, _ := c.Get(ctx, k); v != "" { + t.Errorf("key %s not removed by batch delete", k) + } + } + + // A vanished key in the batch must not fail the call. + if !c.DeleteMany(ctx, "dm:a", "dm:gone") { + t.Error("DeleteMany should tolerate a missing key") + } + + // Empty key list is a no-op and must not error (zero-arg DEL is rejected). + if !c.DeleteMany(ctx) { + t.Error("DeleteMany with no keys should be a successful no-op") + } +} + +// TestKvrocksDeleteManyChunks verifies DeleteMany splits a key list larger than +// deleteBatchSize into multiple DEL calls and still removes every key. Seeding +// five keys with a batch size of two forces three batches; a wrong chunk boundary +// (e.g. a dropped trailing partial batch) would leave keys alive. +func TestKvrocksDeleteManyChunks(t *testing.T) { + resetKvrocksState() + addr := os.Getenv("KVROCKS_TEST_ADDR") + if addr == "" { + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("miniredis: %v", err) + } + t.Cleanup(mr.Close) + addr = mr.Addr() + } + loadConfigForTest(t, addr) + if err := Init(context.Background()); err != nil { + t.Fatalf("kvrocks.Init: %v", err) + } + defer Close() + + // Shrink the batch so a handful of keys spans several DEL calls. + prev := deleteBatchSize + deleteBatchSize = 2 + defer func() { deleteBatchSize = prev }() + + ctx := context.Background() + c := Get() + if c == nil { + t.Fatal("kvrocks client should be initialized") + } + + keys := []string{"ck:a", "ck:b", "ck:c", "ck:d", "ck:e"} + for _, k := range keys { + if !c.Set(ctx, k, "v", 0) { + t.Fatalf("seed %s failed", k) + } + } + if !c.DeleteMany(ctx, keys...) { + t.Fatal("DeleteMany returned false across multiple batches") + } + for _, k := range keys { + if v, _ := c.Get(ctx, k); v != "" { + t.Errorf("key %s not removed by chunked delete", k) + } + } +} diff --git a/internal/engine/kvrocks/delete_many_test.go b/internal/engine/kvrocks/delete_many_test.go new file mode 100644 index 0000000000..a686adff57 --- /dev/null +++ b/internal/engine/kvrocks/delete_many_test.go @@ -0,0 +1,99 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package kvrocks + +import ( + "reflect" + "testing" +) + +// TestSplitIntoBatches exercises the chunking math used by DeleteMany: every +// key must appear exactly once, input order must be preserved within and across +// chunks, and the boundary cases (empty, exact multiple, trailing partial +// batch, size 1, non-positive size) must split without dropping or duplicating. +// A naive split that computed the end index wrong would drop the trailing +// partial batch or emit an out-of-range slice. +func TestSplitIntoBatches(t *testing.T) { + cases := []struct { + name string + keys []string + size int + want [][]string + }{ + { + name: "empty", + keys: nil, + size: 1024, + want: nil, + }, + { + name: "single key within a large batch", + keys: []string{"a"}, + size: 1024, + want: [][]string{{"a"}}, + }, + { + name: "exact single batch", + keys: []string{"a", "b", "c"}, + size: 3, + want: [][]string{{"a", "b", "c"}}, + }, + { + name: "one element over the boundary spills into a second batch", + keys: []string{"a", "b", "c", "d"}, + size: 3, + want: [][]string{{"a", "b", "c"}, {"d"}}, + }, + { + name: "trailing partial batch is kept, not dropped", + keys: []string{"a", "b", "c", "d", "e"}, + size: 3, + want: [][]string{{"a", "b", "c"}, {"d", "e"}}, + }, + { + name: "exact two full batches", + keys: []string{"a", "b", "c", "d", "e", "f"}, + size: 3, + want: [][]string{{"a", "b", "c"}, {"d", "e", "f"}}, + }, + { + name: "size one yields one key per chunk", + keys: []string{"a", "b", "c"}, + size: 1, + want: [][]string{{"a"}, {"b"}, {"c"}}, + }, + { + name: "non-positive size degrades to one key per chunk", + keys: []string{"a", "b", "c"}, + size: 0, + want: [][]string{{"a"}, {"b"}, {"c"}}, + }, + { + name: "negative size degrades to one key per chunk", + keys: []string{"a", "b", "c"}, + size: -1, + want: [][]string{{"a"}, {"b"}, {"c"}}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := splitIntoBatches(tc.keys, tc.size) + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("splitIntoBatches(%v, %d) = %v, want %v", tc.keys, tc.size, got, tc.want) + } + }) + } +} diff --git a/internal/engine/kvrocks/redis.go b/internal/engine/kvrocks/redis.go index 6198e9917c..173ab1751b 100644 --- a/internal/engine/kvrocks/redis.go +++ b/internal/engine/kvrocks/redis.go @@ -927,6 +927,61 @@ func (r *Client) Delete(ctx context.Context, key string) bool { return true } +// deleteBatchSize bounds a single DEL command issued by DeleteMany. The previous +// unbounded DEL scaled its argument list with the task manifest size; splitting +// keeps each command small (avoids Redis/kvrocks proto-size limits and long +// single-threaded stalls) at the cost of losing strict atomicity across the whole +// batch. That loss is acceptable because cache reclaim is best-effort and TTL is +// the fallback. Declared as a var (not a const) so tests can shrink it to +// exercise the multi-batch path; production always uses the default. +var deleteBatchSize = 1024 + +// splitIntoBatches partitions keys into chunks of at most size, preserving input +// order within and across chunks. A non-positive size degrades to one key per +// chunk so the helper never emits an infinite or empty loop. No key is dropped +// or duplicated; an empty input yields nil. +func splitIntoBatches(keys []string, size int) [][]string { + if size <= 0 { + size = 1 + } + if len(keys) == 0 { + return nil + } + out := make([][]string, 0, (len(keys)+size-1)/size) + for start := 0; start < len(keys); start += size { + end := start + size + if end > len(keys) { + end = len(keys) + } + out = append(out, keys[start:end]) + } + return out +} + +// DeleteMany deletes multiple keys by issuing one DEL per batch of at most +// deleteBatchSize keys. A key that no longer exists is not an error (DEL reports +// only the count removed), so a purge that races a TTL expiration still completes +// instead of failing on a vanished key. An empty key list is a no-op: go-redis +// rejects a zero-argument DEL, and there is nothing to remove anyway. The overall +// result is false if any batch fails, but batches already issued are not rolled +// back — TTL reclaims the rest. +func (r *Client) DeleteMany(ctx context.Context, keys ...string) bool { + if r.client == nil { + return false + } + if len(keys) == 0 { + return true + } + ok := true + for _, batch := range splitIntoBatches(keys, deleteBatchSize) { + if err := r.client.Del(ctx, batch...).Err(); err != nil { + common.Warn("Redis DeleteMany error", zap.Strings("keys", batch), zap.Error(err)) + ok = false + } + } + return ok +} + // Expire sets expiration on a key func (r *Client) Expire(ctx context.Context, key string, exp time.Duration) bool { if r.client == nil { diff --git a/internal/ingestion/chunkcache/chunkcache.go b/internal/ingestion/chunkcache/chunkcache.go index f691e9274b..2d5aa0e657 100644 --- a/internal/ingestion/chunkcache/chunkcache.go +++ b/internal/ingestion/chunkcache/chunkcache.go @@ -59,6 +59,7 @@ type Store interface { SAdd(ctx context.Context, key, member string) bool SMembers(ctx context.Context, key string) ([]string, error) Delete(ctx context.Context, key string) bool + DeleteMany(ctx context.Context, keys ...string) bool Expire(ctx context.Context, key string, exp time.Duration) bool } @@ -180,9 +181,16 @@ func Set(ctx context.Context, s Store, key, value string) { // the task row must not make resumable cleanup look complete while Redis is // unavailable. // -// The manifest is deleted last: losing it first would orphan the entries it -// still lists, leaving them to expire. Best-effort — a failed reclaim only -// means the entries wait out their TTL. +// Entries are reclaimed in one batched DeleteMany; the manifest is deleted on +// its own (a separate Delete) only after that batch succeeds. This ordering +// keeps the manifest intact when the entry batch fails or is only partially +// applied (a large manifest is split across several DELs, and a later batch may +// fail after an earlier one already removed its keys) so the task stays +// resumable instead of orphaning surviving entries under a deleted manifest. +// Any entry a failed batch already dropped is reclaimed by its TTL. A key that +// has already expired is not an error (DEL reports only the count removed), so +// a purge that races TTL eviction still completes instead of failing on a +// vanished entry. func PurgeTask(ctx context.Context, s Store, taskID string) error { if s == nil || taskID == "" { return nil @@ -192,13 +200,11 @@ func PurgeTask(ctx context.Context, s Store, taskID string) error { if err != nil { return fmt.Errorf("read chunk cache manifest %s: %w", mk, err) } - for _, k := range keys { - if !s.Delete(ctx, k) { - return fmt.Errorf("delete chunk cache entry %s", k) - } + if !s.DeleteMany(ctx, keys...) { + return fmt.Errorf("purge chunk cache entries for %s", taskID) } if !s.Delete(ctx, mk) { - return fmt.Errorf("delete chunk cache manifest %s", mk) + return fmt.Errorf("delete chunk cache manifest %s", taskID) } return nil } diff --git a/internal/ingestion/chunkcache/chunkcache_test.go b/internal/ingestion/chunkcache/chunkcache_test.go index 06a07527c7..4ffa8f90eb 100644 --- a/internal/ingestion/chunkcache/chunkcache_test.go +++ b/internal/ingestion/chunkcache/chunkcache_test.go @@ -31,14 +31,16 @@ import ( // manifest bookkeeping, with call counters so tests can assert the exact // command sequence (a real Redis is an integration-tier dependency). type fakeStore struct { - kv map[string]string - sets map[string]map[string]bool - ttl map[string]time.Duration - deleted []string - failSet bool - expireFails bool - membersErr error - deleteFails map[string]bool + kv map[string]string + sets map[string]map[string]bool + ttl map[string]time.Duration + deleted []string + deleteManyCalls int + lastDeleteMany []string + failSet bool + expireFails bool + membersErr error + deleteFails map[string]bool } func newFakeStore() *fakeStore { @@ -96,6 +98,22 @@ func (f *fakeStore) Delete(_ context.Context, key string) bool { return true } +func (f *fakeStore) DeleteMany(_ context.Context, keys ...string) bool { + f.deleteManyCalls++ + f.lastDeleteMany = append([]string(nil), keys...) + for _, k := range keys { + if f.deleteFails[k] { + return false + } + } + for _, k := range keys { + delete(f.kv, k) + delete(f.sets, k) + delete(f.ttl, k) + } + return true +} + func (f *fakeStore) Expire(_ context.Context, key string, exp time.Duration) bool { if f.expireFails { return false @@ -240,9 +258,9 @@ func TestGet_HitAndMiss(t *testing.T) { } // TestPurgeTask_DeletesEveryEntryThenManifest asserts persist-time reclaim -// removes all recorded entries and finally the manifest itself, leaving no -// key behind that would otherwise wait out the 7-day TTL. -func TestPurgeTask_DeletesEveryEntryThenManifest(t *testing.T) { +// removes all recorded entries in one batch and then the manifest itself, +// leaving no key behind that would otherwise wait out the 7-day TTL. +func TestPurgeTaskDeletesEveryEntryAndManifest(t *testing.T) { f := newFakeStore() ctx := taskCtx("task-1") Set(ctx, f, "kc:extractor:keywords:a", "1") @@ -251,21 +269,25 @@ func TestPurgeTask_DeletesEveryEntryThenManifest(t *testing.T) { PurgeTask(context.Background(), f, "task-1") - want := []string{"kc:emb:c", "kc:extractor:keywords:a", "kc:meta:b", "kc:manifest:task-1"} - if len(f.deleted) != len(want) { - t.Fatalf("deleted = %v, want %v", f.deleted, want) + // Entries are reclaimed in one batch; the manifest is deleted separately + // (so a failed entry batch leaves the task resumable). Nothing may remain. + if f.deleteManyCalls != 1 { + t.Fatalf("DeleteMany calls = %d, want 1 (single batch)", f.deleteManyCalls) } - // Entry deletes may arrive in any order, but the manifest must be last: - // dropping it early would orphan the remaining entries if we crash midway. - if f.deleted[len(f.deleted)-1] != "kc:manifest:task-1" { - t.Errorf("last delete = %q, want the manifest key", f.deleted[len(f.deleted)-1]) + if len(f.deleted) != 1 || f.deleted[0] != manifestKey("task-1") { + t.Fatalf("per-key Delete = %v, want exactly the manifest key", f.deleted) } - entries := append([]string(nil), f.deleted[:len(f.deleted)-1]...) - sort.Strings(entries) - for i, w := range want[:len(want)-1] { - if entries[i] != w { - t.Errorf("entry deletes = %v, want %v", entries, want[:len(want)-1]) - break + want := map[string]bool{ + "kc:emb:c": true, + "kc:extractor:keywords:a": true, + "kc:meta:b": true, + } + if len(f.lastDeleteMany) != len(want) { + t.Fatalf("batch size = %d, want %d (%v)", len(f.lastDeleteMany), len(want), f.lastDeleteMany) + } + for _, k := range f.lastDeleteMany { + if !want[k] { + t.Errorf("batch carried unexpected key %q", k) } } if len(f.kv) != 0 { diff --git a/internal/ingestion/chunkcache/e2_purge_batch_test.go b/internal/ingestion/chunkcache/e2_purge_batch_test.go new file mode 100644 index 0000000000..7500163a4b --- /dev/null +++ b/internal/ingestion/chunkcache/e2_purge_batch_test.go @@ -0,0 +1,150 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package chunkcache + +import ( + "context" + "fmt" + "testing" +) + +// TestPurgeTaskUsesSingleBatchDelete pins the E2 batch call shape: entries are +// reclaimed in exactly ONE Store.DeleteMany call, and the manifest is deleted +// separately (one per-key Delete) ONLY after that batch succeeds. A naive +// loop-based implementation would pass a happy-path test, so we assert the +// batch carries the entries but NOT the manifest, and that the manifest's own +// Delete is the sole per-key call. +func TestPurgeTaskUsesSingleBatchDelete(t *testing.T) { + f := newFakeStore() + ctx := taskCtx("task-1") + entryKeys := []string{"kc:extractor:keywords:a", "kc:extractor:keywords:b", "kc:extractor:keywords:c"} + for i, k := range entryKeys { + Set(ctx, f, k, fmt.Sprintf("v%d", i)) + } + + if err := PurgeTask(context.Background(), f, "task-1"); err != nil { + t.Fatalf("PurgeTask error = %v, want nil", err) + } + + // Entries travel in one batch. + if f.deleteManyCalls != 1 { + t.Errorf("DeleteMany calls = %d, want 1 (single batch)", f.deleteManyCalls) + } + // Exactly one per-key Delete, for the manifest only. + if len(f.deleted) != 1 || f.deleted[0] != manifestKey("task-1") { + t.Errorf("per-key Delete = %v, want exactly the manifest key", f.deleted) + } + if len(f.kv) != 0 { + t.Errorf("kv not fully reclaimed: %v", f.kv) + } + if _, ok := f.sets[manifestKey("task-1")]; ok { + t.Errorf("manifest %q not removed", manifestKey("task-1")) + } + + // The batch must carry every entry key, but NOT the manifest: the manifest + // is deleted on its own after the batch succeeds (so a failed batch leaves + // it intact and the task resumable). + want := make(map[string]bool, len(entryKeys)) + for _, k := range entryKeys { + want[k] = true + } + if len(f.lastDeleteMany) != len(want) { + t.Fatalf("batch size = %d, want %d (%v)", len(f.lastDeleteMany), len(want), f.lastDeleteMany) + } + for _, k := range f.lastDeleteMany { + if !want[k] { + t.Errorf("batch carried unexpected key %q", k) + } + if k == manifestKey("task-1") { + t.Errorf("batch must not include the manifest key; it is deleted separately") + } + } +} + +// TestPurgeTaskBatchHandlesEmptyManifest asserts that even a task whose manifest +// lists no entries still removes the manifest key in one batch call. The previous +// loop would iterate zero times and then issue a separate Delete for the +// manifest; the batch form must not regress to that two-step shape. +func TestPurgeTaskBatchHandlesEmptyManifest(t *testing.T) { + f := newFakeStore() + + if err := PurgeTask(context.Background(), f, "task-1"); err != nil { + t.Fatalf("PurgeTask error = %v, want nil", err) + } + if f.deleteManyCalls != 1 { + t.Errorf("DeleteMany calls = %d, want 1 for an empty manifest", f.deleteManyCalls) + } + if _, ok := f.sets[manifestKey("task-1")]; ok { + t.Errorf("manifest %q not removed", manifestKey("task-1")) + } +} + +// TestPurgeTaskBatchPropagatesFailure asserts that when an entry batch cannot be +// fully deleted, PurgeTask returns an error and leaves the manifest intact +// rather than deleting it alongside a partially-reclaimed batch. +func TestPurgeTaskBatchPropagatesFailure(t *testing.T) { + f := newFakeStore() + ctx := taskCtx("task-1") + key := "kc:extractor:keywords:a" + Set(ctx, f, key, "result") + f.deleteFails = map[string]bool{key: true} + + if err := PurgeTask(context.Background(), f, "task-1"); err == nil { + t.Fatal("PurgeTask error = nil, want failure when an entry cannot be deleted") + } + if _, ok := f.sets[manifestKey("task-1")]; !ok { + t.Fatal("manifest must survive a failed entry batch (resumable)") + } +} + +// TestPurgeTaskManifestSurvivesEntryFailure is the adversarial test for the +// large-manifest safety property: when an entry batch reports failure (e.g. a +// store outage that hits a later DEL batch after an earlier one already +// succeeded), PurgeTask must NOT delete the manifest, so the task stays +// resumable and the surviving entries are reclaimed on retry. This must hold +// regardless of how many entry batches the manifest is split into. +func TestPurgeTaskManifestSurvivesEntryFailure(t *testing.T) { + f := newFakeStore() + ctx := taskCtx("task-1") + entryKeys := make([]string, 0, 5) + for i := 0; i < 5; i++ { + k := fmt.Sprintf("kc:extractor:keywords:%d", i) + Set(ctx, f, k, fmt.Sprintf("v%d", i)) + entryKeys = append(entryKeys, k) + } + // Fail one entry so the entry batch reports failure. + f.deleteFails = map[string]bool{entryKeys[3]: true} + + if err := PurgeTask(context.Background(), f, "task-1"); err == nil { + t.Fatal("PurgeTask error = nil, want failure when an entry cannot be deleted") + } + // Manifest must survive so the task remains resumable. + if _, ok := f.sets[manifestKey("task-1")]; !ok { + t.Fatal("manifest must survive a failed entry batch") + } +} + +// TestPurgeTaskBatchNoopWithoutTaskOrStore mirrors the existing no-op guard and +// pins that the batch path issues no DeleteMany when there is nothing to purge. +func TestPurgeTaskBatchNoopWithoutTaskOrStore(t *testing.T) { + f := newFakeStore() + PurgeTask(context.Background(), f, "") + if f.deleteManyCalls != 0 { + t.Errorf("DeleteMany calls = %d, want 0 for an empty task id", f.deleteManyCalls) + } + PurgeTask(context.Background(), nil, "task-1") +} diff --git a/internal/ingestion/component/delete_many_fake_test.go b/internal/ingestion/component/delete_many_fake_test.go new file mode 100644 index 0000000000..3501dfd221 --- /dev/null +++ b/internal/ingestion/component/delete_many_fake_test.go @@ -0,0 +1,76 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package component + +import ( + "context" + "testing" + "time" +) + +// TestMemCacheStoreDeleteMany verifies the batch delete clears both the data +// map and the set map (mirroring Delete) and reports success. A naive +// DeleteMany that only clears data would leave set members orphaned. +func TestMemCacheStoreDeleteMany(t *testing.T) { + m := newMemCacheStore() + m.mu.Lock() + m.data["a"] = "1" + m.data["b"] = "2" + m.sets["a"] = map[string]struct{}{"x": {}} + m.sets["b"] = map[string]struct{}{"y": {}} + m.mu.Unlock() + + if !m.DeleteMany(context.Background(), "a", "b") { + t.Fatal("DeleteMany returned false") + } + + m.mu.Lock() + defer m.mu.Unlock() + for _, k := range []string{"a", "b"} { + if _, ok := m.data[k]; ok { + t.Errorf("data[%s] not deleted", k) + } + if _, ok := m.sets[k]; ok { + t.Errorf("sets[%s] not deleted", k) + } + } +} + +// TestFakeCacheStoreDeleteMany verifies DeleteMany clears both kv and sets. +func TestFakeCacheStoreDeleteMany(t *testing.T) { + f := &fakeCacheStore{ + kv: map[string]string{}, + ttl: map[string]time.Duration{}, + sets: map[string]map[string]bool{}, + } + f.kv["a"] = "1" + f.kv["b"] = "2" + f.sets["a"] = map[string]bool{"x": true} + f.sets["b"] = map[string]bool{"y": true} + + if !f.DeleteMany(context.Background(), "a", "b") { + t.Fatal("DeleteMany returned false") + } + + for _, k := range []string{"a", "b"} { + if _, ok := f.kv[k]; ok { + t.Errorf("kv[%s] not deleted", k) + } + if _, ok := f.sets[k]; ok { + t.Errorf("sets[%s] not deleted", k) + } + } +} diff --git a/internal/ingestion/component/extractor_cache_test.go b/internal/ingestion/component/extractor_cache_test.go index 0c4e1af452..7a7b3d6f46 100644 --- a/internal/ingestion/component/extractor_cache_test.go +++ b/internal/ingestion/component/extractor_cache_test.go @@ -97,6 +97,16 @@ func (f *fakeCacheStore) Delete(_ context.Context, key string) bool { return true } +func (f *fakeCacheStore) DeleteMany(_ context.Context, keys ...string) bool { + f.mu.Lock() + defer f.mu.Unlock() + for _, k := range keys { + delete(f.kv, k) + delete(f.sets, k) + } + return true +} + func (f *fakeCacheStore) Expire(_ context.Context, key string, exp time.Duration) bool { f.mu.Lock() defer f.mu.Unlock() diff --git a/internal/ingestion/component/tokenizer_cache_test.go b/internal/ingestion/component/tokenizer_cache_test.go index df29a1e864..4d4b405c31 100644 --- a/internal/ingestion/component/tokenizer_cache_test.go +++ b/internal/ingestion/component/tokenizer_cache_test.go @@ -94,6 +94,16 @@ func (m *memCacheStore) Delete(_ context.Context, key string) bool { return true } +func (m *memCacheStore) DeleteMany(_ context.Context, keys ...string) bool { + m.mu.Lock() + defer m.mu.Unlock() + for _, k := range keys { + delete(m.data, k) + delete(m.sets, k) + } + return true +} + func (m *memCacheStore) Expire(_ context.Context, key string, _ time.Duration) bool { return true } diff --git a/internal/ingestion/mediacache/e2_purge_batch_test.go b/internal/ingestion/mediacache/e2_purge_batch_test.go new file mode 100644 index 0000000000..3e5af92870 --- /dev/null +++ b/internal/ingestion/mediacache/e2_purge_batch_test.go @@ -0,0 +1,153 @@ +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mediacache + +import ( + "context" + "testing" +) + +// TestPurgeTaskUsesSingleBatchDelete pins the E2 batch call shape in the media +// cache: every entry key plus the quota (digest/state) keys travel in exactly +// ONE Store.DeleteMany call, and the manifest is deleted separately (one +// per-key Delete) only after that batch succeeds. A loop-based implementation +// would pass a happy-path test, so we assert the batch carries the entries and +// quota but NOT the manifest, and the manifest's own Delete is the sole +// per-key call. +func TestPurgeTaskUsesSingleBatchDelete(t *testing.T) { + backend := newFakeStore() + ctx := taskContext("task-1") + store := New(backend, "task-1") + ref, err := store.Put(ctx, "html", "image/png", []byte("hello")) + if err != nil { + t.Fatal(err) + } + + if err := PurgeTask(ctx, backend, "task-1"); err != nil { + t.Fatalf("PurgeTask error = %v, want nil", err) + } + + if backend.deleteManyCalls != 1 { + t.Errorf("DeleteMany calls = %d, want 1 (single batch)", backend.deleteManyCalls) + } + // Exactly one per-key Delete, for the manifest only. + if len(backend.deleted) != 1 || backend.deleted[0] != manifestKey("task-1") { + t.Errorf("per-key Delete = %v, want exactly the manifest key", backend.deleted) + } + if len(backend.kv) != 0 || len(backend.sets) != 0 { + t.Errorf("leftover after purge kv=%v sets=%v", backend.kv, backend.sets) + } + + // The batch must carry the entry key plus the quota keys, but NOT the + // manifest (deleted separately after the batch succeeds). + digestKey, stateKey := quotaKeys("task-1") + entryKey := store.key(ref.SHA256) + required := []string{entryKey, digestKey, stateKey} + for _, k := range required { + found := false + for _, got := range backend.lastDeleteMany { + if got == k { + found = true + break + } + } + if !found { + t.Errorf("batch did not include required key %q (batch=%v)", k, backend.lastDeleteMany) + } + } + for _, got := range backend.lastDeleteMany { + if got == manifestKey("task-1") { + t.Errorf("batch must not include the manifest key; it is deleted separately") + } + } +} + +// TestPurgeTaskBatchHandlesEmptyManifest asserts an empty media manifest is still +// reclaimed (manifest + quota keys) in one batch call, not a zero-iteration loop +// followed by separate deletes. +func TestPurgeTaskBatchHandlesEmptyManifest(t *testing.T) { + backend := newFakeStore() + ctx := taskContext("task-1") + + if err := PurgeTask(ctx, backend, "task-1"); err != nil { + t.Fatalf("PurgeTask error = %v, want nil", err) + } + if backend.deleteManyCalls != 1 { + t.Errorf("DeleteMany calls = %d, want 1 for an empty manifest", backend.deleteManyCalls) + } + digestKey, stateKey := quotaKeys("task-1") + if _, ok := backend.kv[digestKey]; ok { + t.Errorf("quota digest key %q not removed", digestKey) + } + if _, ok := backend.kv[stateKey]; ok { + t.Errorf("quota state key %q not removed", stateKey) + } +} + +// TestPurgeTaskBatchPropagatesFailure asserts that when an entry batch fails, the +// media manifest is preserved (resumable), not deleted alongside a partially +// reclaimed batch. +func TestPurgeTaskBatchPropagatesFailure(t *testing.T) { + backend := newFakeStore() + ctx := taskContext("task-1") + store := New(backend, "task-1") + ref, err := store.Put(ctx, "html", "image/png", []byte("first")) + if err != nil { + t.Fatal(err) + } + key := store.key(ref.SHA256) + backend.deleteFails = map[string]bool{key: true} + + if err := PurgeTask(ctx, backend, "task-1"); err == nil { + t.Fatal("expected error when an entry cannot be deleted") + } + if len(backend.sets[manifestKey("task-1")]) == 0 { + t.Fatal("manifest must survive a failed entry batch (resumable)") + } +} + +// TestPurgeTaskManifestSurvivesEntryFailure asserts that when an entry batch +// fails, the media manifest is preserved (resumable), not deleted alongside a +// partially-reclaimed batch. This must hold regardless of how many DEL batches +// the entry list is split into. +func TestPurgeTaskManifestSurvivesEntryFailure(t *testing.T) { + backend := newFakeStore() + ctx := taskContext("task-1") + store := New(backend, "task-1") + ref, err := store.Put(ctx, "html", "image/png", []byte("first")) + if err != nil { + t.Fatal(err) + } + key := store.key(ref.SHA256) + backend.deleteFails = map[string]bool{key: true} + + if err := PurgeTask(ctx, backend, "task-1"); err == nil { + t.Fatal("expected error when an entry cannot be deleted") + } + if len(backend.sets[manifestKey("task-1")]) == 0 { + t.Fatal("manifest must survive a failed entry batch (resumable)") + } +} + +// TestPurgeTaskBatchNoopWithoutTaskOrStore pins that the batch path issues no +// DeleteMany when there is nothing to purge. +func TestPurgeTaskBatchNoopWithoutTaskOrStore(t *testing.T) { + backend := newFakeStore() + PurgeTask(context.Background(), backend, "") + if backend.deleteManyCalls != 0 { + t.Errorf("DeleteMany calls = %d, want 0 for an empty task id", backend.deleteManyCalls) + } + PurgeTask(context.Background(), nil, "task-1") +} diff --git a/internal/ingestion/mediacache/store.go b/internal/ingestion/mediacache/store.go index c9c477a240..339c64b486 100644 --- a/internal/ingestion/mediacache/store.go +++ b/internal/ingestion/mediacache/store.go @@ -208,6 +208,13 @@ func quotaKeys(scope string) (digestKey, stateKey string) { } // PurgeTask removes media after durable persistence. Failed runs rely on TTL. +// +// Entries and quota keys are reclaimed in one batched DeleteMany; the manifest +// is deleted on its own (a separate Delete) only after that batch succeeds. This +// keeps the manifest intact if the entry batch fails or is only partially +// applied (a large manifest is split across several DELs, and a later batch may +// fail after an earlier one already removed its keys), so the task stays +// resumable instead of orphaning surviving entries under a deleted manifest. func PurgeTask(ctx context.Context, backend chunkcache.Store, taskID string) error { if backend == nil || taskID == "" { return nil @@ -217,19 +224,14 @@ func PurgeTask(ctx context.Context, backend chunkcache.Store, taskID string) err if err != nil { return fmt.Errorf("mediacache: read manifest: %w", err) } - for _, key := range keys { - if !backend.Delete(ctx, key) { - return fmt.Errorf("mediacache: delete %s", key) - } + digestKey, stateKey := quotaKeys(taskID) + batch := append([]string{}, keys...) + batch = append(batch, digestKey, stateKey) + if !backend.DeleteMany(ctx, batch...) { + return fmt.Errorf("mediacache: purge task %s", taskID) } if !backend.Delete(ctx, mk) { - return fmt.Errorf("mediacache: delete manifest %s", mk) - } - digestKey, stateKey := quotaKeys(taskID) - for _, key := range []string{digestKey, stateKey} { - if !backend.Delete(ctx, key) { - return fmt.Errorf("mediacache: delete quota state %s", key) - } + return fmt.Errorf("mediacache: delete manifest %s", taskID) } return nil } diff --git a/internal/ingestion/mediacache/store_test.go b/internal/ingestion/mediacache/store_test.go index 6aacd586ac..bed57261bb 100644 --- a/internal/ingestion/mediacache/store_test.go +++ b/internal/ingestion/mediacache/store_test.go @@ -30,17 +30,20 @@ import ( ) type fakeStore struct { - kv map[string]string - sets map[string]map[string]bool - ttl map[string]time.Duration - getErr error - failSet bool - expireFails bool - setCalls int - addCalls int - mediaPuts int - deleteFails map[string]bool - quotaBytes map[string]int64 + kv map[string]string + sets map[string]map[string]bool + ttl map[string]time.Duration + getErr error + failSet bool + expireFails bool + setCalls int + addCalls int + mediaPuts int + deleted []string + deleteFails map[string]bool + deleteManyCalls int + lastDeleteMany []string + quotaBytes map[string]int64 } func newFakeStore() *fakeStore { @@ -151,6 +154,7 @@ func (f *fakeStore) SMembers(_ context.Context, key string) ([]string, error) { } func (f *fakeStore) Delete(_ context.Context, key string) bool { + f.deleted = append(f.deleted, key) if f.deleteFails[key] { return false } @@ -160,6 +164,23 @@ func (f *fakeStore) Delete(_ context.Context, key string) bool { return true } +func (f *fakeStore) DeleteMany(_ context.Context, keys ...string) bool { + f.deleteManyCalls++ + f.lastDeleteMany = append([]string(nil), keys...) + for _, k := range keys { + if f.deleteFails[k] { + return false + } + } + for _, k := range keys { + delete(f.kv, k) + delete(f.sets, k) + delete(f.quotaBytes, k) + delete(f.ttl, k) + } + return true +} + func (f *fakeStore) Expire(_ context.Context, key string, exp time.Duration) bool { if f.expireFails { return false diff --git a/internal/ingestion/pipeline/cleanup_test.go b/internal/ingestion/pipeline/cleanup_test.go index b9bff1f8c9..14ec949b49 100644 --- a/internal/ingestion/pipeline/cleanup_test.go +++ b/internal/ingestion/pipeline/cleanup_test.go @@ -134,6 +134,15 @@ func (f *cleanupFakeStore) Delete(_ context.Context, key string) bool { return true } +func (f *cleanupFakeStore) DeleteMany(_ context.Context, keys ...string) bool { + for _, k := range keys { + delete(f.kv, k) + delete(f.sets, k) + delete(f.bytes, k) + } + return true +} + func (f *cleanupFakeStore) Expire(_ context.Context, key string, exp time.Duration) bool { f.ttl[key] = exp return true diff --git a/internal/ingestion/pipeline/delete_many_fake_test.go b/internal/ingestion/pipeline/delete_many_fake_test.go new file mode 100644 index 0000000000..2e450773e7 --- /dev/null +++ b/internal/ingestion/pipeline/delete_many_fake_test.go @@ -0,0 +1,49 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pipeline + +import ( + "context" + "testing" +) + +// TestCleanupFakeStoreDeleteMany verifies DeleteMany clears the kv, sets and +// bytes maps together (mirroring Delete) and reports success. +func TestCleanupFakeStoreDeleteMany(t *testing.T) { + f := newCleanupFakeStore() + f.kv["a"] = "1" + f.kv["b"] = "2" + f.sets["a"] = map[string]bool{"x": true} + f.sets["b"] = map[string]bool{"y": true} + f.bytes["a"] = 10 + f.bytes["b"] = 20 + + if !f.DeleteMany(context.Background(), "a", "b") { + t.Fatal("DeleteMany returned false") + } + + for _, k := range []string{"a", "b"} { + if _, ok := f.kv[k]; ok { + t.Errorf("kv[%s] not deleted", k) + } + if _, ok := f.sets[k]; ok { + t.Errorf("sets[%s] not deleted", k) + } + if _, ok := f.bytes[k]; ok { + t.Errorf("bytes[%s] not deleted", k) + } + } +}