enhance: Add ctx trace for segment load prepare (milvus-io#31203)

Add ctx paramter to `prepare` and add trace id in underline logs. Signed-off-by: Congqi Xia <[email protected]>
bigsheeper · Mar 13, 2024 · 1d96239 · 1d96239
1 parent 5b51c20
commit 1d96239
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go
@@ -133,10 +133,10 @@ func (loader *segmentLoaderV2) Load(ctx context.Context,
 		return nil, nil
 	}
 	// Filter out loaded & loading segments
-	infos := loader.prepare(segmentType, segments...)
+	infos := loader.prepare(ctx, segmentType, segments...)
 	defer loader.unregister(infos...)
 
-	log.With(
+	log = log.With(
 		zap.Int64s("requestSegments", lo.Map(segments, func(s *querypb.SegmentLoadInfo, _ int) int64 { return s.GetSegmentID() })),
 		zap.Int64s("preparedSegments", lo.Map(infos, func(s *querypb.SegmentLoadInfo, _ int) int64 { return s.GetSegmentID() })),
 	)
@@ -559,10 +559,10 @@ func (loader *segmentLoader) Load(ctx context.Context,
 		return nil, nil
 	}
 	// Filter out loaded & loading segments
-	infos := loader.prepare(segmentType, segments...)
+	infos := loader.prepare(ctx, segmentType, segments...)
 	defer loader.unregister(infos...)
 
-	log.With(
+	log = log.With(
 		zap.Int64s("requestSegments", lo.Map(segments, func(s *querypb.SegmentLoadInfo, _ int) int64 { return s.GetSegmentID() })),
 		zap.Int64s("preparedSegments", lo.Map(infos, func(s *querypb.SegmentLoadInfo, _ int) int64 { return s.GetSegmentID() })),
 	)
@@ -694,7 +694,10 @@ func (loader *segmentLoader) Load(ctx context.Context,
 	return result, nil
 }
 
-func (loader *segmentLoader) prepare(segmentType SegmentType, segments ...*querypb.SegmentLoadInfo) []*querypb.SegmentLoadInfo {
+func (loader *segmentLoader) prepare(ctx context.Context, segmentType SegmentType, segments ...*querypb.SegmentLoadInfo) []*querypb.SegmentLoadInfo {
+	log := log.Ctx(ctx).With(
+		zap.Stringer("segmentType", segmentType),
+	)
 	loader.mut.Lock()
 	defer loader.mut.Unlock()
 
@@ -707,7 +710,8 @@ func (loader *segmentLoader) prepare(segmentType SegmentType, segments ...*query
 			infos = append(infos, segment)
 			loader.loadingSegments.Insert(segment.GetSegmentID(), newLoadResult())
 		} else {
-			log.Info("skip loaded/loading segment", zap.Int64("segmentID", segment.GetSegmentID()),
+			log.Info("skip loaded/loading segment",
+				zap.Int64("segmentID", segment.GetSegmentID()),
 				zap.Bool("isLoaded", len(loader.manager.Segment.GetBy(WithType(segmentType), WithID(segment.GetSegmentID()))) > 0),
 				zap.Bool("isLoading", loader.loadingSegments.Contain(segment.GetSegmentID())),
 			)
@@ -1502,7 +1506,7 @@ func (loader *segmentLoader) LoadIndex(ctx context.Context, segment *LocalSegmen
 
 	// Filter out LOADING segments only
 	// use None to avoid loaded check
-	infos := loader.prepare(commonpb.SegmentState_SegmentStateNone, loadInfo)
+	infos := loader.prepare(ctx, commonpb.SegmentState_SegmentStateNone, loadInfo)
 	defer loader.unregister(infos...)
 
 	indexInfo := lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *querypb.SegmentLoadInfo {

diff --git a/internal/querynodev2/segments/segment_loader_test.go b/internal/querynodev2/segments/segment_loader_test.go
@@ -714,7 +714,7 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
 			return nil
 		})
 		suite.segmentManager.EXPECT().UpdateBy(mock.Anything, mock.Anything, mock.Anything).Return(0)
-		infos = suite.loader.prepare(SegmentTypeSealed, &querypb.SegmentLoadInfo{
+		infos = suite.loader.prepare(context.Background(), SegmentTypeSealed, &querypb.SegmentLoadInfo{
 			SegmentID:    suite.segmentID,
 			PartitionID:  suite.partitionID,
 			CollectionID: suite.collectionID,
@@ -742,7 +742,7 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
 
 			return nil
 		})
-		infos = suite.loader.prepare(SegmentTypeSealed, &querypb.SegmentLoadInfo{
+		infos = suite.loader.prepare(context.Background(), SegmentTypeSealed, &querypb.SegmentLoadInfo{
 			SegmentID:    suite.segmentID,
 			PartitionID:  suite.partitionID,
 			CollectionID: suite.collectionID,
@@ -760,7 +760,7 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
 		suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
 			return nil
 		})
-		suite.loader.prepare(SegmentTypeSealed, &querypb.SegmentLoadInfo{
+		suite.loader.prepare(context.Background(), SegmentTypeSealed, &querypb.SegmentLoadInfo{
 			SegmentID:    suite.segmentID,
 			PartitionID:  suite.partitionID,
 			CollectionID: suite.collectionID,