diff --git a/docs/agent_guides/streaming-system/coordination/broadcaster.md b/docs/agent_guides/streaming-system/coordination/broadcaster.md index cafc1c53af9..10a091d5708 100644 --- a/docs/agent_guides/streaming-system/coordination/broadcaster.md +++ b/docs/agent_guides/streaming-system/coordination/broadcaster.md @@ -17,6 +17,8 @@ Non-primary clusters reject all broadcasts with `ErrNotPrimary`. 5. **AckCallback**: CChannel ACK enqueues the task into `ackCallbackScheduler`. The callback executes only after all VChannels are ACKed. For tasks with conflicting ResourceKeys, callbacks execute in CChannel TimeTick order. Callbacks retry with exponential backoff until success. 6. **Tombstone & GC**: After callbacks complete and TOMBSTONE is persisted, release the callback resource locks before handing the task to `tombstoneScheduler`. Handoff appends to an in-memory queue and coalesces wakeups; catalog deletion runs outside the queue lock. Every queued ID is retained until drained. Recovery rebuilds this queue from durable TOMBSTONE tasks if shutdown interrupts handoff. GC applies the existing count and lifetime limits; a sustained deletion deficit can still grow the queue. +GC removes eligible tombstones in batches bounded by `metastore.maxEtcdTxnNum` (default 64), using exact broadcast-task keys. Each successful batch retires its in-memory tasks and advances the queue; failures retain the batch for idempotent retry, including when the deletion result is ambiguous. Late ACKs on TOMBSTONE or DONE tasks are ignored, so GC deletion holds neither task nor manager locks. Manager shutdown cancels an in-flight deletion; recovery only enqueues records still present in the catalog and does not replay their completed callbacks. + ## Resource Key Locking Each ResourceKey has: **Domain** (resource type), **Key** (entity identifier), **Shared** (read vs exclusive). Every broadcast automatically acquires SharedCluster. diff --git a/internal/metastore/kv/streamingcoord/kv_catalog.go b/internal/metastore/kv/streamingcoord/kv_catalog.go index 1c1e24bb5b3..21bce8b2709 100644 --- a/internal/metastore/kv/streamingcoord/kv_catalog.go +++ b/internal/metastore/kv/streamingcoord/kv_catalog.go @@ -203,6 +203,20 @@ func (c *catalog) SaveBroadcastTask(ctx context.Context, broadcastID uint64, tas return c.metaKV.Save(ctx, key, string(v)) } +func (c *catalog) RemoveBroadcastTasks(ctx context.Context, broadcastIDs []uint64) error { + keys := make([]string, 0, len(broadcastIDs)) + for _, id := range broadcastIDs { + keys = append(keys, buildBroadcastTaskPath(id)) + } + maxTxnNum := paramtable.Get().MetaStoreCfg.MaxEtcdTxnNum.GetAsInt() + return etcd.RemoveByBatchWithLimit(keys, maxTxnNum, func(batch []string) error { + if err := ctx.Err(); err != nil { + return err + } + return c.metaKV.MultiRemove(ctx, batch) + }) +} + // buildPChannelInfoPath builds the path for pchannel info. func buildPChannelInfoPath(name string) string { return PChannelMetaPrefix + name diff --git a/internal/metastore/kv/streamingcoord/kv_catalog_batch_test.go b/internal/metastore/kv/streamingcoord/kv_catalog_batch_test.go new file mode 100644 index 00000000000..c619ecf2ff1 --- /dev/null +++ b/internal/metastore/kv/streamingcoord/kv_catalog_batch_test.go @@ -0,0 +1,118 @@ +package streamingcoord + +import ( + "context" + "slices" + "testing" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus/pkg/v3/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/v3/util/merr" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" +) + +func setBroadcastDeletionBatchSize(t *testing.T, value string) { + t.Helper() + paramtable.Init() + limit := ¶mtable.Get().MetaStoreCfg.MaxEtcdTxnNum + previous := limit.SwapTempValue(value) + t.Cleanup(func() { limit.SwapTempValue(previous) }) +} + +func TestRemoveBroadcastTasksUsesBoundedExactKeys(t *testing.T) { + setBroadcastDeletionBatchSize(t, "2") + catalog, stored, kv := newTestCatalog(t) + ctx := context.Background() + ids := []uint64{1, 2, 3, 4, 5} + for _, id := range append(slices.Clone(ids), 11) { + require.NoError(t, catalog.SaveBroadcastTask(ctx, id, &streamingpb.BroadcastTask{ + State: streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, + })) + } + stored["querycoord-collection-loadinfo/1"] = "keep" + var batches [][]string + kv.EXPECT().MultiRemove(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, keys []string) error { + batches = append(batches, slices.Clone(keys)) + for _, key := range keys { + delete(stored, key) + } + return nil + }).Times(3) + require.NoError(t, catalog.RemoveBroadcastTasks(ctx, nil)) + require.NoError(t, catalog.RemoveBroadcastTasks(ctx, ids)) + require.Equal(t, [][]string{ + {buildBroadcastTaskPath(1), buildBroadcastTaskPath(2)}, + {buildBroadcastTaskPath(3), buildBroadcastTaskPath(4)}, + {buildBroadcastTaskPath(5)}, + }, batches) + require.Len(t, stored, 2) + require.Contains(t, stored, buildBroadcastTaskPath(11), "deletion must not match key prefixes") + require.Equal(t, "keep", stored["querycoord-collection-loadinfo/1"]) +} + +func TestRemoveBroadcastTasksRetriesLostCommitResponse(t *testing.T) { + setBroadcastDeletionBatchSize(t, "2") + catalog, stored, kv := newTestCatalog(t) + ids := []uint64{1, 2, 3} + for _, id := range ids { + stored[buildBroadcastTaskPath(id)] = "tombstone" + } + var attempts [][]string + kv.EXPECT().MultiRemove(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, keys []string) error { + attempts = append(attempts, slices.Clone(keys)) + for _, key := range keys { + delete(stored, key) + } + if len(attempts) == 1 { + return merr.WrapErrServiceUnavailable("commit response lost") + } + return nil + }).Times(3) + require.NoError(t, catalog.RemoveBroadcastTasks(context.Background(), ids)) + require.Equal(t, attempts[0], attempts[1], "the real reliable-write wrapper must retry the same deletion") + require.Equal(t, []string{buildBroadcastTaskPath(3)}, attempts[2]) + require.Empty(t, stored) +} + +func TestRemoveBroadcastTasksCancellationAfterPartialProgress(t *testing.T) { + setBroadcastDeletionBatchSize(t, "2") + catalog, stored, kv := newTestCatalog(t) + ids := []uint64{1, 2, 3, 4, 5} + for _, id := range ids { + require.NoError(t, catalog.SaveBroadcastTask(context.Background(), id, &streamingpb.BroadcastTask{ + State: streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, + })) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var attempts [][]string + kv.EXPECT().MultiRemove(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, keys []string) error { + attempts = append(attempts, slices.Clone(keys)) + if len(attempts) == 2 { + cancel() + return ctx.Err() + } + for _, key := range keys { + delete(stored, key) + } + return nil + }) + require.ErrorIs(t, catalog.RemoveBroadcastTasks(ctx, ids), context.Canceled) + require.Len(t, attempts, 2, "do not issue further batches after cancellation") + tasks, err := catalog.ListBroadcastTask(context.Background()) + require.NoError(t, err) + require.Len(t, tasks, 3) + for _, task := range tasks { + require.Equal(t, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, task.State) + } + require.NoError(t, catalog.RemoveBroadcastTasks(context.Background(), ids)) + require.Equal(t, attempts[0], attempts[2], "retry tolerates already removed IDs") + require.Empty(t, stored) + require.ErrorIs(t, catalog.RemoveBroadcastTasks(ctx, ids), context.Canceled) + require.Len(t, attempts, 5, "an already canceled request must not access the KV store") +} diff --git a/internal/metastore/streamingcoord_catalog.go b/internal/metastore/streamingcoord_catalog.go index 4afa185621a..bb8257980f2 100644 --- a/internal/metastore/streamingcoord_catalog.go +++ b/internal/metastore/streamingcoord_catalog.go @@ -43,6 +43,10 @@ type StreamingCoordCataLog interface { // Only return error if the ctx is canceled, otherwise it will retry until success. SaveBroadcastTask(ctx context.Context, broadcastID uint64, task *streamingpb.BroadcastTask) error + // RemoveBroadcastTasks removes completed tombstones with distinct IDs in bounded batches. + // An error may follow partially completed deletion; retrying the same IDs is safe. + RemoveBroadcastTasks(ctx context.Context, broadcastIDs []uint64) error + // SaveReplicateConfiguration saves the replicate configuration to metastore. // Only return error if the ctx is canceled, otherwise it will retry until success. SaveReplicateConfiguration(ctx context.Context, config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta) error diff --git a/internal/mocks/mock_metastore/mock_StreamingCoordCataLog.go b/internal/mocks/mock_metastore/mock_StreamingCoordCataLog.go index eb2366ba821..f4adf30fdf6 100644 --- a/internal/mocks/mock_metastore/mock_StreamingCoordCataLog.go +++ b/internal/mocks/mock_metastore/mock_StreamingCoordCataLog.go @@ -5,9 +5,8 @@ package mock_metastore import ( context "context" - mock "github.com/stretchr/testify/mock" - streamingpb "github.com/milvus-io/milvus/pkg/v3/proto/streamingpb" + mock "github.com/stretchr/testify/mock" ) // MockStreamingCoordCataLog is an autogenerated mock type for the StreamingCoordCataLog type @@ -313,6 +312,53 @@ func (_c *MockStreamingCoordCataLog_ListPChannel_Call) RunAndReturn(run func(con return _c } +// RemoveBroadcastTasks provides a mock function with given fields: ctx, broadcastIDs +func (_m *MockStreamingCoordCataLog) RemoveBroadcastTasks(ctx context.Context, broadcastIDs []uint64) error { + ret := _m.Called(ctx, broadcastIDs) + + if len(ret) == 0 { + panic("no return value specified for RemoveBroadcastTasks") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, []uint64) error); ok { + r0 = rf(ctx, broadcastIDs) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingCoordCataLog_RemoveBroadcastTasks_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveBroadcastTasks' +type MockStreamingCoordCataLog_RemoveBroadcastTasks_Call struct { + *mock.Call +} + +// RemoveBroadcastTasks is a helper method to define mock.On call +// - ctx context.Context +// - broadcastIDs []uint64 +func (_e *MockStreamingCoordCataLog_Expecter) RemoveBroadcastTasks(ctx interface{}, broadcastIDs interface{}) *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call { + return &MockStreamingCoordCataLog_RemoveBroadcastTasks_Call{Call: _e.mock.On("RemoveBroadcastTasks", ctx, broadcastIDs)} +} + +func (_c *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call) Run(run func(ctx context.Context, broadcastIDs []uint64)) *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].([]uint64)) + }) + return _c +} + +func (_c *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call) Return(_a0 error) *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call) RunAndReturn(run func(context.Context, []uint64) error) *MockStreamingCoordCataLog_RemoveBroadcastTasks_Call { + _c.Call.Return(run) + return _c +} + // SaveBroadcastTask provides a mock function with given fields: ctx, broadcastID, task func (_m *MockStreamingCoordCataLog) SaveBroadcastTask(ctx context.Context, broadcastID uint64, task *streamingpb.BroadcastTask) error { ret := _m.Called(ctx, broadcastID, task) diff --git a/internal/streamingcoord/server/broadcaster/ack_callback_gc_order_test.go b/internal/streamingcoord/server/broadcaster/ack_callback_gc_order_test.go new file mode 100644 index 00000000000..addcd88b989 --- /dev/null +++ b/internal/streamingcoord/server/broadcaster/ack_callback_gc_order_test.go @@ -0,0 +1,169 @@ +package broadcaster + +import ( + "context" + "slices" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus/internal/mocks/mock_metastore" + "github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/registry" + "github.com/milvus-io/milvus/internal/streamingcoord/server/resource" + "github.com/milvus-io/milvus/pkg/v3/mlog" + "github.com/milvus-io/milvus/pkg/v3/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/v3/streaming/util/message" +) + +// Exercises the real ACK scheduler and load/release message types with controlled +// business callbacks. It deliberately overlaps a retry, blocked persistence, +// queued same-key successors, and blocked GC catalog deletion. +func TestAckCallbacksSameCollectionRemainOrderedWithBatchGC(t *testing.T) { + configureTombstoneGCTest(t) + registry.ResetRegistration() + defer registry.ResetRegistration() + s := newAckCallbackScheduler(mlog.With()) + bm := newTombstoneGCTestManager() + gcStarted := make(chan struct{}) + gcGate := make(chan struct{}) + persistStarted := make(chan struct{}) + persistGate := make(chan struct{}) + releaseGC := sync.OnceFunc(func() { close(gcGate) }) + releasePersist := sync.OnceFunc(func() { close(persistGate) }) + defer releaseGC() + defer releasePersist() + + var eventsMu sync.Mutex + var appliedIDs []uint64 + var appliedStates []bool + var active atomic.Int32 + var firstAttempts atomic.Int32 + apply := func(id uint64, loaded bool) error { + if active.Add(1) != 1 { + t.Error("same-collection business callbacks overlapped") + } + defer active.Add(-1) + if id == 1 && firstAttempts.Add(1) == 1 { + return context.DeadlineExceeded + } + eventsMu.Lock() + appliedIDs = append(appliedIDs, id) + appliedStates = append(appliedStates, loaded) + eventsMu.Unlock() + return nil + } + registry.RegisterAlterLoadConfigV2AckCallback(func(_ context.Context, result message.BroadcastResultAlterLoadConfigMessageV2) error { + return apply(result.Message.BroadcastHeader().BroadcastID, true) + }) + registry.RegisterDropLoadConfigV2AckCallback(func(_ context.Context, result message.BroadcastResultDropLoadConfigMessageV2) error { + return apply(result.Message.BroadcastHeader().BroadcastID, false) + }) + + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + meta.EXPECT().SaveBroadcastTask(mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, id uint64, task *streamingpb.BroadcastTask) error { + if id == 1 && task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE { + close(persistStarted) + select { + case <-persistGate: + case <-ctx.Done(): + return ctx.Err() + } + } + return nil + }) + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + if len(ids) == 1 && ids[0] == 0 { + close(gcStarted) + select { + case <-gcGate: + case <-ctx.Done(): + return ctx.Err() + } + } + return nil + }) + resource.InitForTest(resource.OptStreamingCatalog(meta)) + + rk := message.NewExclusiveCollectionNameResourceKey("db", "same_collection") + tasks := make([]*broadcastTask, 0, 3) + for id := uint64(0); id <= 3; id++ { + var msg message.BroadcastMutableMessage + if id == 2 { + msg = message.NewDropLoadConfigMessageBuilderV2(). + WithHeader(&message.DropLoadConfigMessageHeader{CollectionId: 42}). + WithBody(&message.DropLoadConfigMessageBody{}). + WithBroadcast([]string{"by-dev-0_vcchan"}).MustBuildBroadcast() + } else { + msg = message.NewAlterLoadConfigMessageBuilderV2(). + WithHeader(&message.AlterLoadConfigMessageHeader{CollectionId: 42}). + WithBody(&message.AlterLoadConfigMessageBody{}). + WithBroadcast([]string{"by-dev-0_vcchan"}).MustBuildBroadcast() + } + msg = msg.OverwriteBroadcastHeader(id, rk) + state := streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_PENDING + if id == 0 { + state = streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE + } + p := createNewWaitAckBroadcastTaskFromMessage(msg, state, []byte{1}) + p.AckedCheckpoints[0].TimeTick = id + 100 + task := newBroadcastTaskFromProto(p, newBroadcasterMetrics(), s) + task.SetLogger(mlog.With()) + bm.tasks[id] = task + if id != 0 { + tasks = append(tasks, task) + } + } + s.Initialize([]*broadcastTask{tasks[2], tasks[0], tasks[1]}, []uint64{0}, bm) + defer func() { + releaseGC() + releasePersist() + s.Close() + }() + for _, gate := range []<-chan struct{}{gcStarted, persistStarted} { + select { + case <-gate: + case <-time.After(5 * time.Second): + t.Fatal("expected blocked phase was not reached") + } + } + + guards, ok := s.rkLocker.TryLock(rk) + if ok { + guards.Unlock() + t.Fatal("ACK lock released before durable completion") + } + eventsMu.Lock() + ids := slices.Clone(appliedIDs) + eventsMu.Unlock() + require.Equal(t, []uint64{1}, ids) + releasePersist() + + // GC is still blocked, but all three same-collection operations must finish + // business execution in WAL order and enqueue all of their GC work. + require.Eventually(t, func() bool { + s.tombstoneScheduler.pendingMu.Lock() + defer s.tombstoneScheduler.pendingMu.Unlock() + return len(s.tombstoneScheduler.pending) == 3 + }, 5*time.Second, time.Millisecond) + eventsMu.Lock() + ids = slices.Clone(appliedIDs) + states := slices.Clone(appliedStates) + eventsMu.Unlock() + require.Equal(t, []uint64{1, 2, 3}, ids) + require.Equal(t, []bool{true, false, true}, states) + require.EqualValues(t, 2, firstAttempts.Load()) + require.Zero(t, active.Load()) + + releaseGC() + require.Eventually(t, func() bool { + bm.mu.Lock() + defer bm.mu.Unlock() + return len(bm.tasks) == 0 + }, 5*time.Second, time.Millisecond) +} diff --git a/internal/streamingcoord/server/broadcaster/broadcast_manager.go b/internal/streamingcoord/server/broadcaster/broadcast_manager.go index b7253866a0c..e9102bfc4ae 100644 --- a/internal/streamingcoord/server/broadcaster/broadcast_manager.go +++ b/internal/streamingcoord/server/broadcaster/broadcast_manager.go @@ -272,22 +272,49 @@ func (bm *broadcastTaskManager) Ack(ctx context.Context, msg message.ImmutableMe return t.Ack(ctx, msg) } -// DropTombstone drops the tombstone task from the manager. -func (bm *broadcastTaskManager) DropTombstone(ctx context.Context, broadcastID uint64) error { +// DropTombstones removes durable tombstones before retiring their in-memory tasks. +// The GC scheduler enforces retention before calling this method. Once removed, +// these records no longer provide broadcast deduplication. +func (bm *broadcastTaskManager) DropTombstones(ctx context.Context, broadcastIDs []uint64) error { if !bm.lifetime.Add(typeutil.LifetimeStateWorking) { return status.NewOnShutdownError("broadcaster is closing") } defer bm.lifetime.Done() + ctx, cancel := bm.withLifecycleContext(ctx) + defer cancel() - t, ok := bm.getBroadcastTaskByID(broadcastID) - if !ok { - bm.Logger().Debug(ctx, "task is not found, ignored the drop tombstone request", mlog.Uint64("broadcastID", broadcastID)) + // Snapshot and deduplicate IDs without taking task locks under the manager lock. + tasks := make(map[uint64]*broadcastTask, len(broadcastIDs)) + bm.mu.Lock() + for _, id := range broadcastIDs { + if task, ok := bm.tasks[id]; ok { + tasks[id] = task + } + } + bm.mu.Unlock() + ids := make([]uint64, 0, len(tasks)) + for id, task := range tasks { + if state := task.State(); state != streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE { + return merr.WrapErrServiceInternalMsg("cannot drop broadcast task %d in state %s", id, state.String()) + } + ids = append(ids, id) + } + if len(ids) == 0 { return nil } - if err := t.DropTombstone(ctx); err != nil { + // Terminal tasks ignore late ACKs, so deletion needs neither task locks nor + // the manager lock. Keep all tasks on failure, including an ambiguous commit. + if err := resource.Resource().StreamingCatalog().RemoveBroadcastTasks(ctx, ids); err != nil { return err } - bm.removeBroadcastTask(broadcastID) + for _, task := range tasks { + task.markTombstoneDropped() + } + bm.mu.Lock() + for id := range tasks { + delete(bm.tasks, id) + } + bm.mu.Unlock() return nil } @@ -324,7 +351,9 @@ func (bm *broadcastTaskManager) getOrCreateBroadcastTask(msg message.ImmutableMe t, ok := bm.tasks[bh.BroadcastID] if ok { bm.mu.Unlock() - return t, t.State() != streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE + state := t.State() + return t, state != streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE && + state != streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE } if msg.ReplicateHeader() == nil { bm.mu.Unlock() @@ -348,14 +377,6 @@ func (bm *broadcastTaskManager) getBroadcastTaskByID(broadcastID uint64) (*broad return t, ok } -// removeBroadcastTask removes the broadcast task by the broadcastID. -func (bm *broadcastTaskManager) removeBroadcastTask(broadcastID uint64) { - bm.mu.Lock() - defer bm.mu.Unlock() - - delete(bm.tasks, broadcastID) -} - // getIncompleteBroadcastTasks returns all incomplete broadcast tasks that have pending messages. // Tasks in PENDING or REPLICATED state with pending messages are considered incomplete. func (bm *broadcastTaskManager) getIncompleteBroadcastTasks() []*broadcastTask { diff --git a/internal/streamingcoord/server/broadcaster/broadcast_task.go b/internal/streamingcoord/server/broadcaster/broadcast_task.go index 48938ef4c7e..4503cf2c514 100644 --- a/internal/streamingcoord/server/broadcaster/broadcast_task.go +++ b/internal/streamingcoord/server/broadcaster/broadcast_task.go @@ -244,7 +244,8 @@ func (b *broadcastTask) PendingSchemaFileResourceSnapshot() (int64, []int64, boo b.mu.Lock() defer b.mu.Unlock() - if b.task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE { + if b.task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE || + b.task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE { return 0, nil, false } switch b.messageTypeWithVersion { @@ -387,6 +388,12 @@ func (b *broadcastTask) Ack(ctx context.Context, msgs message.ImmutableMessage) // ack acknowledges the message at the specified vchannel. // Caller must resolve callback before acquiring b.mu and hold b.mu while calling. func (b *broadcastTask) ack(ctx context.Context, callback registry.ResolvedMessageAckOnceCallback, msgs ...message.ImmutableMessage) (err error) { + // A caller may have obtained this task before its callback completed. Never + // let a late ACK recreate a tombstone concurrently with catalog deletion. + if b.task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE || + b.task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE { + return nil + } isControlChannelAcked := b.copyAndSetAckedCheckpoints(msgs...) if !b.dirty { return nil @@ -546,16 +553,15 @@ func (b *broadcastTask) FastAck(ctx context.Context, broadcastResult map[string] return b.ack(ctx, callback, msgs...) } -// DropTombstone drops the tombstone of the broadcast task. -// It will remove the tombstone of the broadcast task in recovery storage. -// After the tombstone is dropped, the idempotency and deduplication can not be guaranteed. -func (b *broadcastTask) DropTombstone(ctx context.Context) error { +// markTombstoneDropped retires a task after its catalog record was deleted. +func (b *broadcastTask) markTombstoneDropped() { b.mu.Lock() defer b.mu.Unlock() - b.task.State = streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE - b.dirty = true - return b.saveTaskIfDirty(ctx, b.Logger()) + if b.task.State != streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE { + b.task.State = streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE + b.ObserveStateChanged(b.task.State) + } } // isAllDone check if all the vchannels are acked. diff --git a/internal/streamingcoord/server/broadcaster/broadcaster_test.go b/internal/streamingcoord/server/broadcaster/broadcaster_test.go index adb7d63d750..b4bfcebbbc4 100644 --- a/internal/streamingcoord/server/broadcaster/broadcaster_test.go +++ b/internal/streamingcoord/server/broadcaster/broadcaster_test.go @@ -85,6 +85,7 @@ func TestBroadcaster(t *testing.T) { }, nil }).Times(1) done := typeutil.NewConcurrentSet[uint64]() + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything).Return(nil).Maybe() meta.EXPECT().SaveBroadcastTask(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, broadcastID uint64, bt *streamingpb.BroadcastTask) error { if ctx.Err() != nil { return ctx.Err() @@ -264,6 +265,7 @@ func registerDropCollectionNoopCallbacks() { return nil }) } + func createNewBroadcastTask(broadcastID uint64, vchannels []string, rks ...message.ResourceKey) *streamingpb.BroadcastTask { msg := createNewBroadcastMsg(vchannels).OverwriteBroadcastHeader(broadcastID, rks...) pb := msg.IntoMessageProto() diff --git a/internal/streamingcoord/server/broadcaster/force_promote_failover_test.go b/internal/streamingcoord/server/broadcaster/force_promote_failover_test.go index 0e838c0ea33..6d3caebfaa0 100644 --- a/internal/streamingcoord/server/broadcaster/force_promote_failover_test.go +++ b/internal/streamingcoord/server/broadcaster/force_promote_failover_test.go @@ -243,6 +243,7 @@ func setupForcePromoteTest( walBehaviors []appendBehavior, ) *forcePromoteTestEnv { catalog := mock_metastore.NewMockStreamingCoordCataLog(t) + catalog.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything).Return(nil).Maybe() catalog.EXPECT().ListBroadcastTask(mock.Anything). Return(recoveryTasks, nil).Times(1) diff --git a/internal/streamingcoord/server/broadcaster/tombstone_batch_test.go b/internal/streamingcoord/server/broadcaster/tombstone_batch_test.go new file mode 100644 index 00000000000..a52f36c5457 --- /dev/null +++ b/internal/streamingcoord/server/broadcaster/tombstone_batch_test.go @@ -0,0 +1,293 @@ +package broadcaster + +import ( + "context" + "slices" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "github.com/milvus-io/milvus/internal/mocks/mock_metastore" + "github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/registry" + "github.com/milvus-io/milvus/internal/streamingcoord/server/resource" + "github.com/milvus-io/milvus/pkg/v3/mlog" + "github.com/milvus-io/milvus/pkg/v3/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/v3/streaming/util/types" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" + "github.com/milvus-io/milvus/pkg/v3/util/typeutil" +) + +func newBatchGCTestScheduler(t *testing.T, count int) *tombstoneScheduler { + t.Helper() + configureTombstoneGCTest(t) + limit := ¶mtable.Get().MetaStoreCfg.MaxEtcdTxnNum + previous := limit.SwapTempValue("4") + t.Cleanup(func() { limit.SwapTempValue(previous) }) + s := newTombstoneScheduler(mlog.With()) + t.Cleanup(s.notifier.Cancel) + s.bm = newTombstoneGCTestManager() + for id := uint64(1); id <= uint64(count); id++ { + msg := createNewBroadcastMsg([]string{"v1"}).WithBroadcastID(id) + task := newBroadcastTaskFromProto(createNewWaitAckBroadcastTaskFromMessage( + msg, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, []byte{1}), newBroadcasterMetrics(), nil) + task.SetLogger(mlog.With()) + s.bm.tasks[id] = task + s.tombstones = append(s.tombstones, tombstoneItem{broadcastID: id, createTime: time.Now()}) + } + return s +} + +func TestTombstoneGCBatchRetention(t *testing.T) { + for _, test := range []struct { + name string + maxCount string + aged int + batches []int + }{ + {"count", "2", 0, []int{4, 3}}, + {"lifetime", "20", 5, []int{4, 1}}, + {"both", "4", 7, []int{4, 3}}, + {"all", "0", 0, []int{4, 4, 1}}, + {"exact_batch", "5", 0, []int{4}}, + {"none", "9", 0, nil}, + } { + t.Run(test.name, func(t *testing.T) { + s := newBatchGCTestScheduler(t, 9) + paramtable.Get().StreamingCfg.WALBroadcasterTombstoneMaxCount.SwapTempValue(test.maxCount) + for i := range test.aged { + s.tombstones[i].createTime = time.Now().Add(-2 * time.Hour) + } + var batches []int + var removed []uint64 + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + batches = append(batches, len(ids)) + removed = append(removed, ids...) + return nil + }).Maybe() + resource.InitForTest(resource.OptStreamingCatalog(meta)) + s.triggerGCTombstone() + require.Equal(t, test.batches, batches) + slices.Sort(removed) + for i, id := range removed { + require.Equal(t, uint64(i+1), id) + require.NotContains(t, s.bm.tasks, id) + } + require.Len(t, s.tombstones, 9-len(removed)) + require.Len(t, s.bm.tasks, len(s.tombstones)) + for i, item := range s.tombstones { + require.Equal(t, uint64(len(removed)+i+1), item.broadcastID) + require.Equal(t, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, s.bm.tasks[item.broadcastID].State()) + } + }) + } +} + +func TestTombstoneGCBatchFailurePreservesProgress(t *testing.T) { + s := newBatchGCTestScheduler(t, 10) + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + var batches [][]uint64 + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + batch := slices.Clone(ids) + slices.Sort(batch) + batches = append(batches, batch) + if len(batches) == 2 { + return context.DeadlineExceeded + } + return nil + }) + resource.InitForTest(resource.OptStreamingCatalog(meta)) + s.triggerGCTombstone() + require.Equal(t, [][]uint64{{1, 2, 3, 4}, {5, 6, 7, 8}}, batches) + require.Len(t, s.tombstones, 6) + require.Len(t, s.bm.tasks, 6) + for id := uint64(5); id <= 10; id++ { + require.Equal(t, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, s.bm.tasks[id].State()) + } + s.triggerGCTombstone() + require.Equal(t, [][]uint64{{1, 2, 3, 4}, {5, 6, 7, 8}, {5, 6, 7, 8}, {9, 10}}, batches) + require.Empty(t, s.tombstones) + require.Empty(t, s.bm.tasks) +} + +func TestTombstoneGCBatchShutdownRecovery(t *testing.T) { + for _, committed := range []bool{false, true} { + name := "canceled_before_commit" + if committed { + name = "commit_response_lost" + } + t.Run(name, func(t *testing.T) { + s := newBatchGCTestScheduler(t, 10) + registry.ResetRegistration() + var mu sync.Mutex + persisted := make(map[uint64]*streamingpb.BroadcastTask) + for id, task := range s.bm.tasks { + persisted[id] = proto.Clone(task.task).(*streamingpb.BroadcastTask) + } + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + meta.EXPECT().ListBroadcastTask(mock.Anything). + RunAndReturn(func(ctx context.Context) ([]*streamingpb.BroadcastTask, error) { + mu.Lock() + defer mu.Unlock() + var tasks []*streamingpb.BroadcastTask + for _, task := range persisted { + tasks = append(tasks, proto.Clone(task).(*streamingpb.BroadcastTask)) + } + return tasks, nil + }) + blocked := make(chan struct{}) + attempts := 0 + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + attempts++ + if attempts != 2 || committed { + mu.Lock() + for _, id := range ids { + delete(persisted, id) + } + mu.Unlock() + } + if attempts == 2 { + close(blocked) + <-ctx.Done() + return ctx.Err() + } + return nil + }) + resource.InitForTest(resource.OptStreamingCatalog(meta)) + bc, err := RecoverBroadcaster(context.Background()) + require.NoError(t, err) + closed := sync.OnceFunc(bc.Close) + defer func() { + // Bound cleanup even if the shutdown-cancellation assertion fails. + bc.(*broadcastTaskManager).ackScheduler.tombstoneScheduler.notifier.Cancel() + closed() + }() + select { + case <-blocked: + case <-time.After(5 * time.Second): + t.Fatal("second batch did not start") + } + closeDone := make(chan struct{}) + go func() { closed(); close(closeDone) }() + select { + case <-closeDone: + case <-time.After(5 * time.Second): + t.Fatal("manager shutdown did not cancel in-flight GC") + } + old := bc.(*broadcastTaskManager) + require.Len(t, old.tasks, 6, "the canceled batch must remain in memory") + for _, task := range old.tasks { + require.Equal(t, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE, task.State()) + } + wantRemaining := 6 + if committed { + wantRemaining = 2 + } + require.Len(t, persisted, wantRemaining) + // No business callbacks are registered. Remaining TOMBSTONE records + // must recover directly into GC, including after an ambiguous commit. + recovered, err := RecoverBroadcaster(context.Background()) + require.NoError(t, err) + defer recovered.Close() + bm := recovered.(*broadcastTaskManager) + require.Eventually(t, func() bool { + bm.mu.Lock() + defer bm.mu.Unlock() + return len(bm.tasks) == 0 + }, 5*time.Second, time.Millisecond) + mu.Lock() + remaining := len(persisted) + mu.Unlock() + require.Zero(t, remaining) + }) + } +} + +func TestTombstoneGCBatchDoesNotBlockResultsOrReviveLateACKs(t *testing.T) { + s := newBatchGCTestScheduler(t, 2) + registry.ResetRegistration() + registerDropCollectionNoopCallbacks() + // Legacy tombstones may lack checkpoints. A late ACK must not fill them + // and save the record again while its deletion is in progress. + legacy := s.bm.tasks[2] + legacy.task.AckedCheckpoints[0] = nil + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + started := make(chan struct{}) + release := make(chan struct{}) + unblock := sync.OnceFunc(func() { close(release) }) + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + close(started) + select { + case <-release: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }).Once() + resource.InitForTest(resource.OptStreamingCatalog(meta)) + gcDone := make(chan struct{}) + go func() { s.triggerGCTombstone(); close(gcDone) }() + defer func() { unblock(); <-gcDone }() + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("GC did not start") + } + resultDone := make(chan error, 1) + go func() { + task, _ := s.bm.getBroadcastTaskByID(1) + _, err := task.BlockUntilDone(context.Background()) + resultDone <- err + }() + select { + case err := <-resultDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("RPC result waited for GC I/O") + } + _, result := s.bm.tasks[1].BroadcastResult() + lateACK := legacy.GetImmutableMessageFromVChannel("v1") + ackDone := make(chan error, 1) + go func() { + if err := legacy.FastAck(context.Background(), map[string]*types.AppendResult{"v1": result["v1"]}); err != nil { + ackDone <- err + return + } + ackDone <- legacy.Ack(context.Background(), lateACK) + }() + select { + case err := <-ackDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("late ACK waited for GC I/O") + } + unblock() + <-gcDone + require.NoError(t, legacy.Ack(context.Background(), lateACK)) + require.Equal(t, streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_DONE, legacy.State()) + require.Nil(t, legacy.task.AckedCheckpoints[0]) + require.Empty(t, s.bm.tasks) +} + +func TestDropTombstonesValidatesAndDeduplicatesIDs(t *testing.T) { + s := newBatchGCTestScheduler(t, 2) + meta := mock_metastore.NewMockStreamingCoordCataLog(t) + resource.InitForTest(resource.OptStreamingCatalog(meta)) + s.bm.tasks[2].task.State = streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_PENDING + require.Error(t, s.bm.DropTombstones(context.Background(), []uint64{1, 2})) + require.Len(t, s.bm.tasks, 2) + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, []uint64{1}).Return(nil).Once() + require.NoError(t, s.bm.DropTombstones(context.Background(), []uint64{1, 1, 99})) + require.Len(t, s.bm.tasks, 1) + require.NoError(t, s.bm.DropTombstones(context.Background(), nil)) + s.bm.lifetime.SetState(typeutil.LifetimeStateStopped) + require.Error(t, s.bm.DropTombstones(context.Background(), []uint64{2})) +} diff --git a/internal/streamingcoord/server/broadcaster/tombstone_scheduler.go b/internal/streamingcoord/server/broadcaster/tombstone_scheduler.go index 127bee7396c..1f4f2926b11 100644 --- a/internal/streamingcoord/server/broadcaster/tombstone_scheduler.go +++ b/internal/streamingcoord/server/broadcaster/tombstone_scheduler.go @@ -111,31 +111,44 @@ func (s *tombstoneScheduler) background() { // triggerGCTombstone triggers the garbage collection of the tombstone. func (s *tombstoneScheduler) triggerGCTombstone() { + ctx := s.notifier.Context() maxTombstoneLifetime := paramtable.Get().StreamingCfg.WALBroadcasterTombstoneMaxLifetime.GetAsDurationByParse() maxTombstoneCount := paramtable.Get().StreamingCfg.WALBroadcasterTombstoneMaxCount.GetAsInt() + batchSize := paramtable.Get().MetaStoreCfg.MaxEtcdTxnNum.GetAsInt() expiredTime := time.Now().Add(-maxTombstoneLifetime) expiredOffset := 0 if len(s.tombstones) > maxTombstoneCount { expiredOffset = len(s.tombstones) - maxTombstoneCount } - s.Logger().Info(context.TODO(), + s.Logger().Info(ctx, "triggerGCTombstone", mlog.Int("tombstone count", len(s.tombstones)), mlog.Int("expired offset", expiredOffset), mlog.Time("expired time", expiredTime)) - for idx, tombstone := range s.tombstones { - // drop tombstone until the expired time or until the expired offset. - if idx >= expiredOffset && tombstone.createTime.After(expiredTime) { - s.tombstones = s.tombstones[idx:] + ids := make([]uint64, 0, min(batchSize, len(s.tombstones))) + for len(s.tombstones) > 0 && ctx.Err() == nil { + ids = ids[:0] + for idx, tombstone := range s.tombstones[:min(batchSize, len(s.tombstones))] { + if idx >= expiredOffset && tombstone.createTime.After(expiredTime) { + break + } + ids = append(ids, tombstone.broadcastID) + } + if len(ids) == 0 { return } - if err := s.bm.DropTombstone(s.notifier.Context(), tombstone.broadcastID); err != nil { - s.Logger().Error(context.TODO(), "failed to drop tombstone", mlog.Err(err)) - s.tombstones = s.tombstones[idx:] + if err := s.bm.DropTombstones(ctx, ids); err != nil { + s.Logger().Warn(ctx, "failed to drop tombstone batch", mlog.Int("batchSize", len(ids)), mlog.Err(err)) return } + // Advance only after the whole batch succeeds. A failed batch remains + // queued for idempotent retry, while earlier successful batches stay gone. + clear(s.tombstones[:len(ids)]) + s.tombstones = s.tombstones[len(ids):] + expiredOffset -= len(ids) + } + if len(s.tombstones) == 0 { + s.tombstones = nil } - // all the tombstones are dropped, reset the tombstones. - s.tombstones = make([]tombstoneItem, 0) } diff --git a/internal/streamingcoord/server/broadcaster/tombstone_scheduler_test.go b/internal/streamingcoord/server/broadcaster/tombstone_scheduler_test.go index b5df7bfb295..622b16faaa6 100644 --- a/internal/streamingcoord/server/broadcaster/tombstone_scheduler_test.go +++ b/internal/streamingcoord/server/broadcaster/tombstone_scheduler_test.go @@ -30,9 +30,10 @@ func TestAckCallbacksCompleteWhileTombstoneGCBlocked(t *testing.T) { gcContext, releaseGC := context.WithCancel(context.Background()) defer releaseGC() meta := mock_metastore.NewMockStreamingCoordCataLog(t) - meta.EXPECT().SaveBroadcastTask(mock.Anything, mock.Anything, mock.Anything). - RunAndReturn(func(ctx context.Context, id uint64, task *streamingpb.BroadcastTask) error { - if id == 0 { + meta.EXPECT().SaveBroadcastTask(mock.Anything, mock.Anything, mock.Anything).Return(nil) + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + if len(ids) == 1 && ids[0] == 0 { close(gcStarted) select { case <-gcContext.Done(): @@ -116,13 +117,14 @@ func TestTombstoneRecoveryAfterShutdownBeforeHandoff(t *testing.T) { meta := mock_metastore.NewMockStreamingCoordCataLog(t) meta.EXPECT().SaveBroadcastTask(mock.Anything, uint64(1), mock.Anything). RunAndReturn(func(ctx context.Context, id uint64, task *streamingpb.BroadcastTask) error { - if task.State == streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_TOMBSTONE { - persisted = proto.Clone(task).(*streamingpb.BroadcastTask) - } else { - close(deleted) - } + persisted = proto.Clone(task).(*streamingpb.BroadcastTask) + return nil + }).Once() + meta.EXPECT().RemoveBroadcastTasks(mock.Anything, []uint64{1}). + RunAndReturn(func(ctx context.Context, ids []uint64) error { + close(deleted) return nil - }).Times(2) + }).Once() resource.InitForTest(resource.OptStreamingCatalog(meta)) task := newBroadcastTaskFromProto(createNewWaitAckBroadcastTaskFromMessage( createNewBroadcastMsg([]string{"v1"}).WithBroadcastID(1),