milvus/internal/querycoordv2/meta/leader_view_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 meta

import (
	"sync"

	"github.com/samber/lo"

	"github.com/milvus-io/milvus/internal/proto/querypb"
)

type LeaderViewFilter interface {
	Match(*LeaderView) bool
	Node() (int64, bool)
	ChannelName() (string, bool)
}

type lvFilterFunc func(view *LeaderView) bool

func (f lvFilterFunc) Match(view *LeaderView) bool {
	return f(view)
}

func (f lvFilterFunc) Node() (int64, bool) {
	return -1, false
}

func (f lvFilterFunc) ChannelName() (string, bool) {
	return "", false
}

type lvChannelNameFilter string

func (f lvChannelNameFilter) Match(v *LeaderView) bool {
	return v.Channel == string(f)
}

func (f lvChannelNameFilter) Node() (int64, bool) {
	return -1, false
}

func (f lvChannelNameFilter) ChannelName() (string, bool) {
	return string(f), true
}

type lvNodeFilter int64

func (f lvNodeFilter) Match(v *LeaderView) bool {
	return v.ID == int64(f)
}

func (f lvNodeFilter) Node() (int64, bool) {
	return int64(f), true
}

func (f lvNodeFilter) ChannelName() (string, bool) {
	return "", false
}

func WithNodeID2LeaderView(nodeID int64) LeaderViewFilter {
	return lvNodeFilter(nodeID)
}

func WithChannelName2LeaderView(channelName string) LeaderViewFilter {
	return lvChannelNameFilter(channelName)
}

func WithCollectionID2LeaderView(collectionID int64) LeaderViewFilter {
	return lvFilterFunc(func(view *LeaderView) bool {
		return view.CollectionID == 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
}

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

type channelViews map[string]*LeaderView

type LeaderViewManager struct {
	rwmutex sync.RWMutex
	views   map[int64]channelViews // LeaderID -> Views (one per shard)
}

func NewLeaderViewManager() *LeaderViewManager {
	return &LeaderViewManager{
		views: make(map[int64]channelViews),
	}
}

// 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()
	mgr.views[leaderID] = make(channelViews, len(views))
	for _, view := range views {
		mgr.views[leaderID][view.Channel] = view
	}
}

func (mgr *LeaderViewManager) GetLeaderShardView(id int64, shard string) *LeaderView {
	mgr.rwmutex.RLock()
	defer mgr.rwmutex.RUnlock()

	return mgr.views[id][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 {
	otherFilters := make([]LeaderViewFilter, 0, len(filters))
	var nodeID int64
	var channelName string
	var hasNodeID, hasChannelName bool

	for _, filter := range filters {
		if node, ok := filter.Node(); ok {
			nodeID, hasNodeID = node, true
			continue
		}
		if channel, ok := filter.ChannelName(); ok {
			channelName, hasChannelName = channel, true
			continue
		}
		otherFilters = append(otherFilters, filter)
	}
	mergedFilter := func(view *LeaderView) bool {
		for _, filter := range otherFilters {
			if !filter.Match(view) {
				return false
			}
		}
		return true
	}

	var candidates []channelViews
	if hasNodeID {
		nodeView, ok := mgr.views[nodeID]
		if ok {
			candidates = append(candidates, nodeView)
		}
	} else {
		candidates = lo.Values(mgr.views)
	}

	var result []*LeaderView
	for _, candidate := range candidates {
		if hasChannelName {
			if view, ok := candidate[channelName]; ok && mergedFilter(view) {
				result = append(result, view)
			}
		} else {
			for _, view := range candidate {
				if mergedFilter(view) {
					result = append(result, view)
				}
			}
		}
	}
	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
	})
}