use heap in dispatcher

downstreamadapter/dispatcher/dispatcher.go

@@ -102,11 +102,9 @@ type Dispatcher struct {
 	// if syncPointInfo is not nil, means enable Sync Point feature,
 	syncPointConfig *syncpoint.SyncPointConfig
 
-	// the max resolvedTs received by the dispatcher
-	resolvedTs *TsWithMutex
-
 	// blockEventStatus is used to store the current pending ddl/sync point event and its block status.
 	blockEventStatus BlockEventStatus
+	ResolvedTs       *TsItem // The largest commitTs - 1 of the events received by the dispatcher.
 
 	// tableProgress is used to calculate the checkpointTs of the dispatcher
 	tableProgress *types.TableProgress
@@ -151,7 +149,7 @@ func NewDispatcher(
 		blockStatusesChan:     blockStatusesChan,
 		syncPointConfig:       syncPointConfig,
 		componentStatus:       newComponentStateWithMutex(heartbeatpb.ComponentState_Working),
-		resolvedTs:            newTsWithMutex(startTs),
+		ResolvedTs:            NewTsItem(startTs),
 		filterConfig:          filterConfig,
 		isRemoving:            atomic.Bool{},
 		blockEventStatus:      BlockEventStatus{blockPendingEvent: nil},
@@ -245,7 +243,7 @@ func (d *Dispatcher) HandleEvents(dispatcherEvents []DispatcherEvent, wakeCallba
 	for _, dispatcherEvent := range dispatcherEvents {
 		event := dispatcherEvent.Event
 		// Pre-check, make sure the event is not stale
-		if event.GetCommitTs() < d.resolvedTs.Get() {
+		if event.GetCommitTs() < d.ResolvedTs.Get() {
 			log.Info("Received a stale event, should ignore it",
 				zap.Any("commitTs", event.GetCommitTs()),
 				zap.Any("seq", event.GetSeq()),
@@ -257,7 +255,7 @@ func (d *Dispatcher) HandleEvents(dispatcherEvents []DispatcherEvent, wakeCallba
 
 		switch event.GetType() {
 		case commonEvent.TypeResolvedEvent:
-			d.resolvedTs.Set(event.(commonEvent.ResolvedEvent).ResolvedTs)
+			d.ResolvedTs.Set(event.(commonEvent.ResolvedEvent).ResolvedTs)
 		case commonEvent.TypeDMLEvent:
 			block = true
 			dml := event.(*commonEvent.DMLEvent)
@@ -449,7 +447,7 @@ func (d *Dispatcher) GetStartTs() uint64 {
 }
 
 func (d *Dispatcher) GetResolvedTs() uint64 {
-	return d.resolvedTs.Get()
+	return d.ResolvedTs.Get()
 }
 
 func (d *Dispatcher) GetCheckpointTs() uint64 {

downstreamadapter/dispatcher/helper.go

@@ -178,29 +178,6 @@ func (s *ComponentStateWithMutex) Get() heartbeatpb.ComponentState {
 	return s.componentStatus
 }
 
-type TsWithMutex struct {
-	mutex sync.Mutex
-	ts    uint64
-}
-
-func newTsWithMutex(ts uint64) *TsWithMutex {
-	return &TsWithMutex{
-		ts: ts,
-	}
-}
-
-func (r *TsWithMutex) Set(ts uint64) {
-	r.mutex.Lock()
-	defer r.mutex.Unlock()
-	r.ts = ts
-}
-
-func (r *TsWithMutex) Get() uint64 {
-	r.mutex.Lock()
-	defer r.mutex.Unlock()
-	return r.ts
-}
-
 /*
 HeartBeatInfo is used to collect the message for HeartBeatRequest for each dispatcher.
 Mainly about the progress of each dispatcher:
@@ -345,3 +322,49 @@ func GetDispatcherStatusDynamicStream() dynstream.DynamicStream[common.GID, comm
 func SetDispatcherStatusDynamicStream(dynamicStream dynstream.DynamicStream[common.GID, common.DispatcherID, DispatcherStatusWithID, *Dispatcher, *DispatcherStatusHandler]) {
 	dispatcherStatusDynamicStream = dynamicStream
 }
+
+type TsItem struct {
+	sync.RWMutex
+	ts      uint64
+	heapIdx int
+
+	onUpdate func(*TsItem)
+}
+
+func NewTsItem(ts uint64) *TsItem {
+	return &TsItem{
+		ts: ts,
+	}
+}
+
+func (d *TsItem) SetHeapIndex(idx int) {
+	d.heapIdx = idx
+}
+
+func (d *TsItem) GetHeapIndex() int {
+	return d.heapIdx
+}
+
+func (d *TsItem) LessThan(other *TsItem) bool {
+	return d.Get() < other.Get()
+}
+
+func (r *TsItem) Set(ts uint64) {
+	r.Lock()
+	r.ts = ts
+	r.Unlock()
+
+	if r.onUpdate != nil {
+		r.onUpdate(r)
+	}
+}
+
+func (r *TsItem) Get() uint64 {
+	r.RLock()
+	defer r.RUnlock()
+	return r.ts
+}
+
+func (r *TsItem) SetOnUpdate(f func(*TsItem)) {
+	r.onUpdate = f
+}

downstreamadapter/dispatchermanager/event_dispatcher_manager.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/pingcap/ticdc/pkg/apperror"
 	"github.com/pingcap/ticdc/pkg/node"
+	"github.com/pingcap/ticdc/utils/heap"
 
 	"github.com/pingcap/log"
 	"github.com/pingcap/ticdc/pkg/filter"
@@ -345,9 +346,9 @@ func (e *EventDispatcherManager) newDispatchers(infos []dispatcherCreateInfo) er
 			e.schemaIDToDispatchers.Set(schemaIds[idx], id)
 		}
 
+		e.dispatcherMap.Set(id, d)
 		appcontext.GetService[*eventcollector.EventCollector](appcontext.EventCollector).AddDispatcher(d, int(e.config.MemoryQuota))
 
-		e.dispatcherMap.Set(id, d)
 		e.statusesChan <- &heartbeatpb.TableSpanStatus{
 			ID:              id.ToPB(),
 			ComponentStatus: heartbeatpb.ComponentState_Working,
@@ -505,37 +506,42 @@ func (e *EventDispatcherManager) aggregateDispatcherHeartbeats(needCompleteStatu
 	removedDispatcherSchemaIDs := make([]int64, 0)
 	heartBeatInfo := &dispatcher.HeartBeatInfo{}
 
-	e.dispatcherMap.ForEach(func(id common.DispatcherID, dispatcherItem *dispatcher.Dispatcher) {
-		dispatcherItem.GetHeartBeatInfo(heartBeatInfo)
-		// If the dispatcher is in removing state, we need to check if it's closed successfully.
-		// If it's closed successfully, we could clean it up.
-		// TODO: we need to consider how to deal with the checkpointTs of the removed dispatcher if the message will be discarded.
-		if heartBeatInfo.IsRemoving {
-			watermark, ok := dispatcherItem.TryClose()
-			if ok {
-				// remove successfully
-				message.Watermark.UpdateMin(watermark)
-				// If the dispatcher is removed successfully, we need to add the tableSpan into message whether needCompleteStatus is true or not.
+	if needCompleteStatus {
+		e.dispatcherMap.ForEach(func(id common.DispatcherID, dispatcherItem *dispatcher.Dispatcher) {
+			dispatcherItem.GetHeartBeatInfo(heartBeatInfo)
+			// If the dispatcher is in removing state, we need to check if it's closed successfully.
+			// If it's closed successfully, we could clean it up.
+			// TODO: we need to consider how to deal with the checkpointTs of the removed dispatcher if the message will be discarded.
+			if heartBeatInfo.IsRemoving {
+				watermark, ok := dispatcherItem.TryClose()
+				if ok {
+					// remove successfully
+					message.Watermark.UpdateMin(watermark)
+					// If the dispatcher is removed successfully, we need to add the tableSpan into message whether needCompleteStatus is true or not.
+					message.Statuses = append(message.Statuses, &heartbeatpb.TableSpanStatus{
+						ID:              id.ToPB(),
+						ComponentStatus: heartbeatpb.ComponentState_Stopped,
+						CheckpointTs:    watermark.CheckpointTs,
+					})
+					toRemoveDispatcherIDs = append(toRemoveDispatcherIDs, id)
+					removedDispatcherSchemaIDs = append(removedDispatcherSchemaIDs, dispatcherItem.GetSchemaID())
+				}
+			}
+
+			message.Watermark.UpdateMin(heartBeatInfo.Watermark)
+			if needCompleteStatus {
 				message.Statuses = append(message.Statuses, &heartbeatpb.TableSpanStatus{
-					ID:              id.ToPB(),
-					ComponentStatus: heartbeatpb.ComponentState_Stopped,
-					CheckpointTs:    watermark.CheckpointTs,
+					ID:                 id.ToPB(),
+					ComponentStatus:    heartBeatInfo.ComponentStatus,
+					CheckpointTs:       heartBeatInfo.Watermark.CheckpointTs,
+					EventSizePerSecond: heartBeatInfo.EventSizePerSecond,
 				})
-				toRemoveDispatcherIDs = append(toRemoveDispatcherIDs, id)
-				removedDispatcherSchemaIDs = append(removedDispatcherSchemaIDs, dispatcherItem.GetSchemaID())
 			}
-		}
-
-		message.Watermark.UpdateMin(heartBeatInfo.Watermark)
-		if needCompleteStatus {
-			message.Statuses = append(message.Statuses, &heartbeatpb.TableSpanStatus{
-				ID:                 id.ToPB(),
-				ComponentStatus:    heartBeatInfo.ComponentStatus,
-				CheckpointTs:       heartBeatInfo.Watermark.CheckpointTs,
-				EventSizePerSecond: heartBeatInfo.EventSizePerSecond,
-			})
-		}
-	})
+		})
+	} else {
+		minWatermark := e.dispatcherMap.GetMinWatermark()
+		message.Watermark = &minWatermark
+	}
 
 	for idx, id := range toRemoveDispatcherIDs {
 		e.cleanTableEventDispatcher(id, removedDispatcherSchemaIDs[idx])
@@ -617,12 +623,27 @@ func (e *EventDispatcherManager) GetAllDispatchers(schemaID int64) []common.Disp
 }
 
 type DispatcherMap struct {
-	m sync.Map
+	m              sync.Map
+	resolvedTsHeap struct {
+		sync.RWMutex
+		h *heap.Heap[*dispatcher.TsItem]
+	}
+	// checkpointsHeap struct {
+	// 	sync.RWMutex
+	// 	h *heap.Heap[*dispatcher.TsItem]
+	// }
 }
 
 func newDispatcherMap() *DispatcherMap {
 	return &DispatcherMap{
 		m: sync.Map{},
+		resolvedTsHeap: struct {
+			sync.RWMutex
+			h *heap.Heap[*dispatcher.TsItem]
+		}{
+			RWMutex: sync.RWMutex{},
+			h:       heap.NewHeap[*dispatcher.TsItem](),
+		},
 	}
 }
 
@@ -644,11 +665,21 @@ func (d *DispatcherMap) Get(id common.DispatcherID) (*dispatcher.Dispatcher, boo
 }
 
 func (d *DispatcherMap) Delete(id common.DispatcherID) {
+	dispatcherItem, ok := d.m.Load(id)
+	if ok {
+		log.Error("delete an empty dispatcher", zap.Any("dispatcher id", id))
+	}
 	d.m.Delete(id)
+
+	d.resolvedTsHeap.Lock()
+	defer d.resolvedTsHeap.Unlock()
+	d.resolvedTsHeap.h.Remove(dispatcherItem.(*dispatcher.Dispatcher).ResolvedTs)
 }
 
 func (d *DispatcherMap) Set(id common.DispatcherID, dispatcher *dispatcher.Dispatcher) {
 	d.m.Store(id, dispatcher)
+	dispatcher.ResolvedTs.SetOnUpdate(d.onUpdateResolvedTs)
+	d.onUpdateResolvedTs(dispatcher.ResolvedTs)
 }
 
 func (d *DispatcherMap) ForEach(fn func(id common.DispatcherID, dispatcher *dispatcher.Dispatcher)) {
@@ -657,3 +688,34 @@ func (d *DispatcherMap) ForEach(fn func(id common.DispatcherID, dispatcher *disp
 		return true
 	})
 }
+
+func (d *DispatcherMap) onUpdateResolvedTs(tsItem *dispatcher.TsItem) {
+	d.resolvedTsHeap.Lock()
+	defer d.resolvedTsHeap.Unlock()
+	d.resolvedTsHeap.h.AddOrUpdate(tsItem)
+}
+
+// func (d *DispatcherMap) onUpdateCheckpointTs(tsItem *dispatcher.TsItem) {
+// 	d.resolvedTsHeap.Lock()
+// 	defer d.resolvedTsHeap.Unlock()
+// 	d.checkpointsHeap.h.AddOrUpdate(tsItem)
+// }
+
+func (d *DispatcherMap) GetMinWatermark() (w heartbeatpb.Watermark) {
+	d.resolvedTsHeap.RLock()
+	r, ok := d.resolvedTsHeap.h.PeekTop()
+	d.resolvedTsHeap.RUnlock()
+	if !ok {
+		return
+	}
+	w.ResolvedTs = r.Get()
+	w.CheckpointTs = w.ResolvedTs
+
+	// TODO: implement checkpointsHeap
+	// c, ok := d.checkpointsHeap.PeekTop()
+	// if ok {
+	// 	return
+	// }
+	// w.CheckpointTs = c.Get()
+	return
+}

downstreamadapter/dispatchermanager/event_dispatcher_manager_test.go

@@ -0,0 +1,61 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dispatchermanager
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/pingcap/ticdc/downstreamadapter/dispatcher"
+	"github.com/pingcap/ticdc/pkg/common"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDispatcherResolvedTsHeap(t *testing.T) {
+	dMap := newDispatcherMap()
+	dispatchers := []*dispatcher.Dispatcher{
+		{
+			ResolvedTs: dispatcher.NewTsItem(1),
+		},
+		{
+			ResolvedTs: dispatcher.NewTsItem(2),
+		},
+		{
+			ResolvedTs: dispatcher.NewTsItem(3),
+		},
+		{
+			ResolvedTs: dispatcher.NewTsItem(4),
+		},
+		{
+			ResolvedTs: dispatcher.NewTsItem(5),
+		},
+	}
+	for _, d := range dispatchers {
+		dMap.Set(common.NewDispatcherID(), d)
+	}
+
+	for idx, d := range dispatchers {
+		fmt.Println(idx, ":", d.ResolvedTs)
+	}
+
+	require.Equal(t, 5, dMap.Len())
+	require.Equal(t, uint64(1), dMap.GetMinWatermark().ResolvedTs)
+	d := dispatchers[0]
+
+	d.ResolvedTs.Set(2)
+	require.Equal(t, uint64(2), dMap.GetMinWatermark().ResolvedTs)
+
+	d.ResolvedTs.Set(3)
+	require.Equal(t, uint64(2), dMap.GetMinWatermark().ResolvedTs)
+}