// Licensed to the LF AI & Data foundation under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package segments /* #cgo pkg-config: milvus_core #include "segcore/collection_c.h" #include "segcore/segment_c.h" */ import "C" import ( "context" "fmt" "sync" "go.uber.org/atomic" "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" "github.com/milvus-io/milvus/pkg/v3/eventlog" "github.com/milvus-io/milvus/pkg/v3/metrics" "github.com/milvus-io/milvus/pkg/v3/mlog" "github.com/milvus-io/milvus/pkg/v3/proto/datapb" "github.com/milvus-io/milvus/pkg/v3/proto/querypb" "github.com/milvus-io/milvus/pkg/v3/util/lock" "github.com/milvus-io/milvus/pkg/v3/util/merr" "github.com/milvus-io/milvus/pkg/v3/util/metautil" "github.com/milvus-io/milvus/pkg/v3/util/paramtable" "github.com/milvus-io/milvus/pkg/v3/util/typeutil" ) // TODO maybe move to manager and change segment constructor var channelMapper = metautil.NewDynChannelMapper() type SegmentAction func(segment Segment) bool func IncreaseVersion(version int64) SegmentAction { return func(segment Segment) bool { log := mlog.With( mlog.FieldSegmentID(segment.ID()), mlog.String("type", segment.Type().String()), mlog.Int64("segmentVersion", segment.Version()), mlog.Int64("updateVersion", version), ) for oldVersion := segment.Version(); oldVersion < version; { if segment.CASVersion(oldVersion, version) { return true } } log.Warn(context.TODO(), "segment version cannot go backwards, skip update") return false } } type Manager struct { Collection CollectionManager Segment SegmentManager Loader Loader } func NewManager() *Manager { segMgr := NewSegmentManager() manager := &Manager{ Collection: NewCollectionManager(), Segment: segMgr, } return manager } func (mgr *Manager) SetLoader(loader Loader) { mgr.Loader = loader } type SegmentManager interface { // Put puts the given segments in, // and increases the ref count of the corresponding collection, // dup segments will not increase the ref count Put(ctx context.Context, segmentType SegmentType, segments ...Segment) UpdateBy(action SegmentAction, filters ...SegmentFilter) int Get(segmentID typeutil.UniqueID) Segment GetWithType(segmentID typeutil.UniqueID, typ SegmentType) Segment GetBy(filters ...SegmentFilter) []Segment // Get segments and acquire the read locks GetAndPinBy(filters ...SegmentFilter) ([]Segment, error) GetAndPin(segments []int64, filters ...SegmentFilter) ([]Segment, error) Unpin(segments []Segment) GetSealed(segmentID typeutil.UniqueID) Segment GetGrowing(segmentID typeutil.UniqueID) Segment Empty() bool // Remove removes the given segment, // and decreases the ref count of the corresponding collection, // will not decrease the ref count if the given segment not exists Remove(ctx context.Context, segmentID typeutil.UniqueID, scope querypb.DataScope) (int, int) RemoveBy(ctx context.Context, filters ...SegmentFilter) (int, int) Clear(ctx context.Context) // ReleaseDetached completes the release of a segment previously taken out of // the active maps via DetachStreaming. It runs the same bookkeeping as the // normal release path (release callback, segment teardown, metric Dec and // clearing the on-releasing set). ReleaseDetached(ctx context.Context, segment Segment) // Deprecated: quick fix critical issue: #30857 // TODO: All Segment assigned to querynode should be managed by SegmentManager, including loading or releasing to perform a transaction. Exist(segmentID typeutil.UniqueID, typ SegmentType) bool AddLogicalResource(usage ResourceUsage) SubLogicalResource(usage ResourceUsage) GetLogicalResource() ResourceUsage AddLoadedBinlogSize(size int64) SubLoadedBinlogSize(size int64) GetLoadedBinlogSize() int64 } var _ SegmentManager = (*segmentManager)(nil) type secondarySegmentIndex struct { keyLock *lock.KeyLock[metautil.Channel] shardSegments *typeutil.ConcurrentMap[metautil.Channel, segments] } func newSecondarySegmentIndex() secondarySegmentIndex { return secondarySegmentIndex{ keyLock: lock.NewKeyLock[metautil.Channel](), shardSegments: typeutil.NewConcurrentMap[metautil.Channel, segments](), } } func (si secondarySegmentIndex) Put(ctx context.Context, segmentType SegmentType, segment Segment) { shard := segment.Shard() si.keyLock.Lock(shard) defer si.keyLock.Unlock(shard) segments, _ := si.shardSegments.GetOrInsert(shard, newSegments()) segments.Put(ctx, segmentType, segment) } func (si secondarySegmentIndex) Remove(s Segment) { shard := s.Shard() si.keyLock.Lock(shard) defer si.keyLock.Unlock(shard) segments, ok := si.shardSegments.Get(shard) if !ok { return } segments.Remove(s) if segments.Empty() { si.shardSegments.Remove(shard) } } type segments struct { growingSegments *typeutil.ConcurrentMap[typeutil.UniqueID, Segment] sealedSegments *typeutil.ConcurrentMap[typeutil.UniqueID, Segment] } func (segments segments) Put(_ context.Context, segmentType SegmentType, segment Segment) { switch segmentType { case SegmentTypeGrowing: segments.growingSegments.Insert(segment.ID(), segment) case SegmentTypeSealed: segments.sealedSegments.Insert(segment.ID(), segment) } } func (segments segments) Get(segmentID int64) (growing Segment, sealed Segment) { growing, _ = segments.growingSegments.Get(segmentID) sealed, _ = segments.sealedSegments.Get(segmentID) return growing, sealed } func (segments segments) GetWithType(segmentID int64, segmentType SegmentType) (Segment, bool) { // var targetMap map[int64]Segment var segment Segment var ok bool switch segmentType { case SegmentTypeGrowing: segment, ok = segments.growingSegments.Get(segmentID) case SegmentTypeSealed: segment, ok = segments.sealedSegments.Get(segmentID) } return segment, ok } func (segments segments) RemoveWithType(segmentID int64, segmentType SegmentType) (Segment, bool) { var segment Segment var ok bool switch segmentType { case SegmentTypeGrowing: segment, ok = segments.growingSegments.GetAndRemove(segmentID) case SegmentTypeSealed: segment, ok = segments.sealedSegments.GetAndRemove(segmentID) } return segment, ok } func (segments segments) Remove(segment Segment) { switch segment.Type() { case SegmentTypeGrowing: segments.growingSegments.Remove(segment.ID()) case SegmentTypeSealed: segments.sealedSegments.Remove(segment.ID()) } } func (segments segments) RangeWithFilter(criterion *segmentCriterion, process func(id int64, segType SegmentType, segment Segment) bool) { if criterion.segmentIDs != nil { for id := range criterion.segmentIDs { // var segment Segment // var ok bool var segs []Segment if criterion.segmentType == commonpb.SegmentState_SegmentStateNone { growing, sealed := segments.Get(id) if growing != nil { segs = append(segs, growing) } if sealed != nil { segs = append(segs, sealed) } } else { segment, ok := segments.GetWithType(id, criterion.segmentType) if ok { segs = append(segs, segment) } } for _, segment := range segs { if criterion.Match(segment) { if !process(id, segment.Type(), segment) { return } } } } return } var candidates []*typeutil.ConcurrentMap[typeutil.UniqueID, Segment] switch criterion.segmentType { case SegmentTypeGrowing: candidates = []*typeutil.ConcurrentMap[typeutil.UniqueID, Segment]{segments.growingSegments} case SegmentTypeSealed: candidates = []*typeutil.ConcurrentMap[typeutil.UniqueID, Segment]{segments.sealedSegments} default: candidates = []*typeutil.ConcurrentMap[typeutil.UniqueID, Segment]{segments.growingSegments, segments.sealedSegments} } for _, candidate := range candidates { candidate.Range(func(id typeutil.UniqueID, segment Segment) bool { if criterion.Match(segment) { if !process(id, segment.Type(), segment) { return false } } return true }) } } func (segments segments) Empty() bool { return segments.growingSegments.Len() == 0 && segments.sealedSegments.Len() == 0 } func newSegments() segments { return segments{ growingSegments: typeutil.NewConcurrentMap[typeutil.UniqueID, Segment](), sealedSegments: typeutil.NewConcurrentMap[typeutil.UniqueID, Segment](), } } // Manager manages all collections and segments type segmentManager struct { globalSegments segments secondaryIndex secondarySegmentIndex // releaseCallback is the callback function when a segment is released. releaseCallback func(s Segment) growingOnReleasingSegments *typeutil.ConcurrentSet[int64] sealedOnReleasingSegments *typeutil.ConcurrentSet[int64] // logicalResource is the logical resource usage for all loaded segments of this querynode segment manager, // which is to avoid memory and disk pressure when loading too many segments after eviction is enabled. // only MemorySize and DiskSize are used, other fields are ignored. logicalResource ResourceUsage logicalResourceLock sync.Mutex // loadedBinlogSize stats the total binlog size of all loaded segments of this querynode. loadedBinlogSize atomic.Int64 } func NewSegmentManager() *segmentManager { return &segmentManager{ globalSegments: newSegments(), secondaryIndex: newSecondarySegmentIndex(), growingOnReleasingSegments: typeutil.NewConcurrentSet[int64](), sealedOnReleasingSegments: typeutil.NewConcurrentSet[int64](), logicalResourceLock: sync.Mutex{}, } } func (mgr *segmentManager) AddLogicalResource(usage ResourceUsage) { mgr.logicalResourceLock.Lock() defer mgr.logicalResourceLock.Unlock() mgr.logicalResource.MemorySize += usage.MemorySize mgr.logicalResource.DiskSize += usage.DiskSize } func (mgr *segmentManager) SubLogicalResource(usage ResourceUsage) { mgr.logicalResourceLock.Lock() defer mgr.logicalResourceLock.Unlock() // avoid overflow of memory and disk size if mgr.logicalResource.MemorySize < usage.MemorySize { mgr.logicalResource.MemorySize = 0 mlog.Warn(context.TODO(), "Logical memory size would be negative, setting to 0") } else { mgr.logicalResource.MemorySize -= usage.MemorySize } if mgr.logicalResource.DiskSize < usage.DiskSize { mgr.logicalResource.DiskSize = 0 mlog.Warn(context.TODO(), "Logical disk size would be negative, setting to 0") } else { mgr.logicalResource.DiskSize -= usage.DiskSize } } func (mgr *segmentManager) GetLogicalResource() ResourceUsage { mgr.logicalResourceLock.Lock() defer mgr.logicalResourceLock.Unlock() return mgr.logicalResource } func (mgr *segmentManager) AddLoadedBinlogSize(size int64) { mgr.loadedBinlogSize.Add(size) } func (mgr *segmentManager) SubLoadedBinlogSize(size int64) { // Clamp to zero to avoid negative values on concurrent or duplicate subtractions for { current := mgr.loadedBinlogSize.Load() newVal := current - size if newVal < 0 { newVal = 0 } if mgr.loadedBinlogSize.CompareAndSwap(current, newVal) { if current < size { mlog.Warn(context.TODO(), "Loaded binlog size subtraction exceeds current value, clamped to 0", mlog.Int64("current", current), mlog.Int64("subtracted", size)) } return } // retry on CompareAndSwap failure } } func (mgr *segmentManager) GetLoadedBinlogSize() int64 { current := mgr.loadedBinlogSize.Load() if current < 0 { mlog.Warn(context.TODO(), "Loaded binlog size is negative, returning 0", mlog.Int64("current", current)) return 0 } return current } // put is the internal put method updating both global segments and secondary index. func (mgr *segmentManager) put(ctx context.Context, segmentType SegmentType, segment Segment) { mgr.globalSegments.Put(ctx, segmentType, segment) mgr.secondaryIndex.Put(ctx, segmentType, segment) } func (mgr *segmentManager) Put(ctx context.Context, segmentType SegmentType, segments ...Segment) { var replacedSegment []Segment for _, segment := range segments { oldSegment, ok := mgr.globalSegments.GetWithType(segment.ID(), segmentType) if ok { if oldSegment.Version() >= segment.Version() { mlog.Warn(ctx, "Invalid segment distribution changed, skip it", mlog.FieldSegmentID(segment.ID()), mlog.Int64("oldVersion", oldSegment.Version()), mlog.Int64("newVersion", segment.Version()), ) // delete redundant segment segment.Release(ctx) continue } replacedSegment = append(replacedSegment, oldSegment) } mgr.put(ctx, segmentType, segment) eventlog.Record(eventlog.NewRawEvt(eventlog.Level_Info, fmt.Sprintf("Segment %d[%d] loaded", segment.ID(), segment.Collection()))) metrics.QueryNodeNumSegments.WithLabelValues( paramtable.GetStringNodeID(), fmt.Sprint(segment.Collection()), segment.Type().String(), segment.Level().String(), ).Inc() } // release replaced segment if len(replacedSegment) > 0 { go func() { for _, segment := range replacedSegment { mgr.release(ctx, segment) } }() } } func (mgr *segmentManager) UpdateBy(action SegmentAction, filters ...SegmentFilter) int { updated := 0 mgr.rangeWithFilter(func(_ int64, _ SegmentType, segment Segment) bool { if action(segment) { updated++ } return true }, filters...) return updated } // Deprecated: // TODO: All Segment assigned to querynode should be managed by SegmentManager, including loading or releasing to perform a transaction. func (mgr *segmentManager) Exist(segmentID typeutil.UniqueID, typ SegmentType) bool { _, ok := mgr.globalSegments.GetWithType(segmentID, typ) if ok { return true } switch typ { case SegmentTypeGrowing: if mgr.growingOnReleasingSegments.Contain(segmentID) { return true } case SegmentTypeSealed: if mgr.sealedOnReleasingSegments.Contain(segmentID) { return true } } return false } func (mgr *segmentManager) Get(segmentID typeutil.UniqueID) Segment { growing, sealed := mgr.globalSegments.Get(segmentID) if growing != nil { return growing } return sealed } func (mgr *segmentManager) GetWithType(segmentID typeutil.UniqueID, typ SegmentType) Segment { segment, _ := mgr.globalSegments.GetWithType(segmentID, typ) return segment } func (mgr *segmentManager) GetBy(filters ...SegmentFilter) []Segment { var ret []Segment mgr.rangeWithFilter(func(id int64, _ SegmentType, segment Segment) bool { ret = append(ret, segment) return true }, filters...) return ret } func (mgr *segmentManager) GetAndPinBy(filters ...SegmentFilter) ([]Segment, error) { var ret []Segment var err error defer func() { if err != nil { for _, segment := range ret { segment.Unpin() } ret = nil } }() mgr.rangeWithFilter(func(id int64, _ SegmentType, segment Segment) bool { if segment.Level() == datapb.SegmentLevel_L0 { return true } err = segment.PinIfNotReleased() if err != nil { return false } ret = append(ret, segment) return true }, filters...) return ret, err } func (mgr *segmentManager) GetAndPin(segments []int64, filters ...SegmentFilter) ([]Segment, error) { lockedSegments := make([]Segment, 0, len(segments)) var err error defer func() { if err != nil { for _, segment := range lockedSegments { segment.Unpin() } lockedSegments = nil } }() criterion := &segmentCriterion{} for _, filter := range filters { filter.AddFilter(criterion) } for _, id := range segments { var segments []Segment if criterion.segmentType == commonpb.SegmentState_SegmentStateNone { growing, sealed := mgr.globalSegments.Get(id) if growing == nil && sealed == nil { err = merr.WrapErrSegmentNotLoaded(id, "segment not found") return nil, err } segments = []Segment{growing, sealed} } else { segment, ok := mgr.globalSegments.GetWithType(id, criterion.segmentType) if !ok { err = merr.WrapErrSegmentNotLoaded(id, "segment not found") return nil, err } segments = []Segment{segment} } for _, segment := range segments { if segment == nil { continue } // L0 Segment should not be queryable. if segment.Level() == datapb.SegmentLevel_L0 { continue } if !filter(segment, filters...) { continue } err = segment.PinIfNotReleased() if err != nil { return nil, err } lockedSegments = append(lockedSegments, segment) } } return lockedSegments, nil } func (mgr *segmentManager) Unpin(segments []Segment) { for _, segment := range segments { segment.Unpin() } } func (mgr *segmentManager) rangeWithFilter(process func(id int64, segType SegmentType, segment Segment) bool, filters ...SegmentFilter) { criterion := &segmentCriterion{} for _, filter := range filters { filter.AddFilter(criterion) } target := mgr.globalSegments var ok bool if !criterion.channel.IsZero() { target, ok = mgr.secondaryIndex.shardSegments.Get(criterion.channel) if !ok { return } } target.RangeWithFilter(criterion, process) } func filter(segment Segment, filters ...SegmentFilter) bool { for _, filter := range filters { if !filter.Match(segment) { return false } } return true } func (mgr *segmentManager) GetSealed(segmentID typeutil.UniqueID) Segment { return mgr.GetWithType(segmentID, SegmentTypeSealed) } func (mgr *segmentManager) GetGrowing(segmentID typeutil.UniqueID) Segment { return mgr.GetWithType(segmentID, SegmentTypeGrowing) } func (mgr *segmentManager) Empty() bool { return mgr.globalSegments.growingSegments.Len()+mgr.globalSegments.sealedSegments.Len() == 0 } // returns true if the segment exists, // false otherwise func (mgr *segmentManager) Remove(ctx context.Context, segmentID typeutil.UniqueID, scope querypb.DataScope) (int, int) { var removeGrowing, removeSealed int var growing, sealed Segment switch scope { case querypb.DataScope_Streaming: growing = mgr.removeSegmentWithType(SegmentTypeGrowing, segmentID) if growing != nil { removeGrowing = 1 } case querypb.DataScope_Historical: sealed = mgr.removeSegmentWithType(SegmentTypeSealed, segmentID) if sealed != nil { removeSealed = 1 } case querypb.DataScope_All: growing = mgr.removeSegmentWithType(SegmentTypeGrowing, segmentID) if growing != nil { removeGrowing = 1 } sealed = mgr.removeSegmentWithType(SegmentTypeSealed, segmentID) if sealed != nil { removeSealed = 1 } } if growing != nil { mgr.release(ctx, growing) } if sealed != nil { mgr.release(ctx, sealed) } return removeGrowing, removeSealed } // DetachStreaming removes the given growing segment from the active maps // WITHOUT releasing it. The segment is staged into the on-releasing set so // Exist() keeps reporting it while growing-source flush handoff keeps the // segment alive. The owner MUST call ReleaseDetached once it is done, // otherwise the on-releasing set entry, the segment gauge and the release // callback are never reconciled. func (mgr *segmentManager) DetachStreaming(ctx context.Context, segmentID typeutil.UniqueID) int { removeGrowing := 0 if mgr.removeSegmentWithType(SegmentTypeGrowing, segmentID) != nil { removeGrowing = 1 } mlog.Info(ctx, "detached segment from active segment manager", mlog.Int64("segmentID", segmentID), mlog.Int("growingCount", removeGrowing)) return removeGrowing } // ReleaseDetached completes the release of a segment previously taken out of the // active maps via DetachStreaming. DetachStreaming deliberately skips release() // so an out-of-band owner (growing-source flush handoff) can keep the segment // alive; that owner MUST call ReleaseDetached when it is done. Otherwise the // on-releasing set entry, the segment gauge (QueryNodeNumSegments) and the // release callback are never reconciled, and Exist() would keep returning true // for the segmentID. func (mgr *segmentManager) ReleaseDetached(ctx context.Context, segment Segment) { if segment == nil { return } mgr.release(ctx, segment) } func (mgr *segmentManager) removeSegmentWithType(typ SegmentType, segmentID typeutil.UniqueID) Segment { segment, ok := mgr.globalSegments.RemoveWithType(segmentID, typ) if !ok { return nil } switch typ { case SegmentTypeGrowing: mgr.growingOnReleasingSegments.Insert(segmentID) case SegmentTypeSealed: mgr.sealedOnReleasingSegments.Insert(segmentID) } mgr.secondaryIndex.Remove(segment) return segment } func (mgr *segmentManager) RemoveBy(ctx context.Context, filters ...SegmentFilter) (int, int) { var removeSegments []Segment var removeGrowing, removeSealed int mgr.rangeWithFilter(func(id int64, segType SegmentType, segment Segment) bool { s := mgr.removeSegmentWithType(segType, id) if s != nil { removeSegments = append(removeSegments, s) switch segType { case SegmentTypeGrowing: removeGrowing++ case SegmentTypeSealed: removeSealed++ } } return true }, filters...) for _, s := range removeSegments { mgr.release(ctx, s) } return removeGrowing, removeSealed } func (mgr *segmentManager) Clear(ctx context.Context) { mgr.globalSegments.growingSegments.Range(func(id typeutil.UniqueID, _ Segment) bool { mgr.growingOnReleasingSegments.Insert(id) return true }) growingWaitForRelease := mgr.globalSegments.growingSegments mgr.globalSegments.sealedSegments.Range(func(id typeutil.UniqueID, _ Segment) bool { mgr.sealedOnReleasingSegments.Insert(id) return true }) sealedWaitForRelease := mgr.globalSegments.sealedSegments mgr.globalSegments = newSegments() mgr.secondaryIndex = newSecondarySegmentIndex() growingWaitForRelease.Range(func(_ typeutil.UniqueID, segment Segment) bool { mgr.release(ctx, segment) return true }) sealedWaitForRelease.Range(func(_ typeutil.UniqueID, segment Segment) bool { mgr.release(ctx, segment) return true }) } // registerReleaseCallback registers the callback function when a segment is released. // TODO: bad implementation for keep consistency with DiskCache, need to be refactor. func (mgr *segmentManager) registerReleaseCallback(callback func(s Segment)) { mgr.releaseCallback = callback } func (mgr *segmentManager) release(ctx context.Context, segment Segment) { if mgr.releaseCallback != nil { mgr.releaseCallback(segment) mlog.Info(ctx, "remove segment from cache", mlog.FieldSegmentID(segment.ID())) } segment.Release(ctx) metrics.QueryNodeNumSegments.WithLabelValues( paramtable.GetStringNodeID(), fmt.Sprint(segment.Collection()), segment.Type().String(), segment.Level().String(), ).Dec() switch segment.Type() { case SegmentTypeGrowing: mgr.growingOnReleasingSegments.Remove(segment.ID()) case SegmentTypeSealed: mgr.sealedOnReleasingSegments.Remove(segment.ID()) } }