fix: unittest

Signed-off-by: chyezh <[email protected]>

Author: chyezh

internal/flushcommon/pipeline/data_sync_service.go

@@ -410,6 +410,7 @@ func NewDataSyncServiceWithMetaCache(metaCache metacache.MetaCache) *DataSyncSer
 	return &DataSyncService{metacache: metaCache}
 }
 
+// NewEmptyStreamingNodeDataSyncService is used to create a new data sync service when incoming create collection message.
 func NewEmptyStreamingNodeDataSyncService(
 	initCtx context.Context,
 	pipelineParams *util.PipelineParams,

internal/flushcommon/util/checkpoint_updater.go

@@ -48,8 +48,9 @@ type ChannelCheckpointUpdater struct {
 	tasks      map[string]*channelCPUpdateTask
 	notifyChan chan struct{}
 
-	closeCh   chan struct{}
-	closeOnce sync.Once
+	closeCh            chan struct{}
+	closeOnce          sync.Once
+	updateDoneCallback func(*msgpb.MsgPosition)
 }
 
 func NewChannelCheckpointUpdater(broker broker.Broker) *ChannelCheckpointUpdater {
@@ -61,6 +62,17 @@ func NewChannelCheckpointUpdater(broker broker.Broker) *ChannelCheckpointUpdater
 	}
 }
 
+// NewChannelCheckpointUpdaterWithCallback creates a ChannelCheckpointUpdater with a callback function
+func NewChannelCheckpointUpdaterWithCallback(broker broker.Broker, updateDoneCallback func(*msgpb.MsgPosition)) *ChannelCheckpointUpdater {
+	return &ChannelCheckpointUpdater{
+		broker:             broker,
+		tasks:              make(map[string]*channelCPUpdateTask),
+		closeCh:            make(chan struct{}),
+		notifyChan:         make(chan struct{}, 1),
+		updateDoneCallback: updateDoneCallback,
+	}
+}
+
 func (ccu *ChannelCheckpointUpdater) Start() {
 	log.Info("channel checkpoint updater start")
 	ticker := time.NewTicker(paramtable.Get().DataNodeCfg.ChannelCheckpointUpdateTickInSeconds.GetAsDuration(time.Second))
@@ -134,6 +146,9 @@ func (ccu *ChannelCheckpointUpdater) updateCheckpoints(tasks []*channelCPUpdateT
 				for _, task := range tasks {
 					task.callback()
 					finished.Insert(task.pos.GetChannelName(), task)
+					if ccu.updateDoneCallback != nil {
+						ccu.updateDoneCallback(task.pos)
+					}
 				}
 			}(tasks)
 		}

internal/rootcoord/meta_table.go

@@ -925,6 +925,11 @@ func (mt *MetaTable) GetPChannelInfo(ctx context.Context, pchannel string) *root
 		Collections: make([]*rootcoordpb.CollectionInfoOnPChannel, 0),
 	}
 	for _, collInfo := range mt.collID2Meta {
+		if collInfo.State != pb.CollectionState_CollectionCreated {
+			// streamingnode, skip non-created collections when recovering
+			// streamingnode will receive the createCollectionMessage to recover if the collection is creating.
+			continue
+		}
 		if idx := lo.IndexOf(collInfo.PhysicalChannelNames, pchannel); idx >= 0 {
 			partitions := make([]*rootcoordpb.PartitionInfoOnPChannel, 0, len(collInfo.Partitions))
 			for _, part := range collInfo.Partitions {

internal/streamingnode/server/flusher/flusherimpl/data_service_wrapper.go

@@ -9,12 +9,21 @@ import (
 	"github.com/milvus-io/milvus/pkg/streaming/util/message/adaptor"
 )
 
-// dataSyncServiceWrapper wraps DataSyncService with related fields.
+// newDataSyncServiceWrapper creates a new data sync service wrapper.
+func newDataSyncServiceWrapper(input chan<- *msgstream.MsgPack, ds *pipeline.DataSyncService) *dataSyncServiceWrapper {
+	handler := adaptor.NewBaseMsgPackAdaptorHandler()
+	return &dataSyncServiceWrapper{
+		input:   input,
+		handler: handler,
+		ds:      ds,
+	}
+}
+
+// dataSyncServiceWrapper wraps DataSyncService and related input channel.
 type dataSyncServiceWrapper struct {
-	input          chan<- *msgstream.MsgPack
-	handler        *adaptor.BaseMsgPackAdaptorHandler
-	ds             *pipeline.DataSyncService
-	startMessageID message.MessageID
+	input   chan<- *msgstream.MsgPack
+	handler *adaptor.BaseMsgPackAdaptorHandler
+	ds      *pipeline.DataSyncService
 }
 
 // Start starts the data sync service.
@@ -30,6 +39,7 @@ func (ds *dataSyncServiceWrapper) HandleMessage(ctx context.Context, msg message
 		case <-ctx.Done():
 			return ctx.Err()
 		case ds.input <- ds.handler.PendingMsgPack.Next():
+			// The input channel will never get stuck because the data sync service will consume the message continuously.
 			ds.handler.PendingMsgPack.UnsafeAdvance()
 		}
 	}
@@ -38,6 +48,7 @@ func (ds *dataSyncServiceWrapper) HandleMessage(ctx context.Context, msg message
 
 // Close close the input channel and gracefully close the data sync service.
 func (ds *dataSyncServiceWrapper) Close() {
+	// The input channel should be closed first, otherwise the flowgraph in datasync service will be blocked.
 	close(ds.input)
 	ds.ds.GracefullyClose()
 }

internal/streamingnode/server/flusher/flusherimpl/flusher_components.go

@@ -24,88 +24,28 @@ import (
 	"github.com/milvus-io/milvus/pkg/streaming/util/message"
 	"github.com/milvus-io/milvus/pkg/streaming/util/message/adaptor"
 	"github.com/milvus-io/milvus/pkg/util/conc"
-	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/retry"
 )
 
-// buildFlusherComponents builds the components of the flusher.
-func (impl *WALFlusherImpl) buildFlusherComponents(ctx context.Context, l wal.WAL) (*flusherComponents, error) {
-	// Get all existed vchannels of the pchannel.
-	vchannels, err := impl.getVchannels(ctx, l.Channel().Name)
-	if err != nil {
-		impl.logger.Warn("get vchannels failed", zap.Error(err))
-		return nil, err
-	}
-	impl.logger.Info("fetch vchannel done", zap.Int("vchannelNum", len(vchannels)))
-
-	// Get all the recovery info of the recoverable vchannels.
-	recoverInfos, err := impl.getRecoveryInfos(ctx, vchannels)
-	if err != nil {
-		impl.logger.Warn("get recovery info failed", zap.Error(err))
-		return nil, err
-	}
-	impl.logger.Info("fetch recovery info done", zap.Int("recoveryInfoNum", len(recoverInfos)))
-
-	// build up components
-	dc, err := resource.Resource().DataCoordClient().GetWithContext(ctx)
-	if err != nil {
-		impl.logger.Warn("flusher recovery is canceled before data coord client ready", zap.Error(err))
-		return nil, err
-	}
-	impl.logger.Info("data coord client ready")
-
-	broker := broker.NewCoordBroker(dc, paramtable.GetNodeID())
-	chunkManager := resource.Resource().ChunkManager()
-	syncMgr := syncmgr.NewSyncManager(chunkManager)
-	wbMgr := writebuffer.NewManager(syncMgr)
-	wbMgr.Start()
-	cpUpdater := util.NewChannelCheckpointUpdater(broker)
-	go cpUpdater.Start()
-
-	fc := &flusherComponents{
-		wal:          l,
-		broker:       broker,
-		syncMgr:      syncMgr,
-		wbMgr:        wbMgr,
-		cpUpdater:    cpUpdater,
-		chunkManager: chunkManager,
-		dataServices: make(map[string]*dataSyncServiceWrapper),
-		logger:       impl.logger,
-	}
-	impl.logger.Info("flusher components intiailizing done")
-	if err := fc.recover(ctx, recoverInfos); err != nil {
-		impl.logger.Warn("flusher recovery is canceled before recovery done, recycle the resource", zap.Error(err))
-		fc.Close()
-		impl.logger.Info("flusher recycle the resource done")
-		return nil, err
-	}
-	impl.logger.Info("flusher recovery done")
-	return fc, nil
-}
-
 // flusherComponents is the components of the flusher.
 type flusherComponents struct {
-	wal          wal.WAL
-	broker       broker.Broker
-	syncMgr      syncmgr.SyncManager
-	wbMgr        writebuffer.BufferManager
-	cpUpdater    *util.ChannelCheckpointUpdater
-	chunkManager storage.ChunkManager
-	dataServices map[string]*dataSyncServiceWrapper
-	logger       *log.MLogger
+	wal               wal.WAL
+	broker            broker.Broker
+	syncMgr           syncmgr.SyncManager
+	wbMgr             writebuffer.BufferManager
+	cpUpdater         *util.ChannelCheckpointUpdater
+	chunkManager      storage.ChunkManager
+	dataServices      map[string]*dataSyncServiceWrapper
+	checkpointManager *pchannelCheckpointManager
+	logger            *log.MLogger
 }
 
-// GetMinimumStartMessage gets the minimum start message of all the data services.
-func (impl *flusherComponents) GetMinimumStartMessage() message.MessageID {
-	var startMessageID message.MessageID
-	for _, ds := range impl.dataServices {
-		if startMessageID == nil || ds.startMessageID.LT(startMessageID) {
-			startMessageID = ds.startMessageID
-		}
-	}
-	return startMessageID
+// StartMessageID returns the start message id of the flusher after recovering.
+func (impl *flusherComponents) StartMessageID() message.MessageID {
+	return impl.checkpointManager.StartMessageID()
 }
 
+// WhenCreateCollection handles the create collection message.
 func (impl *flusherComponents) WhenCreateCollection(createCollectionMsg message.ImmutableCreateCollectionMessageV1) {
 	if _, ok := impl.dataServices[createCollectionMsg.VChannel()]; ok {
 		impl.logger.Info("the data sync service of current vchannel is built, skip it", zap.String("vchannel", createCollectionMsg.VChannel()))
@@ -172,8 +112,10 @@ func (impl *flusherComponents) WhenDropCollection(vchannel string) {
 		delete(impl.dataServices, vchannel)
 		impl.logger.Info("drop data sync service", zap.String("vchannel", vchannel))
 	}
+	impl.checkpointManager.DropVChannel(vchannel)
 }
 
+// HandleMessage handles the plain message.
 func (impl *flusherComponents) HandleMessage(ctx context.Context, msg message.ImmutableMessage) error {
 	vchannel := msg.VChannel()
 	if vchannel == "" {
@@ -185,6 +127,7 @@ func (impl *flusherComponents) HandleMessage(ctx context.Context, msg message.Im
 	return impl.dataServices[vchannel].HandleMessage(ctx, msg)
 }
 
+// broadcastToAllDataSyncService broadcasts the message to all data sync services.
 func (impl *flusherComponents) broadcastToAllDataSyncService(ctx context.Context, msg message.ImmutableMessage) error {
 	for _, ds := range impl.dataServices {
 		if err := ds.HandleMessage(ctx, msg); err != nil {
@@ -194,17 +137,14 @@ func (impl *flusherComponents) broadcastToAllDataSyncService(ctx context.Context
 	return nil
 }
 
+// addNewDataSyncService adds a new data sync service to the components when new collection is created.
 func (impl *flusherComponents) addNewDataSyncService(
 	createCollectionMsg message.ImmutableCreateCollectionMessageV1,
 	input chan<- *msgstream.MsgPack,
 	ds *pipeline.DataSyncService,
 ) {
-	newDS := &dataSyncServiceWrapper{
-		input:          input,
-		handler:        adaptor.NewBaseMsgPackAdaptorHandler(),
-		ds:             ds,
-		startMessageID: createCollectionMsg.LastConfirmedMessageID(),
-	}
+	impl.checkpointManager.AddVChannel(createCollectionMsg.VChannel(), createCollectionMsg.LastConfirmedMessageID())
+	newDS := newDataSyncServiceWrapper(input, ds)
 	newDS.Start()
 	impl.dataServices[createCollectionMsg.VChannel()] = newDS
 	impl.logger.Info("create data sync service done", zap.String("vchannel", createCollectionMsg.VChannel()))
@@ -219,6 +159,7 @@ func (impl *flusherComponents) Close() {
 	impl.wbMgr.Stop()
 	impl.cpUpdater.Close()
 	impl.syncMgr.Close()
+	impl.checkpointManager.Close()
 }
 
 // recover recover the components of the flusher.
@@ -254,6 +195,7 @@ func (impl *flusherComponents) recover(ctx context.Context, recoverInfos map[str
 	return nil
 }
 
+// buildDataSyncServiceWithRetry builds the data sync service with retry.
 func (impl *flusherComponents) buildDataSyncServiceWithRetry(ctx context.Context, recoverInfo *datapb.GetChannelRecoveryInfoResponse) (*dataSyncServiceWrapper, error) {
 	var ds *dataSyncServiceWrapper
 	err := retry.Do(ctx, func() error {
@@ -267,6 +209,7 @@ func (impl *flusherComponents) buildDataSyncServiceWithRetry(ctx context.Context
 	return ds, nil
 }
 
+// buildDataSyncService builds the data sync service with given recovery info.
 func (impl *flusherComponents) buildDataSyncService(ctx context.Context, recoverInfo *datapb.GetChannelRecoveryInfoResponse) (*dataSyncServiceWrapper, error) {
 	// Build and add pipeline.
 	input := make(chan *msgstream.MsgPack, 10)
@@ -300,10 +243,5 @@ func (impl *flusherComponents) buildDataSyncService(ctx context.Context, recover
 	if err != nil {
 		return nil, err
 	}
-	return &dataSyncServiceWrapper{
-		input:          input,
-		handler:        adaptor.NewBaseMsgPackAdaptorHandler(),
-		ds:             ds,
-		startMessageID: adaptor.MustGetMessageIDFromMQWrapperIDBytes(impl.wal.WALName(), recoverInfo.GetInfo().GetSeekPosition().GetMsgID()),
-	}, nil
+	return newDataSyncServiceWrapper(input, ds), nil
 }