Fix potential crash bug of bulkinsert (#24763)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
milvus-io · Jun 13, 2023 · 5c3b744 · 5c3b744
1 parent a3437e0
commit 5c3b744
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/internal/datanode/mock_test.go b/internal/datanode/mock_test.go
@@ -230,10 +230,14 @@ type DataCoordFactory struct {
 
 	AddSegmentError      bool
 	AddSegmentNotSuccess bool
+	AddSegmentEmpty      bool
 }
 
 func (ds *DataCoordFactory) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentIDRequest) (*datapb.AssignSegmentIDResponse, error) {
-	return &datapb.AssignSegmentIDResponse{
+	if ds.AddSegmentError {
+		return nil, errors.New("Error")
+	}
+	res := &datapb.AssignSegmentIDResponse{
 		Status: &commonpb.Status{
 			ErrorCode: commonpb.ErrorCode_Success,
 		},
@@ -242,7 +246,15 @@ func (ds *DataCoordFactory) AssignSegmentID(ctx context.Context, req *datapb.Ass
 				SegID: 666,
 			},
 		},
-	}, nil
+	}
+	if ds.AddSegmentNotSuccess {
+		res.Status = &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_UnexpectedError,
+		}
+	} else if ds.AddSegmentEmpty {
+		res.SegIDAssignments = []*datapb.SegmentIDAssignment{}
+	}
+	return res, nil
 }
 
 func (ds *DataCoordFactory) CompleteCompaction(ctx context.Context, req *datapb.CompactionResult) (*commonpb.Status, error) {

diff --git a/internal/datanode/services.go b/internal/datanode/services.go
@@ -642,6 +642,9 @@ func assignSegmentFunc(node *DataNode, req *datapb.ImportTaskRequest) importutil
 		if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
 			return 0, "", fmt.Errorf("syncSegmentID Failed:%s", resp.Status.Reason)
 		}
+		if len(resp.SegIDAssignments) == 0 || resp.SegIDAssignments[0] == nil {
+			return 0, "", fmt.Errorf("syncSegmentID Failed: the collection was dropped")
+		}
 		segmentID := resp.SegIDAssignments[0].SegID
 		log.Info("new segment assigned",
 			zap.Int64("task ID", importTaskID),

diff --git a/internal/datanode/services_test.go b/internal/datanode/services_test.go
@@ -452,6 +452,11 @@ func (s *DataNodeServicesSuite) TestImport() {
 		s.Assert().NoError(err)
 		s.node.dataCoord.(*DataCoordFactory).AddSegmentNotSuccess = false
 
+		s.node.dataCoord.(*DataCoordFactory).AddSegmentEmpty = true
+		_, err = s.node.Import(context.WithValue(s.ctx, ctxKey{}, ""), req)
+		s.Assert().NoError(err)
+		s.node.dataCoord.(*DataCoordFactory).AddSegmentEmpty = false
+
 		stat, err := s.node.Import(context.WithValue(s.ctx, ctxKey{}, ""), req)
 		s.Assert().NoError(err)
 		s.Assert().True(merr.Ok(stat))