milvus/internal/querycoordv2/meta/target_manager.go

209 lines
5.6 KiB
Go
Raw Normal View History

// 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 meta
import (
"sync"
"time"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/util/funcutil"
"github.com/milvus-io/milvus/internal/util/tsoutil"
"go.uber.org/zap"
)
type TargetManager struct {
rwmutex sync.RWMutex
segments map[int64]*datapb.SegmentInfo
dmChannels map[string]*DmChannel
}
func NewTargetManager() *TargetManager {
return &TargetManager{
segments: make(map[int64]*datapb.SegmentInfo),
dmChannels: make(map[string]*DmChannel),
}
}
// RemoveCollection removes all channels and segments in the given collection
func (mgr *TargetManager) RemoveCollection(collectionID int64) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
log.Info("remove collection from targets")
for _, segment := range mgr.segments {
if segment.CollectionID == collectionID {
mgr.removeSegment(segment.GetID())
}
}
for _, dmChannel := range mgr.dmChannels {
if dmChannel.CollectionID == collectionID {
mgr.removeDmChannel(dmChannel.GetChannelName())
}
}
}
// RemovePartition removes all segment in the given partition,
// NOTE: this doesn't remove any channel even the given one is the only partition
func (mgr *TargetManager) RemovePartition(partitionID int64) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
log.Info("remove partition from targets",
zap.Int64("partitionID", partitionID))
for _, segment := range mgr.segments {
if segment.GetPartitionID() == partitionID {
mgr.removeSegment(segment.GetID())
}
}
}
func (mgr *TargetManager) RemoveSegment(segmentID int64) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
delete(mgr.segments, segmentID)
}
func (mgr *TargetManager) removeSegment(segmentID int64) {
delete(mgr.segments, segmentID)
log.Info("segment removed from targets", zap.Int64("segment", segmentID))
}
// AddSegment adds segment into target set,
// requires CollectionID, PartitionID, InsertChannel, SegmentID are set
func (mgr *TargetManager) AddSegment(segments ...*datapb.SegmentInfo) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
mgr.addSegment(segments...)
}
func (mgr *TargetManager) addSegment(segments ...*datapb.SegmentInfo) {
for _, segment := range segments {
log.Info("add segment into targets",
zap.Int64("segmentID", segment.GetID()),
zap.Int64("collectionID", segment.GetCollectionID()),
)
mgr.segments[segment.GetID()] = segment
}
}
func (mgr *TargetManager) ContainSegment(id int64) bool {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
return mgr.containSegment(id)
}
func (mgr *TargetManager) containSegment(id int64) bool {
_, ok := mgr.segments[id]
return ok
}
func (mgr *TargetManager) GetSegmentsByCollection(collection int64, partitions ...int64) []*datapb.SegmentInfo {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
segments := make([]*datapb.SegmentInfo, 0)
for _, segment := range mgr.segments {
if segment.CollectionID == collection &&
(len(partitions) == 0 || funcutil.SliceContain(partitions, segment.PartitionID)) {
segments = append(segments, segment)
}
}
return segments
}
func (mgr *TargetManager) HandoffSegment(dest *datapb.SegmentInfo, sources ...int64) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
// add dest to target
dest.CompactionFrom = sources
mgr.addSegment(dest)
}
// AddDmChannel adds a channel into target set,
// requires CollectionID, ChannelName are set
func (mgr *TargetManager) AddDmChannel(channels ...*DmChannel) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
for _, channel := range channels {
ts := channel.GetSeekPosition().GetTimestamp()
log.Info("add channel into targets",
zap.String("channel", channel.GetChannelName()),
zap.Uint64("checkpoint", ts),
zap.Duration("sinceCheckpoint", time.Since(tsoutil.PhysicalTime(ts))),
)
mgr.dmChannels[channel.ChannelName] = channel
}
}
func (mgr *TargetManager) GetDmChannel(channel string) *DmChannel {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
for _, ch := range mgr.dmChannels {
if ch.ChannelName == channel {
return ch
}
}
return nil
}
func (mgr *TargetManager) ContainDmChannel(channel string) bool {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
_, ok := mgr.dmChannels[channel]
return ok
}
func (mgr *TargetManager) removeDmChannel(channel string) {
delete(mgr.dmChannels, channel)
log.Info("remove channel from targets", zap.String("channel", channel))
}
func (mgr *TargetManager) GetDmChannelsByCollection(collectionID int64) []*DmChannel {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
channels := make([]*DmChannel, 0)
for _, channel := range mgr.dmChannels {
if channel.GetCollectionID() == collectionID {
channels = append(channels, channel)
}
}
return channels
}
func (mgr *TargetManager) GetSegment(id int64) *datapb.SegmentInfo {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
for _, s := range mgr.segments {
if s.GetID() == id {
return s
}
}
return nil
}