milvus/internal/querycoordv2/meta/leader_view_manager.go

299 lines
7.5 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"
"github.com/samber/lo"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/pkg/util/typeutil"
)
type lvCriterion struct {
nodeID int64
channelName string
collectionID int64
hasOtherFilter bool
}
type LeaderViewFilter interface {
Match(*LeaderView) bool
AddFilter(*lvCriterion)
}
type lvFilterFunc func(view *LeaderView) bool
func (f lvFilterFunc) Match(view *LeaderView) bool {
return f(view)
}
func (f lvFilterFunc) AddFilter(c *lvCriterion) {
c.hasOtherFilter = true
}
type lvChannelNameFilter string
func (f lvChannelNameFilter) Match(v *LeaderView) bool {
return v.Channel == string(f)
}
func (f lvChannelNameFilter) AddFilter(c *lvCriterion) {
c.channelName = string(f)
}
type lvNodeFilter int64
func (f lvNodeFilter) Match(v *LeaderView) bool {
return v.ID == int64(f)
}
func (f lvNodeFilter) AddFilter(c *lvCriterion) {
c.nodeID = int64(f)
}
type lvCollectionFilter int64
func (f lvCollectionFilter) Match(v *LeaderView) bool {
return v.CollectionID == int64(f)
}
func (f lvCollectionFilter) AddFilter(c *lvCriterion) {
c.collectionID = int64(f)
}
func WithNodeID2LeaderView(nodeID int64) LeaderViewFilter {
return lvNodeFilter(nodeID)
}
func WithChannelName2LeaderView(channelName string) LeaderViewFilter {
return lvChannelNameFilter(channelName)
}
func WithCollectionID2LeaderView(collectionID int64) LeaderViewFilter {
return lvCollectionFilter(collectionID)
}
func WithReplica2LeaderView(replica *Replica) LeaderViewFilter {
return lvFilterFunc(func(view *LeaderView) bool {
if replica == nil {
return false
}
return replica.GetCollectionID() == view.CollectionID && replica.Contains(view.ID)
})
}
func WithSegment2LeaderView(segmentID int64, isGrowing bool) LeaderViewFilter {
return lvFilterFunc(func(view *LeaderView) bool {
if isGrowing {
_, ok := view.GrowingSegments[segmentID]
return ok
}
_, ok := view.Segments[segmentID]
return ok
})
}
type LeaderView struct {
ID int64
CollectionID int64
Channel string
Version int64
Segments map[int64]*querypb.SegmentDist
GrowingSegments map[int64]*Segment
TargetVersion int64
NumOfGrowingRows int64
PartitionStatsVersions map[int64]int64
}
func (view *LeaderView) Clone() *LeaderView {
segments := make(map[int64]*querypb.SegmentDist)
for k, v := range view.Segments {
segments[k] = v
}
growings := make(map[int64]*Segment)
for k, v := range view.GrowingSegments {
growings[k] = v
}
return &LeaderView{
ID: view.ID,
CollectionID: view.CollectionID,
Channel: view.Channel,
Version: view.Version,
Segments: segments,
GrowingSegments: growings,
TargetVersion: view.TargetVersion,
NumOfGrowingRows: view.NumOfGrowingRows,
PartitionStatsVersions: view.PartitionStatsVersions,
}
}
type nodeViews struct {
views []*LeaderView
// channel name => LeaderView
channelView map[string]*LeaderView
// collection id => leader views
collectionViews map[int64][]*LeaderView
}
func (v nodeViews) Filter(criterion *lvCriterion, filters ...LeaderViewFilter) []*LeaderView {
mergedFilter := func(view *LeaderView) bool {
for _, filter := range filters {
if !filter.Match(view) {
return false
}
}
return true
}
var views []*LeaderView
switch {
case criterion.channelName != "":
if view, ok := v.channelView[criterion.channelName]; ok {
views = append(views, view)
}
case criterion.collectionID != 0:
views = v.collectionViews[criterion.collectionID]
default:
views = v.views
}
if criterion.hasOtherFilter {
views = lo.Filter(views, func(view *LeaderView, _ int) bool {
return mergedFilter(view)
})
}
return views
}
func composeNodeViews(views ...*LeaderView) nodeViews {
return nodeViews{
views: views,
channelView: lo.SliceToMap(views, func(view *LeaderView) (string, *LeaderView) {
return view.Channel, view
}),
collectionViews: lo.GroupBy(views, func(view *LeaderView) int64 {
return view.CollectionID
}),
}
}
type NotifyDelegatorChanges = func(collectionID ...int64)
type LeaderViewManager struct {
rwmutex sync.RWMutex
views map[int64]nodeViews // LeaderID -> Views (one per shard)
notifyFunc NotifyDelegatorChanges
}
func NewLeaderViewManager() *LeaderViewManager {
return &LeaderViewManager{
views: make(map[int64]nodeViews),
}
}
func (mgr *LeaderViewManager) SetNotifyFunc(notifyFunc NotifyDelegatorChanges) {
mgr.notifyFunc = notifyFunc
}
// Update updates the leader's views, all views have to be with the same leader ID
func (mgr *LeaderViewManager) Update(leaderID int64, views ...*LeaderView) {
mgr.rwmutex.Lock()
defer mgr.rwmutex.Unlock()
oldViews := make(map[string]*LeaderView, 0)
if _, ok := mgr.views[leaderID]; ok {
oldViews = mgr.views[leaderID].channelView
}
newViews := lo.SliceToMap(views, func(v *LeaderView) (string, *LeaderView) {
return v.Channel, v
})
// update leader views
mgr.views[leaderID] = composeNodeViews(views...)
// compute leader location change, find it's correspond collection
if mgr.notifyFunc != nil {
viewChanges := typeutil.NewUniqueSet()
for channel, oldView := range oldViews {
// if channel released from current node
if _, ok := newViews[channel]; !ok {
viewChanges.Insert(oldView.CollectionID)
}
}
for channel, newView := range newViews {
// if channel loaded to current node
if _, ok := oldViews[channel]; !ok {
viewChanges.Insert(newView.CollectionID)
}
}
mgr.notifyFunc(viewChanges.Collect()...)
}
}
func (mgr *LeaderViewManager) GetLeaderShardView(id int64, shard string) *LeaderView {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
return mgr.views[id].channelView[shard]
}
func (mgr *LeaderViewManager) GetByFilter(filters ...LeaderViewFilter) []*LeaderView {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
return mgr.getByFilter(filters...)
}
func (mgr *LeaderViewManager) getByFilter(filters ...LeaderViewFilter) []*LeaderView {
criterion := &lvCriterion{}
for _, filter := range filters {
filter.AddFilter(criterion)
}
var candidates []nodeViews
if criterion.nodeID > 0 {
nodeView, ok := mgr.views[criterion.nodeID]
if ok {
candidates = append(candidates, nodeView)
}
} else {
candidates = lo.Values(mgr.views)
}
var result []*LeaderView
for _, candidate := range candidates {
result = append(result, candidate.Filter(criterion, filters...)...)
}
return result
}
func (mgr *LeaderViewManager) GetLatestShardLeaderByFilter(filters ...LeaderViewFilter) *LeaderView {
mgr.rwmutex.RLock()
defer mgr.rwmutex.RUnlock()
views := mgr.getByFilter(filters...)
return lo.MaxBy(views, func(v1, v2 *LeaderView) bool {
return v1.Version > v2.Version
})
}