Set start positions in datanode (#5757)

Signed-off-by: Nyangxuan <xuan.yang@zilliz.com>

internal/datanode/collection_replica.go

@@ -20,6 +20,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/types"
 )
@@ -42,6 +43,7 @@ type Replica interface {
 	getChannelName(segID UniqueID) (string, error)
 	setStartPositions(segmentID UniqueID, startPos []*internalpb.MsgPosition) error
 	setEndPositions(segmentID UniqueID, endPos []*internalpb.MsgPosition) error
+	getAllStartPositions() []*datapb.SegmentStartPosition
 	getSegmentPositions(segID UniqueID) ([]*internalpb.MsgPosition, []*internalpb.MsgPosition)
 	listOpenSegmentCheckPointAndNumRows(segs []UniqueID) (map[UniqueID]internalpb.MsgPosition, map[UniqueID]int64)
 }
@@ -185,6 +187,32 @@ func (replica *CollectionSegmentReplica) addSegment(
 	return nil
 }
 
+func (replica *CollectionSegmentReplica) getAllStartPositions() []*datapb.SegmentStartPosition {
+	replica.mu.RLock()
+	defer replica.mu.RUnlock()
+
+	result := make([]*datapb.SegmentStartPosition, 0, len(replica.segments))
+	for id, seg := range replica.segments {
+
+		if seg.isNew.Load().(bool) {
+
+			pos, ok := replica.startPositions[id]
+			if !ok {
+				log.Warn("Segment has no start positions")
+				continue
+			}
+
+			result = append(result, &datapb.SegmentStartPosition{
+				SegmentID:     id,
+				StartPosition: pos[0],
+			})
+			seg.isNew.Store(false)
+		}
+
+	}
+	return result
+}
+
 func (replica *CollectionSegmentReplica) removeSegment(segmentID UniqueID) error {
 	replica.mu.Lock()
 	delete(replica.segments, segmentID)

internal/datanode/data_sync_service.go

@@ -119,6 +119,7 @@ func (dsService *dataSyncService) initNodes(vchanInfo *datapb.VchannelInfo) {
 			zap.Int64("SegmentID", fu.segID),
 			zap.Int64("CollectionID", fu.collID),
 			zap.Int("Length of Field2BinlogPaths", len(id2path)),
+			zap.Any("Start Positions", fu.startPositions),
 		)
 
 		req := &datapb.SaveBinlogPathsRequest{
@@ -132,6 +133,7 @@ func (dsService *dataSyncService) initNodes(vchanInfo *datapb.VchannelInfo) {
 			CollectionID:      fu.collID,
 			Field2BinlogPaths: id2path,
 			CheckPoints:       checkPoints,
+			StartPositions:    fu.startPositions,
 			Flushed:           fu.flushed,
 		}
 		rsp, err := dsService.dataService.SaveBinlogPaths(dsService.ctx, req)

internal/datanode/flow_graph_insert_buffer_node.go

@@ -72,11 +72,12 @@ type segmentCheckPoint struct {
 }
 
 type segmentFlushUnit struct {
-	collID     UniqueID
-	segID      UniqueID
-	field2Path map[UniqueID]string
-	checkPoint map[UniqueID]segmentCheckPoint
-	flushed    bool
+	collID         UniqueID
+	segID          UniqueID
+	field2Path     map[UniqueID]string
+	checkPoint     map[UniqueID]segmentCheckPoint
+	startPositions []*datapb.SegmentStartPosition
+	flushed        bool
 }
 
 type insertBuffer struct {
@@ -714,7 +715,9 @@ func flushSegment(
 	_, ep := ibNode.replica.getSegmentPositions(segID)
 	sta, _ := ibNode.replica.getSegmentStatisticsUpdates(segID)
 	ibNode.setSegmentCheckPoint(segID, segmentCheckPoint{sta.NumRows, *ep[0]})
-	flushUnit <- segmentFlushUnit{collID: collID, segID: segID, field2Path: field2Path}
+
+	startPos := ibNode.replica.getAllStartPositions()
+	flushUnit <- segmentFlushUnit{collID: collID, segID: segID, field2Path: field2Path, startPositions: startPos}
 	clearFn(true)
 }
 

internal/proto/data_service.proto

@@ -245,9 +245,9 @@ message ID2PathList {
   repeated string Paths = 2;
 }
 
-message PositionPair {
+message SegmentStartPosition {
   internal.MsgPosition start_position = 1;
-  internal.MsgPosition end_position = 2;
+  int64 segmentID = 2;
 }
 
 message SaveBinlogPathsRequest {
@@ -256,7 +256,8 @@ message SaveBinlogPathsRequest {
   int64 collectionID = 3;
   repeated ID2PathList field2BinlogPaths = 4;
   repeated CheckPoint checkPoints = 5;
-  bool flushed = 6;
+  repeated SegmentStartPosition start_positions = 6;
+  bool flushed = 7;
 }
 
 message CheckPoint {