Fix potential crash bug of bulkinsert (milvus-io#24596)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
zhagnlu · Jun 2, 2023 · 9ffcd53 · 9ffcd53
1 parent 1a60fe4
commit 9ffcd53
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
@@ -1233,6 +1233,10 @@ 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/data_node_test.go b/internal/datanode/data_node_test.go
@@ -513,6 +513,11 @@ func TestDataNode(t *testing.T) {
 		assert.NoError(t, err)
 		node.dataCoord.(*DataCoordFactory).AddSegmentNotSuccess = false
 
+		node.dataCoord.(*DataCoordFactory).AddSegmentEmpty = true
+		_, err = node.Import(context.WithValue(ctx, ctxKey{}, ""), req)
+		assert.NoError(t, err)
+		node.dataCoord.(*DataCoordFactory).AddSegmentEmpty = false
+
 		stat, err := node.Import(context.WithValue(ctx, ctxKey{}, ""), req)
 		assert.NoError(t, err)
 		assert.Equal(t, commonpb.ErrorCode_Success, stat.GetErrorCode())

diff --git a/internal/datanode/mock_test.go b/internal/datanode/mock_test.go
@@ -228,10 +228,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,
 		},
@@ -240,7 +244,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) {