downstreamadapter: preserve remove upgrade during close (#4815)

close #4825

Author: hongyunyan

cmd/storage-consumer/main.go

@@ -101,7 +101,7 @@ func main() {
 	deferFunc := func() int {
 		stop()
 		if consumer != nil {
-			consumer.sink.Close(false)
+			consumer.sink.Close()
 		}
 		if err != nil && err != context.Canceled {
 			return 1

downstreamadapter/dispatcher/mock_sink_helper_test.go

@@ -72,7 +72,7 @@ func newDispatcherTestSink(t *testing.T, sinkType common.SinkType) *dispatcherTe
 	}).AnyTimes()
 	testSink.sink.EXPECT().AddCheckpointTs(gomock.Any()).AnyTimes()
 	testSink.sink.EXPECT().SetTableSchemaStore(gomock.Any()).AnyTimes()
-	testSink.sink.EXPECT().Close(gomock.Any()).AnyTimes()
+	testSink.sink.EXPECT().Close().AnyTimes()
 	testSink.sink.EXPECT().Run(gomock.Any()).Return(nil).AnyTimes()
 	return testSink
 }

downstreamadapter/dispatchermanager/dispatcher_manager.go

@@ -133,8 +133,19 @@ type DispatcherManager struct {
 
 	closing atomic.Bool
 	closed  atomic.Bool
-	cancel  context.CancelFunc
-	wg      sync.WaitGroup
+	// removeChangefeedRequested is sticky once any close request asks for removed=true.
+	// A later removed=false request must not downgrade the final cleanup semantics.
+	removeChangefeedRequested atomic.Bool
+	// removeChangefeedCleaned records whether the best-effort remove-only cleanup has finished.
+	// It is intentionally not part of the TryClose success condition so the close contract
+	// stays compatible with the historical behavior.
+	removeChangefeedCleaned atomic.Bool
+	// removeChangefeedCleanupRunning prevents duplicate background cleanup runs
+	// while late remove requests or retries keep asking for remove semantics after
+	// the base close path ends.
+	removeChangefeedCleanupRunning atomic.Bool
+	cancel                         context.CancelFunc
+	wg                             sync.WaitGroup
 
 	// removeTaskHandles stores the task handles for async dispatcher removal
 	// map[common.DispatcherID]*threadpool.TaskHandle
@@ -855,19 +866,23 @@ func (e *DispatcherManager) mergeEventDispatcher(dispatcherIDs []common.Dispatch
 // ==== remove and clean related functions ====
 
 func (e *DispatcherManager) TryClose(removeChangefeed bool) bool {
+	if removeChangefeed {
+		e.removeChangefeedRequested.Store(true)
+	}
 	if e.closed.Load() {
+		e.tryScheduleRemoveChangefeedCleanup()
 		return true
 	}
 	if e.closing.Load() {
 		return e.closed.Load()
 	}
 
 	e.closing.Store(true)
-	go e.close(removeChangefeed)
+	go e.close()
 	return false
 }
 
-func (e *DispatcherManager) close(removeChangefeed bool) {
+func (e *DispatcherManager) close() {
 	log.Info("closing event dispatcher manager",
 		zap.Stringer("changefeedID", e.changefeedID))
 
@@ -919,11 +934,9 @@ func (e *DispatcherManager) close(removeChangefeed bool) {
 	log.Info("shared info closed", zap.Stringer("changefeedID", e.changefeedID))
 
 	if e.IsRedoEnabled() {
-		e.redoSink.Close(removeChangefeed)
-		// FIXME: cleanup redo log when remove the changefeed
-		e.closeRedoMeta(removeChangefeed)
+		e.redoSink.Close()
 	}
-	e.sink.Close(removeChangefeed)
+	e.sink.Close()
 	log.Info("sink closed", zap.Stringer("changefeedID", e.changefeedID))
 
 	e.wg.Wait()
@@ -937,10 +950,58 @@ func (e *DispatcherManager) close(removeChangefeed bool) {
 	e.cleanMetrics()
 
 	e.closed.Store(true)
+	e.tryScheduleRemoveChangefeedCleanup()
 	log.Info("event dispatcher manager closed",
 		zap.Stringer("changefeedID", e.changefeedID))
 }
 
+func (e *DispatcherManager) tryScheduleRemoveChangefeedCleanup() {
+	if !e.removeChangefeedRequested.Load() {
+		return
+	}
+	if e.removeChangefeedCleaned.Load() {
+		return
+	}
+	if !e.removeChangefeedCleanupRunning.CompareAndSwap(false, true) {
+		return
+	}
+
+	go func() {
+		defer e.removeChangefeedCleanupRunning.Store(false)
+
+		if err := e.runRemoveChangefeedCleanup(); err != nil {
+			log.Warn("failed to cleanup removed changefeed",
+				zap.Stringer("changefeedID", e.changefeedID),
+				zap.Error(err))
+			return
+		}
+		e.removeChangefeedCleaned.Store(true)
+	}()
+}
+
+func (e *DispatcherManager) runRemoveChangefeedCleanup() error {
+	if !e.removeChangefeedRequested.Load() || e.removeChangefeedCleaned.Load() {
+		return nil
+	}
+
+	if e.IsRedoEnabled() {
+		// Redo meta cleanup is the remove-only step for redo mode. It is safe to retry
+		// because removeChangefeedCleaned is only set after all remove-only cleanup succeeds.
+		if err := e.closeRedoMeta(true); err != nil {
+			return errors.Trace(err)
+		}
+	}
+
+	if mysqlSink, ok := e.sink.(*mysql.Sink); ok {
+		// MySQL sink may still need ddl_ts cleanup after the base close path has already
+		// released the long-lived DB connection, so keep the retryable remove step here.
+		if err := mysqlSink.CleanupRemovedChangefeed(); err != nil {
+			return errors.Trace(err)
+		}
+	}
+	return nil
+}
+
 // cleanEventDispatcher is called when the event dispatcher is removed successfully.
 func (e *DispatcherManager) cleanEventDispatcher(id common.DispatcherID, schemaID int64) {
 	e.dispatcherMap.Delete(id)

downstreamadapter/dispatchermanager/dispatcher_manager_redo.go

@@ -257,16 +257,19 @@ func (e *DispatcherManager) cleanRedoDispatcher(id common.DispatcherID, schemaID
 	)
 }
 
-func (e *DispatcherManager) closeRedoMeta(removeChangefeed bool) {
+func (e *DispatcherManager) closeRedoMeta(removeChangefeed bool) error {
 	if d := e.GetTableTriggerRedoDispatcher(); d != nil {
 		redoMeta := d.GetRedoMeta()
 		if redoMeta != nil {
 			redoMeta.CleanupMetrics()
 			if removeChangefeed {
-				redoMeta.Cleanup(context.Background())
+				if err := redoMeta.Cleanup(context.Background()); err != nil {
+					return errors.Trace(err)
+				}
 			}
 		}
 	}
+	return nil
 }
 
 func (e *DispatcherManager) InitalizeTableTriggerRedoDispatcher(schemaInfo []*heartbeatpb.SchemaInfo) error {

downstreamadapter/dispatchermanager/dispatcher_manager_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/ticdc/downstreamadapter/eventcollector"
 	"github.com/pingcap/ticdc/downstreamadapter/sink"
 	"github.com/pingcap/ticdc/downstreamadapter/sink/mock"
+	"github.com/pingcap/ticdc/downstreamadapter/sink/mysql"
 	"github.com/pingcap/ticdc/heartbeatpb"
 	"github.com/pingcap/ticdc/pkg/common"
 	appcontext "github.com/pingcap/ticdc/pkg/common/context"
@@ -33,6 +34,7 @@ import (
 	"github.com/pingcap/ticdc/pkg/metrics"
 	"github.com/pingcap/ticdc/pkg/node"
 	"github.com/pingcap/ticdc/pkg/pdutil"
+	mysqlcfg "github.com/pingcap/ticdc/pkg/sink/mysql"
 	"github.com/pingcap/ticdc/pkg/util"
 	"github.com/pingcap/ticdc/utils/threadpool"
 	"github.com/stretchr/testify/require"
@@ -53,7 +55,7 @@ func newDispatcherManagerTestSink(t *testing.T, sinkType common.SinkType) sink.S
 	}).AnyTimes()
 	mockSink.EXPECT().AddCheckpointTs(gomock.Any()).AnyTimes()
 	mockSink.EXPECT().SetTableSchemaStore(gomock.Any()).AnyTimes()
-	mockSink.EXPECT().Close(gomock.Any()).AnyTimes()
+	mockSink.EXPECT().Close().AnyTimes()
 	mockSink.EXPECT().Run(gomock.Any()).Return(nil).AnyTimes()
 	return mockSink
 }
@@ -274,6 +276,36 @@ func TestMergeDispatcherInvalidIDs(t *testing.T) {
 	require.False(t, exists)
 }
 
+func TestTryCloseRemovedRequestAfterClosedReturnsImmediatelyAndTriggersCleanup(t *testing.T) {
+	changefeedID := common.NewChangeFeedIDWithName("test", common.DefaultKeyspaceName)
+	mysqlConfig := mysqlcfg.New()
+	mysqlConfig.EnableDDLTs = false
+	mysqlSink := mysql.NewMySQLSink(
+		context.Background(),
+		changefeedID,
+		mysqlConfig,
+		nil,
+		false,
+		false,
+		time.Minute,
+	)
+	manager := &DispatcherManager{
+		changefeedID: changefeedID,
+		sink:         mysqlSink,
+	}
+	manager.closed.Store(true)
+
+	// Preserve the historical close contract: once the manager is already closed,
+	// late remove requests should not delay TryClose success.
+	closed := manager.TryClose(true)
+	require.True(t, closed)
+	require.True(t, manager.removeChangefeedRequested.Load())
+	require.Eventually(t, func() bool {
+		return manager.removeChangefeedCleaned.Load()
+	}, time.Second, 10*time.Millisecond)
+	require.True(t, manager.TryClose(true))
+}
+
 func TestMergeDispatcherExistingID(t *testing.T) {
 	manager := createTestManager(t)
 

downstreamadapter/dispatcherorchestrator/dispatcher_orchestrator.go

@@ -119,12 +119,11 @@ func getPendingMessageKey(msg *messaging.TargetMessage) (pendingMessageKey, bool
 // handleMessages processes messages from the queue
 func (m *DispatcherOrchestrator) handleMessages() {
 	for {
-		key, ok := m.msgQueue.Pop()
+		msg, ok := m.msgQueue.Pop()
 		if !ok {
 			log.Info("dispatcher orchestrator is shutting down, exit handleMessages")
 			return
 		}
-		msg := m.msgQueue.Get(key)
 
 		// Process the message
 		switch req := msg.Message[0].(type) {
@@ -146,8 +145,6 @@ func (m *DispatcherOrchestrator) handleMessages() {
 				zap.String("type", msg.Type.String()),
 				zap.Any("message", msg.Message))
 		}
-
-		m.msgQueue.Done(key)
 	}
 }
 

downstreamadapter/dispatcherorchestrator/dispatcher_orchestrator_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestPendingMessageQueue_TryEnqueueDropsDuplicatesUntilDone(t *testing.T) {
+func TestPendingMessageQueue_TryEnqueueDropsDuplicatesOnlyWhileQueued(t *testing.T) {
 	t.Parallel()
 
 	q := newPendingMessageQueue()
@@ -38,14 +38,18 @@ func TestPendingMessageQueue_TryEnqueueDropsDuplicatesUntilDone(t *testing.T) {
 	require.True(t, q.TryEnqueue(key, msg))
 	require.False(t, q.TryEnqueue(key, msg))
 
-	poppedKey, ok := q.Pop()
+	poppedMsg, ok := q.Pop()
 	require.True(t, ok)
-	require.Equal(t, key, poppedKey)
+	require.Same(t, msg, poppedMsg)
 
-	// The key remains pending while being processed.
+	// Once the request is popped, allow one queued retry for the next round.
+	require.True(t, q.TryEnqueue(key, msg))
 	require.False(t, q.TryEnqueue(key, msg))
 
-	q.Done(key)
+	nextMsg, ok := q.Pop()
+	require.True(t, ok)
+	require.Same(t, msg, nextMsg)
+
 	require.True(t, q.TryEnqueue(key, msg))
 }
 
@@ -62,15 +66,13 @@ func TestPendingMessageQueue_OrderPreservedAcrossKeys(t *testing.T) {
 	require.True(t, q.TryEnqueue(key1, &messaging.TargetMessage{Type: key1.msgType}))
 	require.True(t, q.TryEnqueue(key2, &messaging.TargetMessage{Type: key2.msgType}))
 
-	poppedKey, ok := q.Pop()
+	poppedMsg, ok := q.Pop()
 	require.True(t, ok)
-	require.Equal(t, key1, poppedKey)
-	q.Done(poppedKey)
+	require.Equal(t, key1.msgType, poppedMsg.Type)
 
-	poppedKey, ok = q.Pop()
+	poppedMsg, ok = q.Pop()
 	require.True(t, ok)
-	require.Equal(t, key2, poppedKey)
-	q.Done(poppedKey)
+	require.Equal(t, key2.msgType, poppedMsg.Type)
 }
 
 func TestPendingMessageQueue_PopReturnsAfterClose(t *testing.T) {
@@ -118,17 +120,14 @@ func TestPendingMessageQueue_CloseRequestRemovedTrueOverridesPendingFalse(t *tes
 	require.True(t, q.TryEnqueue(key, msgFalse))
 	require.True(t, q.TryEnqueue(key, msgTrue))
 
-	poppedKey, ok := q.Pop()
+	poppedMsg, ok := q.Pop()
 	require.True(t, ok)
-	require.Equal(t, key, poppedKey)
-	poppedMsg := q.Get(poppedKey)
 	require.NotNil(t, poppedMsg)
 	req := poppedMsg.Message[0].(*heartbeatpb.MaintainerCloseRequest)
 	require.True(t, req.Removed)
-	q.Done(key)
 }
 
-func TestPendingMessageQueue_CloseRequestUpgradeBetweenPopAndGet(t *testing.T) {
+func TestPendingMessageQueue_CloseRequestUpgradeAfterPopKeepsReturnedMessageStable(t *testing.T) {
 	t.Parallel()
 
 	q := newPendingMessageQueue()
@@ -150,16 +149,67 @@ func TestPendingMessageQueue_CloseRequestUpgradeBetweenPopAndGet(t *testing.T) {
 	)
 
 	require.True(t, q.TryEnqueue(key, msgFalse))
-	poppedKey, ok := q.Pop()
+	poppedMsg, ok := q.Pop()
 	require.True(t, ok)
-	require.Equal(t, key, poppedKey)
+	require.NotNil(t, poppedMsg)
 
 	require.True(t, q.TryEnqueue(key, msgTrue))
-	poppedMsg := q.Get(poppedKey)
-	require.NotNil(t, poppedMsg)
 	req2 := poppedMsg.Message[0].(*heartbeatpb.MaintainerCloseRequest)
-	require.True(t, req2.Removed)
-	q.Done(key)
+	require.False(t, req2.Removed)
+}
+
+func TestPendingMessageQueue_CloseRequestUpgradeAfterPopRequeuesNextRound(t *testing.T) {
+	t.Parallel()
+
+	q := newPendingMessageQueue()
+	cfID := common.NewChangeFeedIDWithName("cf", "default")
+	key := pendingMessageKey{
+		changefeedID: cfID,
+		msgType:      messaging.TypeMaintainerCloseRequest,
+	}
+
+	msgFalse := messaging.NewSingleTargetMessage(
+		node.ID("to"),
+		messaging.DispatcherManagerManagerTopic,
+		&heartbeatpb.MaintainerCloseRequest{ChangefeedID: cfID.ToPB(), Removed: false},
+	)
+	msgTrue := messaging.NewSingleTargetMessage(
+		node.ID("to"),
+		messaging.DispatcherManagerManagerTopic,
+		&heartbeatpb.MaintainerCloseRequest{ChangefeedID: cfID.ToPB(), Removed: true},
+	)
+
+	require.True(t, q.TryEnqueue(key, msgFalse))
+
+	poppedMsg, ok := q.Pop()
+	require.True(t, ok)
+
+	require.NotNil(t, poppedMsg)
+	req := poppedMsg.Message[0].(*heartbeatpb.MaintainerCloseRequest)
+	require.False(t, req.Removed)
+
+	require.True(t, q.TryEnqueue(key, msgTrue))
+
+	type popResult struct {
+		msg *messaging.TargetMessage
+		ok  bool
+	}
+	resultCh := make(chan popResult, 1)
+	go func() {
+		nextMsg, nextOK := q.Pop()
+		resultCh <- popResult{msg: nextMsg, ok: nextOK}
+	}()
+
+	select {
+	case result := <-resultCh:
+		require.True(t, result.ok)
+		require.NotNil(t, result.msg)
+		nextReq := result.msg.Message[0].(*heartbeatpb.MaintainerCloseRequest)
+		require.True(t, nextReq.Removed)
+	case <-time.After(time.Second):
+		q.Close()
+		require.FailNow(t, "upgraded close request was not requeued after the first pop")
+	}
 }
 
 func TestGetPendingMessageKey_SupportedTypes(t *testing.T) {