fix: correct AssignSegmentID return and add retry for loadCollectionF… (#32349)

pr: #32335

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

internal/datacoord/handler.go

@@ -394,9 +394,19 @@ func (h *ServerHandler) GetCollection(ctx context.Context, collectionID UniqueID
 	if coll != nil {
 		return coll, nil
 	}
-	err := h.s.loadCollectionFromRootCoord(ctx, collectionID)
-	if err != nil {
-		log.Warn("failed to load collection from rootcoord", zap.Int64("collectionID", collectionID), zap.Error(err))
+	ctx2, cancel := context.WithTimeout(ctx, time.Second*10)
+	defer cancel()
+	if err := retry.Do(ctx2, func() error {
+		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
+		}
+		return nil
+	}, retry.Attempts(5)); err != nil {
+		log.Ctx(ctx2).Warn("datacoord ServerHandler GetCollection finally failed",
+			zap.Int64("collectionID", collectionID),
+			zap.Error(err))
 		return nil, err
 	}
 

internal/datacoord/services.go

@@ -194,6 +194,12 @@ 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))
+			assigns = append(assigns, &datapb.SegmentIDAssignment{
+				ChannelName:  r.ChannelName,
+				CollectionID: r.CollectionID,
+				PartitionID:  r.PartitionID,
+				Status:       merr.Status(err),
+			})
 			continue
 		}
 

internal/datacoord/services_test.go

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