mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-04 21:09:06 +08:00
107 lines
2.5 KiB
Go
107 lines
2.5 KiB
Go
package proxy
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
|
|
qnClient "github.com/milvus-io/milvus/internal/distributed/querynode/client"
|
|
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type getQueryNodePolicy func(context.Context, string) (types.QueryNode, error)
|
|
|
|
type pickShardPolicy func(ctx context.Context, policy getQueryNodePolicy, query func(UniqueID, types.QueryNode) error, leaders []queryNode) error
|
|
|
|
// TODO add another policy to enbale the use of cache
|
|
// defaultGetQueryNodePolicy creates QueryNode client for every address everytime
|
|
func defaultGetQueryNodePolicy(ctx context.Context, address string) (types.QueryNode, error) {
|
|
qn, err := qnClient.NewClient(ctx, address)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := qn.Init(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := qn.Start(); err != nil {
|
|
return nil, err
|
|
}
|
|
return qn, nil
|
|
}
|
|
|
|
var (
|
|
errBegin = errors.New("begin error")
|
|
errInvalidShardLeaders = errors.New("Invalid shard leader")
|
|
)
|
|
|
|
type queryNode struct {
|
|
nodeID UniqueID
|
|
address string
|
|
}
|
|
|
|
func (q queryNode) String() string {
|
|
return fmt.Sprintf("<NodeID: %d>", q.nodeID)
|
|
}
|
|
|
|
func updateShardsWithRoundRobin(shardsLeaders map[string][]queryNode) map[string][]queryNode {
|
|
|
|
for channelID, leaders := range shardsLeaders {
|
|
if len(leaders) <= 1 {
|
|
continue
|
|
}
|
|
|
|
shardsLeaders[channelID] = append(leaders[1:], leaders[0])
|
|
}
|
|
|
|
return shardsLeaders
|
|
}
|
|
|
|
func roundRobinPolicy(ctx context.Context, getQueryNodePolicy getQueryNodePolicy, query func(UniqueID, types.QueryNode) error, leaders []queryNode) error {
|
|
var (
|
|
err = errBegin
|
|
current = 0
|
|
qn types.QueryNode
|
|
)
|
|
replicaNum := len(leaders)
|
|
|
|
for err != nil && current < replicaNum {
|
|
currentID := leaders[current].nodeID
|
|
if err != errBegin {
|
|
log.Warn("retry with another QueryNode",
|
|
zap.Int("retries numbers", current),
|
|
zap.Int64("nodeID", currentID))
|
|
}
|
|
|
|
qn, err = getQueryNodePolicy(ctx, leaders[current].address)
|
|
if err != nil {
|
|
log.Warn("fail to get valid QueryNode", zap.Int64("nodeID", currentID),
|
|
zap.Error(err))
|
|
current++
|
|
continue
|
|
}
|
|
|
|
defer qn.Stop()
|
|
err = query(currentID, qn)
|
|
if err != nil {
|
|
log.Warn("fail to Query with shard leader",
|
|
zap.Int64("nodeID", currentID),
|
|
zap.Error(err))
|
|
}
|
|
current++
|
|
}
|
|
|
|
if current == replicaNum && err != nil {
|
|
log.Warn("no shard leaders available",
|
|
zap.String("leaders", fmt.Sprintf("%v", leaders)), zap.Error(err))
|
|
// needs to return the error from query
|
|
return err
|
|
}
|
|
return nil
|
|
}
|