milvus/internal/querynodev2/segments/manager.go

// 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/zap"
	"golang.org/x/sync/singleflight"

	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
	"github.com/milvus-io/milvus/internal/querynodev2/segments/metricsutil"
	"github.com/milvus-io/milvus/pkg/v2/eventlog"
	"github.com/milvus-io/milvus/pkg/v2/log"
	"github.com/milvus-io/milvus/pkg/v2/metrics"
	"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v2/util/cache"
	"github.com/milvus-io/milvus/pkg/v2/util/lock"
	"github.com/milvus-io/milvus/pkg/v2/util/merr"
	"github.com/milvus-io/milvus/pkg/v2/util/metautil"
	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v2/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 := log.Ctx(context.Background()).With(
			zap.Int64("segmentID", segment.ID()),
			zap.String("type", segment.Type().String()),
			zap.Int64("segmentVersion", segment.Version()),
			zap.Int64("updateVersion", version),
		)
		for oldVersion := segment.Version(); oldVersion < version; {
			if segment.CASVersion(oldVersion, version) {
				return true
			}
		}
		log.Warn("segment version cannot go backwards, skip update")
		return false
	}
}

type Manager struct {
	Collection CollectionManager
	Segment    SegmentManager
	DiskCache  cache.Cache[int64, Segment]
	Loader     Loader
}

func NewManager() *Manager {
	diskCap := paramtable.Get().QueryNodeCfg.DiskCacheCapacityLimit.GetAsSize()

	segMgr := NewSegmentManager()
	sf := singleflight.Group{}
	manager := &Manager{
		Collection: NewCollectionManager(),
		Segment:    segMgr,
	}

	manager.DiskCache = cache.NewCacheBuilder[int64, Segment]().WithLazyScavenger(func(key int64) int64 {
		segment := segMgr.GetWithType(key, SegmentTypeSealed)
		if segment == nil {
			return 0
		}
		return int64(segment.ResourceUsageEstimate().DiskSize)
	}, diskCap).WithLoader(func(ctx context.Context, key int64) (Segment, error) {
		log := log.Ctx(ctx)
		log.Debug("cache missed segment", zap.Int64("segmentID", key))
		segment := segMgr.GetWithType(key, SegmentTypeSealed)
		if segment == nil {
			// the segment has been released, just ignore it
			log.Warn("segment is not found when loading", zap.Int64("segmentID", key))
			return nil, merr.ErrSegmentNotFound
		}
		info := segment.LoadInfo()
		_, err, _ := sf.Do(fmt.Sprint(segment.ID()), func() (nop interface{}, err error) {
			cacheLoadRecord := metricsutil.NewCacheLoadRecord(getSegmentMetricLabel(segment))
			cacheLoadRecord.WithBytes(segment.ResourceUsageEstimate().DiskSize)
			defer func() {
				cacheLoadRecord.Finish(err)
			}()

			collection := manager.Collection.Get(segment.Collection())
			if collection == nil {
				return nil, merr.WrapErrCollectionNotLoaded(segment.Collection(), "failed to load segment fields")
			}

			err = manager.Loader.LoadLazySegment(ctx, segment, info)
			return nil, err
		})
		if err != nil {
			log.Warn("cache sealed segment failed", zap.Error(err))
			return nil, err
		}
		return segment, nil
	}).WithFinalizer(func(ctx context.Context, key int64, segment Segment) error {
		log := log.Ctx(ctx)
		log.Debug("evict segment from cache", zap.Int64("segmentID", key))
		cacheEvictRecord := metricsutil.NewCacheEvictRecord(getSegmentMetricLabel(segment))
		cacheEvictRecord.WithBytes(segment.ResourceUsageEstimate().DiskSize)
		defer cacheEvictRecord.Finish(nil)
		segment.Release(ctx, WithReleaseScope(ReleaseScopeData))
		return nil
	}).WithReloader(func(ctx context.Context, key int64) (Segment, error) {
		log := log.Ctx(ctx)
		segment := segMgr.GetWithType(key, SegmentTypeSealed)
		if segment == nil {
			// the segment has been released, just ignore it
			log.Debug("segment is not found when reloading", zap.Int64("segmentID", key))
			return nil, merr.ErrSegmentNotFound
		}

		localSegment := segment.(*LocalSegment)
		err := manager.Loader.LoadIndex(ctx, localSegment, segment.LoadInfo(), segment.NeedUpdatedVersion())
		if err != nil {
			log.Warn("reload segment failed", zap.Int64("segmentID", key), zap.Error(err))
			return nil, merr.ErrSegmentLoadFailed
		}
		if err := localSegment.RemoveUnusedFieldFiles(); err != nil {
			log.Warn("remove unused field files failed", zap.Int64("segmentID", key), zap.Error(err))
			return nil, merr.ErrSegmentReduplicate
		}

		return segment, nil
	}).Build()

	segMgr.registerReleaseCallback(func(s Segment) {
		if s.Type() == SegmentTypeSealed {
			// !!! We cannot use ctx of request to call Remove,
			// Once context canceled, the segment will be leak in cache forever.
			// Because it has been cleaned from segment manager.
			manager.DiskCache.Remove(context.Background(), s.ID())
		}
	})

	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)

	// 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
}

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
}

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
		log.Warn("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
		log.Warn("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
}

// 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

	log := log.Ctx(ctx)
	for _, segment := range segments {
		oldSegment, ok := mgr.globalSegments.GetWithType(segment.ID(), segmentType)
		if ok {
			if oldSegment.Version() >= segment.Version() {
				log.Warn("Invalid segment distribution changed, skip it",
					zap.Int64("segmentID", segment.ID()),
					zap.Int64("oldVersion", oldSegment.Version()),
					zap.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(
			fmt.Sprint(paramtable.GetNodeID()),
			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
}

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)
		log.Ctx(ctx).Info("remove segment from cache", zap.Int64("segmentID", segment.ID()))
	}
	segment.Release(ctx)

	metrics.QueryNodeNumSegments.WithLabelValues(
		fmt.Sprint(paramtable.GetNodeID()),
		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())
	}
}