milvus/internal/proxy/task_policies.go
congqixia 4c4db1f556
Display all shard error if search/query fails (#20169)
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2022-10-31 19:09:39 +08:00

160 lines
4.9 KiB
Go

package proxy
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/types"
"go.uber.org/zap"
)
// type pickShardPolicy func(ctx context.Context, mgr *shardClientMgr, query func(UniqueID, types.QueryNode) error, leaders []nodeInfo) error
type pickShardPolicy func(context.Context, *shardClientMgr, func(context.Context, UniqueID, types.QueryNode, []string) error, map[string][]nodeInfo) error
var (
errBegin = errors.New("begin error")
errInvalidShardLeaders = errors.New("Invalid shard leader")
)
func updateShardsWithRoundRobin(shardsLeaders map[string][]nodeInfo) {
for channelID, leaders := range shardsLeaders {
if len(leaders) <= 1 {
continue
}
shardsLeaders[channelID] = append(leaders[1:], leaders[0])
}
}
// mergeErrSet merges all errors in ErrSet
func mergeErrSet(errSet map[string]error) error {
var builder strings.Builder
for channel, err := range errSet {
if err == nil {
continue
}
builder.WriteString(fmt.Sprintf("Channel: %s returns err: %s", channel, err.Error()))
}
return errors.New(builder.String())
}
// group dml shard leader with same nodeID
func groupShardleadersWithSameQueryNode(
ctx context.Context,
shard2leaders map[string][]nodeInfo,
nexts map[string]int, errSet map[string]error,
mgr *shardClientMgr) (map[int64][]string, map[int64]types.QueryNode, error) {
// check if all leaders were checked
for dml, idx := range nexts {
if idx >= len(shard2leaders[dml]) {
log.Ctx(ctx).Warn("no shard leaders were available",
zap.String("channel", dml),
zap.String("leaders", fmt.Sprintf("%v", shard2leaders[dml])))
if _, ok := errSet[dml]; ok {
return nil, nil, mergeErrSet(errSet) // return merged last error recorded
}
return nil, nil, fmt.Errorf("no available shard leader")
}
}
qnSet := make(map[int64]types.QueryNode)
node2dmls := make(map[int64][]string)
updates := make(map[string]int)
for dml, idx := range nexts {
updates[dml] = idx + 1
nodeInfo := shard2leaders[dml][idx]
if _, ok := qnSet[nodeInfo.nodeID]; !ok {
qn, err := mgr.GetClient(ctx, nodeInfo.nodeID)
if err != nil {
log.Ctx(ctx).Warn("failed to get shard leader", zap.Int64("nodeID", nodeInfo.nodeID), zap.Error(err))
// if get client failed, just record error and wait for next round to get client and do query
errSet[dml] = err
continue
}
qnSet[nodeInfo.nodeID] = qn
}
if _, ok := node2dmls[nodeInfo.nodeID]; !ok {
node2dmls[nodeInfo.nodeID] = make([]string, 0)
}
node2dmls[nodeInfo.nodeID] = append(node2dmls[nodeInfo.nodeID], dml)
}
// update idxes
for dml, idx := range updates {
nexts[dml] = idx
}
return node2dmls, qnSet, nil
}
// mergeRoundRobinPolicy first group shard leaders with same querynode, then do the query with multiple dml channels
// if request failed, it finds shard leader for failed dml channels, and again groups shard leaders and do the query
//
// Suppose qn0 is the shard leader for dml-channel0 and dml-channel1, if search for dml-channel0 succeeded, but
// failed for dml-channel1. In this case, an error returned from qn0, and next shard leaders for dml-channel0 and dml-channel1 will be
// retrieved and dml-channel0 therefore will again be searched.
//
// TODO: In this senario, qn0 should return a partial success results for dml-channel0, and only retrys for dml-channel1
func mergeRoundRobinPolicy(
ctx context.Context,
mgr *shardClientMgr,
query func(context.Context, UniqueID, types.QueryNode, []string) error,
dml2leaders map[string][]nodeInfo) error {
nexts := make(map[string]int)
errSet := make(map[string]error) // record err for dml channels
for dml := range dml2leaders {
nexts[dml] = 0
}
for len(nexts) > 0 {
node2dmls, nodeset, err := groupShardleadersWithSameQueryNode(ctx, dml2leaders, nexts, errSet, mgr)
if err != nil {
return err
}
wg := &sync.WaitGroup{}
mu := &sync.Mutex{}
wg.Add(len(node2dmls))
for nodeID, channels := range node2dmls {
nodeID := nodeID
channels := channels
qn := nodeset[nodeID]
go func() {
defer wg.Done()
if err := query(ctx, nodeID, qn, channels); err != nil {
log.Ctx(ctx).Warn("failed to do query with node", zap.Int64("nodeID", nodeID),
zap.Strings("dmlChannels", channels), zap.Error(err))
mu.Lock()
defer mu.Unlock()
for _, ch := range channels {
errSet[ch] = err
}
return
}
mu.Lock()
defer mu.Unlock()
for _, channel := range channels {
delete(nexts, channel)
delete(errSet, channel)
}
}()
}
wg.Wait()
if len(nexts) > 0 {
nextSet := make(map[string]int64)
for dml, idx := range nexts {
if idx >= len(dml2leaders[dml]) {
nextSet[dml] = -1
} else {
nextSet[dml] = dml2leaders[dml][idx].nodeID
}
}
log.Ctx(ctx).Warn("retry another query node with round robin", zap.Any("Nexts", nextSet))
}
}
return nil
}