fix: modify retry error (#32351)

#32322

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>

internal/datacoord/handler.go

@@ -397,7 +397,7 @@ func (h *ServerHandler) GetCollection(ctx context.Context, collectionID UniqueID
 	ctx2, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer cancel()
 	if err := retry.Do(ctx2, func() error {
-		err := h.s.loadCollectionFromRootCoord(ctx, collectionID)
+		err := h.s.loadCollectionFromRootCoord(ctx2, collectionID)
 		if err != nil {
 			log.Warn("failed to load collection from rootcoord", zap.Int64("collectionID", collectionID), zap.Error(err))
 			return err
@@ -407,6 +407,7 @@ func (h *ServerHandler) GetCollection(ctx context.Context, collectionID UniqueID
 		log.Ctx(ctx2).Warn("datacoord ServerHandler GetCollection finally failed",
 			zap.Int64("collectionID", collectionID),
 			zap.Error(err))
+		return nil, err
 	}
 
 	return h.s.meta.GetCollection(collectionID), nil

internal/datacoord/services.go

@@ -194,7 +194,13 @@ func (s *Server) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentI
 			r.CollectionID, r.PartitionID, r.ChannelName, int64(r.Count))
 		if err != nil {
 			log.Warn("failed to alloc segment", zap.Any("request", r), zap.Error(err))
-			return nil, err
+			assigns = append(assigns, &datapb.SegmentIDAssignment{
+				ChannelName:  r.ChannelName,
+				CollectionID: r.CollectionID,
+				PartitionID:  r.PartitionID,
+				Status:       merr.Status(err),
+			})
+			continue
 		}
 
 		log.Info("success to assign segments", zap.Int64("collectionID", r.GetCollectionID()), zap.Any("assignments", segmentAllocations))

internal/datacoord/services_test.go

@@ -950,7 +950,7 @@ func (s *ServerSuite) TestAssignSegmentID() {
 			SegmentIDRequests: []*datapb.SegmentIDRequest{req},
 		})
 		s.NoError(err)
-		s.EqualValues(0, len(resp.SegIDAssignments))
+		s.EqualValues(1, len(resp.SegIDAssignments))
 	})
 }