Verify target id for FlushSegments request in datanode (#20308)

Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>
Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>

internal/datacoord/cluster.go

@@ -104,6 +104,7 @@ func (c *Cluster) Flush(ctx context.Context, nodeID int64, channel string,
 		Base: commonpbutil.NewMsgBase(
 			commonpbutil.WithMsgType(commonpb.MsgType_Flush),
 			commonpbutil.WithSourceID(Params.DataCoordCfg.GetNodeID()),
+			commonpbutil.WithTargetID(nodeID),
 		),
 		CollectionID:   ch.CollectionID,
 		SegmentIDs:     lo.Map(segments, getSegmentID),

internal/datanode/data_node.go

@@ -594,6 +594,18 @@ func (node *DataNode) FlushSegments(ctx context.Context, req *datapb.FlushSegmen
 		return errStatus, nil
 	}
 
+	if req.GetBase().GetTargetID() != node.session.ServerID {
+		log.Warn("flush segment target id not matched",
+			zap.Int64("targetID", req.GetBase().GetTargetID()),
+			zap.Int64("serverID", node.session.ServerID),
+		)
+		status := &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_NodeIDNotMatch,
+			Reason:    common.WrapNodeIDNotMatchMsg(req.GetBase().GetTargetID(), node.session.ServerID),
+		}
+		return status, nil
+	}
+
 	log.Info("receiving FlushSegments request",
 		zap.Int64("collection ID", req.GetCollectionID()),
 		zap.Int64s("segments", req.GetSegmentIDs()),

internal/datanode/data_node_test.go

@@ -227,7 +227,9 @@ func TestDataNode(t *testing.T) {
 		assert.Nil(t, err)
 
 		req := &datapb.FlushSegmentsRequest{
-			Base:         &commonpb.MsgBase{},
+			Base: &commonpb.MsgBase{
+				TargetID: node1.session.ServerID,
+			},
 			DbID:         0,
 			CollectionID: 1,
 			SegmentIDs:   []int64{0},
@@ -288,7 +290,21 @@ func TestDataNode(t *testing.T) {
 
 		// failure call
 		req = &datapb.FlushSegmentsRequest{
-			Base:         &commonpb.MsgBase{},
+			Base: &commonpb.MsgBase{
+				TargetID: -1,
+			},
+			DbID:         0,
+			CollectionID: 1,
+			SegmentIDs:   []int64{1},
+		}
+		status, err = node1.FlushSegments(node1.ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_NodeIDNotMatch, status.ErrorCode)
+
+		req = &datapb.FlushSegmentsRequest{
+			Base: &commonpb.MsgBase{
+				TargetID: node1.session.ServerID,
+			},
 			DbID:         0,
 			CollectionID: 1,
 			SegmentIDs:   []int64{1},
@@ -299,7 +315,9 @@ func TestDataNode(t *testing.T) {
 		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
 
 		req = &datapb.FlushSegmentsRequest{
-			Base:           &commonpb.MsgBase{},
+			Base: &commonpb.MsgBase{
+				TargetID: node1.session.ServerID,
+			},
 			DbID:           0,
 			CollectionID:   1,
 			SegmentIDs:     []int64{},