milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: Add ctx trace for segment load `prepare` (#31203) 

Add ctx paramter to `prepare` and add trace id in underline logs.

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/31170/head
congqixia 2024-03-13 10:01:07 +08:00 committed by GitHub
parent 5b51c20293
commit 1d96239137
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 14 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
internal/querynodev2/segments/segment_loader.go
View File

@ -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 {

6
internal/querynodev2/segments/segment_loader_test.go
View File

@ -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,