diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go index 660537ce62..4f625a3e69 100644 --- 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 index d7901b4ebc..122fbd0e4e 100644 --- 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,