Feat: add load balancer (#1437)
* add balancer * add p2c balancer * add http client selector filter Co-authored-by: yuemoxi <99347745@qq.com> Co-authored-by: chenzhihui <zhihui_chen@foxmail.com>pull/1466/head
longxboy
3 years ago
committed by
GitHub
parent
0184d217cf
commit
20f0a07d36
@ -0,0 +1,30 @@ |
||||
package selector |
||||
|
||||
import ( |
||||
"context" |
||||
"time" |
||||
) |
||||
|
||||
// Balancer is balancer interface
|
||||
type Balancer interface { |
||||
Pick(ctx context.Context, nodes []WeightedNode) (selected WeightedNode, done DoneFunc, err error) |
||||
} |
||||
|
||||
// WeightedNode calculates scheduling weight in real time
|
||||
type WeightedNode interface { |
||||
Node |
||||
|
||||
// Weight is the runtime calculated weight
|
||||
Weight() float64 |
||||
|
||||
// Pick the node
|
||||
Pick() DoneFunc |
||||
|
||||
// PickElapsed is time elapsed since the latest pick
|
||||
PickElapsed() time.Duration |
||||
} |
||||
|
||||
// WeightedNodeBuilder is WeightedNode Builder
|
||||
type WeightedNodeBuilder interface { |
||||
Build(Node) WeightedNode |
||||
} |
||||
@ -0,0 +1,49 @@ |
||||
package selector |
||||
|
||||
import ( |
||||
"context" |
||||
"sync" |
||||
) |
||||
|
||||
// Default is composite selector.
|
||||
type Default struct { |
||||
NodeBuilder WeightedNodeBuilder |
||||
Balancer Balancer |
||||
|
||||
lk sync.RWMutex |
||||
weightedNodes []Node |
||||
} |
||||
|
||||
// Select select one node.
|
||||
func (d *Default) Select(ctx context.Context, opts ...SelectOption) (selected Node, done DoneFunc, err error) { |
||||
d.lk.RLock() |
||||
weightedNodes := d.weightedNodes |
||||
d.lk.RUnlock() |
||||
var options SelectOptions |
||||
for _, o := range opts { |
||||
o(&options) |
||||
} |
||||
for _, f := range options.Filters { |
||||
weightedNodes = f(ctx, weightedNodes) |
||||
} |
||||
candidates := make([]WeightedNode, 0, len(weightedNodes)) |
||||
for _, n := range weightedNodes { |
||||
candidates = append(candidates, n.(WeightedNode)) |
||||
} |
||||
if len(candidates) == 0 { |
||||
return nil, nil, ErrNoAvailable |
||||
} |
||||
return d.Balancer.Pick(ctx, candidates) |
||||
} |
||||
|
||||
// Apply update nodes info.
|
||||
func (d *Default) Apply(nodes []Node) { |
||||
weightedNodes := make([]Node, 0, len(nodes)) |
||||
for _, n := range nodes { |
||||
weightedNodes = append(weightedNodes, d.NodeBuilder.Build(n)) |
||||
} |
||||
d.lk.Lock() |
||||
// TODO: Do not delete unchanged nodes
|
||||
d.weightedNodes = weightedNodes |
||||
d.lk.Unlock() |
||||
} |
||||
@ -0,0 +1,20 @@ |
||||
package filter |
||||
|
||||
import ( |
||||
"context" |
||||
|
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
) |
||||
|
||||
// Version is verion filter.
|
||||
func Version(version string) selector.Filter { |
||||
return func(_ context.Context, nodes []selector.Node) []selector.Node { |
||||
filters := make([]selector.Node, 0, len(nodes)) |
||||
for _, n := range nodes { |
||||
if n.Version() == version { |
||||
filters = append(filters, n) |
||||
} |
||||
} |
||||
return filters |
||||
} |
||||
} |
||||
@ -0,0 +1,60 @@ |
||||
package node |
||||
|
||||
import ( |
||||
"strconv" |
||||
|
||||
"github.com/go-kratos/kratos/v2/registry" |
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
) |
||||
|
||||
// Node is slector node
|
||||
type Node struct { |
||||
addr string |
||||
weight *int64 |
||||
version string |
||||
name string |
||||
metadata map[string]string |
||||
} |
||||
|
||||
// Address is node address
|
||||
func (n *Node) Address() string { |
||||
return n.addr |
||||
} |
||||
|
||||
// ServiceName is node serviceName
|
||||
func (n *Node) ServiceName() string { |
||||
return n.name |
||||
} |
||||
|
||||
// InitialWeight is node initialWeight
|
||||
func (n *Node) InitialWeight() *int64 { |
||||
return n.weight |
||||
} |
||||
|
||||
// Version is node version
|
||||
func (n *Node) Version() string { |
||||
return n.version |
||||
} |
||||
|
||||
// Metadata is node metadata
|
||||
func (n *Node) Metadata() map[string]string { |
||||
return n.metadata |
||||
} |
||||
|
||||
// New node
|
||||
func New(addr string, ins *registry.ServiceInstance) selector.Node { |
||||
n := &Node{ |
||||
addr: addr, |
||||
} |
||||
if ins != nil { |
||||
n.name = ins.Name |
||||
n.version = ins.Version |
||||
n.metadata = ins.Metadata |
||||
if str, ok := ins.Metadata["weight"]; ok { |
||||
if weight, err := strconv.ParseInt(str, 10, 64); err == nil { |
||||
n.weight = &weight |
||||
} |
||||
} |
||||
} |
||||
return n |
||||
} |
||||
@ -0,0 +1,52 @@ |
||||
package direct |
||||
|
||||
import ( |
||||
"context" |
||||
"sync/atomic" |
||||
"time" |
||||
|
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
) |
||||
|
||||
const ( |
||||
defaultWeight = 100 |
||||
) |
||||
|
||||
var ( |
||||
_ selector.WeightedNode = &node{} |
||||
_ selector.WeightedNodeBuilder = &Builder{} |
||||
) |
||||
|
||||
// node is endpoint instance
|
||||
type node struct { |
||||
selector.Node |
||||
|
||||
// last lastPick timestamp
|
||||
lastPick int64 |
||||
} |
||||
|
||||
// Builder is direct node builder
|
||||
type Builder struct{} |
||||
|
||||
// Build create node
|
||||
func (*Builder) Build(n selector.Node) selector.WeightedNode { |
||||
return &node{Node: n, lastPick: 0} |
||||
} |
||||
|
||||
func (n *node) Pick() selector.DoneFunc { |
||||
now := time.Now().UnixNano() |
||||
atomic.StoreInt64(&n.lastPick, now) |
||||
return func(ctx context.Context, di selector.DoneInfo) {} |
||||
} |
||||
|
||||
// Weight is node effective weight
|
||||
func (n *node) Weight() float64 { |
||||
if n.InitialWeight() != nil { |
||||
return float64(*n.InitialWeight()) |
||||
} |
||||
return defaultWeight |
||||
} |
||||
|
||||
func (n *node) PickElapsed() time.Duration { |
||||
return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick)) |
||||
} |
||||
@ -0,0 +1,180 @@ |
||||
package ewma |
||||
|
||||
import ( |
||||
"container/list" |
||||
"context" |
||||
"math" |
||||
"sync" |
||||
"sync/atomic" |
||||
"time" |
||||
|
||||
"github.com/go-kratos/kratos/v2/errors" |
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
) |
||||
|
||||
const ( |
||||
// The mean lifetime of `cost`, it reaches its half-life after Tau*ln(2).
|
||||
tau = int64(time.Millisecond * 600) |
||||
// if statistic not collected,we add a big lag penalty to endpoint
|
||||
penalty = uint64(time.Second * 10) |
||||
) |
||||
|
||||
var ( |
||||
_ selector.WeightedNode = &node{} |
||||
_ selector.WeightedNodeBuilder = &Builder{} |
||||
) |
||||
|
||||
// node is endpoint instance
|
||||
type node struct { |
||||
selector.Node |
||||
|
||||
// client statistic data
|
||||
lag int64 |
||||
success uint64 |
||||
inflight int64 |
||||
inflights *list.List |
||||
// last collected timestamp
|
||||
stamp int64 |
||||
predictTs int64 |
||||
predict int64 |
||||
// request number in a period time
|
||||
reqs int64 |
||||
// last lastPick timestamp
|
||||
lastPick int64 |
||||
|
||||
errHandler func(err error) (isErr bool) |
||||
lk sync.RWMutex |
||||
} |
||||
|
||||
// Builder is ewma node builder.
|
||||
type Builder struct { |
||||
ErrHandler func(err error) (isErr bool) |
||||
} |
||||
|
||||
// Build create a weighted node.
|
||||
func (b *Builder) Build(n selector.Node) selector.WeightedNode { |
||||
s := &node{ |
||||
Node: n, |
||||
lag: 0, |
||||
success: 1000, |
||||
inflight: 1, |
||||
inflights: list.New(), |
||||
errHandler: b.ErrHandler, |
||||
} |
||||
return s |
||||
} |
||||
|
||||
func (n *node) health() uint64 { |
||||
return atomic.LoadUint64(&n.success) |
||||
} |
||||
|
||||
func (n *node) load() (load uint64) { |
||||
now := time.Now().UnixNano() |
||||
avgLag := atomic.LoadInt64(&n.lag) |
||||
lastPredictTs := atomic.LoadInt64(&n.predictTs) |
||||
predicInterval := avgLag / 5 |
||||
if predicInterval < int64(time.Millisecond*5) { |
||||
predicInterval = int64(time.Millisecond * 5) |
||||
} else if predicInterval > int64(time.Millisecond*200) { |
||||
predicInterval = int64(time.Millisecond * 200) |
||||
} |
||||
if now-lastPredictTs > predicInterval { |
||||
if atomic.CompareAndSwapInt64(&n.predictTs, lastPredictTs, now) { |
||||
var ( |
||||
total int64 |
||||
count int |
||||
predict int64 |
||||
) |
||||
n.lk.RLock() |
||||
first := n.inflights.Front() |
||||
for first != nil { |
||||
lag := now - first.Value.(int64) |
||||
if lag > avgLag { |
||||
count++ |
||||
total += lag |
||||
} |
||||
first = first.Next() |
||||
} |
||||
if count > (n.inflights.Len()/2 + 1) { |
||||
predict = total / int64(count) |
||||
} |
||||
n.lk.RUnlock() |
||||
atomic.StoreInt64(&n.predict, predict) |
||||
} |
||||
} |
||||
|
||||
if avgLag == 0 { |
||||
// penalty是node刚启动时没有数据时的惩罚值，默认为1e9 * 10
|
||||
load = penalty * uint64(atomic.LoadInt64(&n.inflight)) |
||||
} else { |
||||
predict := atomic.LoadInt64(&n.predict) |
||||
if predict > avgLag { |
||||
avgLag = predict |
||||
} |
||||
load = uint64(avgLag) * uint64(atomic.LoadInt64(&n.inflight)) |
||||
} |
||||
return |
||||
} |
||||
|
||||
// Pick pick a node.
|
||||
func (n *node) Pick() selector.DoneFunc { |
||||
now := time.Now().UnixNano() |
||||
atomic.StoreInt64(&n.lastPick, now) |
||||
atomic.AddInt64(&n.inflight, 1) |
||||
atomic.AddInt64(&n.reqs, 1) |
||||
n.lk.Lock() |
||||
e := n.inflights.PushBack(now) |
||||
n.lk.Unlock() |
||||
return func(ctx context.Context, di selector.DoneInfo) { |
||||
n.lk.Lock() |
||||
n.inflights.Remove(e) |
||||
n.lk.Unlock() |
||||
atomic.AddInt64(&n.inflight, -1) |
||||
|
||||
now := time.Now().UnixNano() |
||||
// get moving average ratio w
|
||||
stamp := atomic.SwapInt64(&n.stamp, now) |
||||
td := now - stamp |
||||
if td < 0 { |
||||
td = 0 |
||||
} |
||||
w := math.Exp(float64(-td) / float64(tau)) |
||||
|
||||
start := e.Value.(int64) |
||||
lag := now - start |
||||
if lag < 0 { |
||||
lag = 0 |
||||
} |
||||
oldLag := atomic.LoadInt64(&n.lag) |
||||
if oldLag == 0 { |
||||
w = 0.0 |
||||
} |
||||
lag = int64(float64(oldLag)*w + float64(lag)*(1.0-w)) |
||||
atomic.StoreInt64(&n.lag, lag) |
||||
|
||||
success := uint64(1000) // error value ,if error set 1
|
||||
if di.Err != nil { |
||||
if n.errHandler != nil { |
||||
if n.errHandler(di.Err) { |
||||
success = 0 |
||||
} |
||||
} else if errors.Is(context.DeadlineExceeded, di.Err) || errors.Is(context.Canceled, di.Err) || |
||||
errors.IsServiceUnavailable(di.Err) || errors.IsGatewayTimeout(di.Err) { |
||||
success = 0 |
||||
} |
||||
} |
||||
oldSuc := atomic.LoadUint64(&n.success) |
||||
success = uint64(float64(oldSuc)*w + float64(success)*(1.0-w)) |
||||
atomic.StoreUint64(&n.success, success) |
||||
} |
||||
} |
||||
|
||||
// Weight is node effective weight.
|
||||
func (n *node) Weight() (weight float64) { |
||||
weight = float64(n.health()*uint64(time.Second)) / float64(n.load()) |
||||
return |
||||
} |
||||
|
||||
func (n *node) PickElapsed() time.Duration { |
||||
return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick)) |
||||
} |
||||
@ -0,0 +1,21 @@ |
||||
package selector |
||||
|
||||
import "context" |
||||
|
||||
// SelectOptions is Select Options.
|
||||
type SelectOptions struct { |
||||
Filters []Filter |
||||
} |
||||
|
||||
// SelectOption is Selector option.
|
||||
type SelectOption func(*SelectOptions) |
||||
|
||||
// Filter is node filter function.
|
||||
type Filter func(context.Context, []Node) []Node |
||||
|
||||
// WithFilter with filter options
|
||||
func WithFilter(fn ...Filter) SelectOption { |
||||
return func(opts *SelectOptions) { |
||||
opts.Filters = fn |
||||
} |
||||
} |
||||
@ -0,0 +1,74 @@ |
||||
package p2c |
||||
|
||||
import ( |
||||
"context" |
||||
"math/rand" |
||||
"sync/atomic" |
||||
"time" |
||||
|
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
"github.com/go-kratos/kratos/v2/selector/node/ewma" |
||||
) |
||||
|
||||
const ( |
||||
forcePick = time.Second * 3 |
||||
// Name is balancer name
|
||||
Name = "p2c" |
||||
) |
||||
|
||||
var _ selector.Balancer = &Balancer{} |
||||
|
||||
// New creates a p2c selector.
|
||||
func New() selector.Selector { |
||||
return &selector.Default{ |
||||
NodeBuilder: &ewma.Builder{}, |
||||
Balancer: &Balancer{ |
||||
r: rand.New(rand.NewSource(time.Now().UnixNano())), |
||||
}, |
||||
} |
||||
} |
||||
|
||||
// Balancer is p2c selector.
|
||||
type Balancer struct { |
||||
r *rand.Rand |
||||
lk int64 |
||||
} |
||||
|
||||
// choose two distinct nodes.
|
||||
func (s *Balancer) prePick(nodes []selector.WeightedNode) (nodeA selector.WeightedNode, nodeB selector.WeightedNode) { |
||||
a := s.r.Intn(len(nodes)) |
||||
b := s.r.Intn(len(nodes) - 1) |
||||
if b >= a { |
||||
b = b + 1 |
||||
} |
||||
nodeA, nodeB = nodes[a], nodes[b] |
||||
return |
||||
} |
||||
|
||||
// Pick pick a node.
|
||||
func (s *Balancer) Pick(ctx context.Context, nodes []selector.WeightedNode) (selector.WeightedNode, selector.DoneFunc, error) { |
||||
if len(nodes) == 0 { |
||||
return nil, nil, selector.ErrNoAvailable |
||||
} else if len(nodes) == 1 { |
||||
done := nodes[0].Pick() |
||||
return nodes[0], done, nil |
||||
} |
||||
|
||||
var pc, upc selector.WeightedNode |
||||
nodeA, nodeB := s.prePick(nodes) |
||||
// meta.Weight为服务发布者在discovery中设置的权重
|
||||
if nodeB.Weight() > nodeA.Weight() { |
||||
pc, upc = nodeB, nodeA |
||||
} else { |
||||
pc, upc = nodeA, nodeB |
||||
} |
||||
|
||||
// 如果落选节点在forceGap期间内从来没有被选中一次，则强制选一次
|
||||
// 利用强制的机会，来触发成功率、延迟的更新
|
||||
if upc.PickElapsed() > forcePick && atomic.CompareAndSwapInt64(&s.lk, 0, 1) { |
||||
pc = upc |
||||
atomic.StoreInt64(&s.lk, 0) |
||||
} |
||||
done := pc.Pick() |
||||
return pc, done, nil |
||||
} |
||||
@ -0,0 +1,38 @@ |
||||
package random |
||||
|
||||
import ( |
||||
"context" |
||||
"math/rand" |
||||
|
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
"github.com/go-kratos/kratos/v2/selector/node/direct" |
||||
) |
||||
|
||||
var ( |
||||
_ selector.Balancer = &Balancer{} |
||||
|
||||
// Name is balancer name
|
||||
Name = "random" |
||||
) |
||||
|
||||
// Balancer is a random balancer.
|
||||
type Balancer struct{} |
||||
|
||||
// New random a selector.
|
||||
func New() selector.Selector { |
||||
return &selector.Default{ |
||||
Balancer: &Balancer{}, |
||||
NodeBuilder: &direct.Builder{}, |
||||
} |
||||
} |
||||
|
||||
// Pick pick a weighted node.
|
||||
func (p *Balancer) Pick(_ context.Context, nodes []selector.WeightedNode) (selector.WeightedNode, selector.DoneFunc, error) { |
||||
if len(nodes) == 0 { |
||||
return nil, nil, selector.ErrNoAvailable |
||||
} |
||||
cur := rand.Intn(len(nodes)) |
||||
selected := nodes[cur] |
||||
d := selected.Pick() |
||||
return selected, d, nil |
||||
} |
||||
@ -0,0 +1,66 @@ |
||||
package selector |
||||
|
||||
import ( |
||||
"context" |
||||
|
||||
"github.com/go-kratos/kratos/v2/errors" |
||||
) |
||||
|
||||
// ErrNoAvailable is no available node.
|
||||
var ErrNoAvailable = errors.ServiceUnavailable("no_available_node", "") |
||||
|
||||
// Selector is node pick balancer.
|
||||
type Selector interface { |
||||
Rebalancer |
||||
|
||||
// Select nodes
|
||||
// if err == nil, selected and done must not be empty.
|
||||
Select(ctx context.Context, opts ...SelectOption) (selected Node, done DoneFunc, err error) |
||||
} |
||||
|
||||
// Rebalancer is nodes rebalancer.
|
||||
type Rebalancer interface { |
||||
// apply all nodes when any changes happen
|
||||
Apply(nodes []Node) |
||||
} |
||||
|
||||
// Node is node interface.
|
||||
type Node interface { |
||||
// Address is the unique address under the same service
|
||||
Address() string |
||||
|
||||
// ServiceName is service name
|
||||
ServiceName() string |
||||
|
||||
// InitialWeight is the initial value of scheduling weight
|
||||
// if not set return nil
|
||||
InitialWeight() *int64 |
||||
|
||||
// Version is service node version
|
||||
Version() string |
||||
|
||||
// Metadata is the kv pair metadata associated with the service instance.
|
||||
// version,namespace,region,protocol etc..
|
||||
Metadata() map[string]string |
||||
} |
||||
|
||||
// DoneInfo is callback info when RPC invoke done.
|
||||
type DoneInfo struct { |
||||
// Response Error
|
||||
Err error |
||||
// Response Metadata
|
||||
ReplyMeta ReplyMeta |
||||
|
||||
// BytesSent indicates if any bytes have been sent to the server.
|
||||
BytesSent bool |
||||
// BytesReceived indicates if any byte has been received from the server.
|
||||
BytesReceived bool |
||||
} |
||||
|
||||
// ReplyMeta is Reply Metadata.
|
||||
type ReplyMeta interface { |
||||
Get(key string) string |
||||
} |
||||
|
||||
// DoneFunc is callback function when RPC invoke done.
|
||||
type DoneFunc func(ctx context.Context, di DoneInfo) |
||||
@ -0,0 +1,106 @@ |
||||
package balancer |
||||
|
||||
import ( |
||||
"sync" |
||||
|
||||
"github.com/go-kratos/kratos/v2/registry" |
||||
"github.com/go-kratos/kratos/v2/selector" |
||||
"github.com/go-kratos/kratos/v2/selector/node" |
||||
"github.com/go-kratos/kratos/v2/selector/p2c" |
||||
"github.com/go-kratos/kratos/v2/selector/random" |
||||
|
||||
gBalancer "google.golang.org/grpc/balancer" |
||||
"google.golang.org/grpc/balancer/base" |
||||
"google.golang.org/grpc/metadata" |
||||
) |
||||
|
||||
var ( |
||||
_ base.PickerBuilder = &Builder{} |
||||
_ gBalancer.Picker = &Picker{} |
||||
|
||||
mu sync.Mutex |
||||
) |
||||
|
||||
func init() { |
||||
// inject global grpc balancer
|
||||
SetGlobalBalancer(random.Name, random.New()) |
||||
SetGlobalBalancer(p2c.Name, p2c.New()) |
||||
} |
||||
|
||||
// SetGlobalBalancer set grpc balancer with scheme.
|
||||
func SetGlobalBalancer(scheme string, selector selector.Selector) { |
||||
mu.Lock() |
||||
defer mu.Unlock() |
||||
|
||||
b := base.NewBalancerBuilder( |
||||
scheme, |
||||
&Builder{selector}, |
||||
base.Config{HealthCheck: true}, |
||||
) |
||||
gBalancer.Register(b) |
||||
} |
||||
|
||||
// Builder is grpc balancer builder.
|
||||
type Builder struct { |
||||
selector selector.Selector |
||||
} |
||||
|
||||
// Build creates a grpc Picker.
|
||||
func (b *Builder) Build(info base.PickerBuildInfo) gBalancer.Picker { |
||||
nodes := make([]selector.Node, 0) |
||||
subConns := make(map[string]gBalancer.SubConn) |
||||
for conn, info := range info.ReadySCs { |
||||
if _, ok := subConns[info.Address.Addr]; ok { |
||||
continue |
||||
} |
||||
subConns[info.Address.Addr] = conn |
||||
|
||||
ins, _ := info.Address.Attributes.Value("rawServiceInstance").(*registry.ServiceInstance) |
||||
nodes = append(nodes, node.New(info.Address.Addr, ins)) |
||||
} |
||||
p := &Picker{ |
||||
selector: b.selector, |
||||
subConns: subConns, |
||||
} |
||||
p.selector.Apply(nodes) |
||||
return p |
||||
} |
||||
|
||||
// Picker is a grpc picker.
|
||||
type Picker struct { |
||||
subConns map[string]gBalancer.SubConn |
||||
selector selector.Selector |
||||
} |
||||
|
||||
// Pick pick instances.
|
||||
func (p *Picker) Pick(info gBalancer.PickInfo) (gBalancer.PickResult, error) { |
||||
n, done, err := p.selector.Select(info.Ctx) |
||||
if err != nil { |
||||
return gBalancer.PickResult{}, err |
||||
} |
||||
sub := p.subConns[n.Address()] |
||||
|
||||
return gBalancer.PickResult{ |
||||
SubConn: sub, |
||||
Done: func(di gBalancer.DoneInfo) { |
||||
done(info.Ctx, selector.DoneInfo{ |
||||
Err: di.Err, |
||||
BytesSent: di.BytesSent, |
||||
BytesReceived: di.BytesReceived, |
||||
ReplyMeta: Trailer(di.Trailer), |
||||
}) |
||||
}, |
||||
}, nil |
||||
} |
||||
|
||||
// Trailer is a grpc trailder MD.
|
||||
type Trailer metadata.MD |
||||
|
||||
// Get get a grpc trailer value.
|
||||
func (t Trailer) Get(k string) string { |
||||
v := metadata.MD(t).Get(k) |
||||
if len(v) > 0 { |
||||
return v[0] |
||||
} |
||||
return "" |
||||
} |
||||
@ -1,21 +0,0 @@ |
||||
package balancer |
||||
|
||||
import ( |
||||
"context" |
||||
|
||||
"github.com/go-kratos/kratos/v2/registry" |
||||
) |
||||
|
||||
// DoneInfo is callback when rpc done
|
||||
type DoneInfo struct { |
||||
Err error |
||||
Trailer map[string]string |
||||
} |
||||
|
||||
// Balancer is node pick balancer
|
||||
type Balancer interface { |
||||
// Pick one node
|
||||
Pick(ctx context.Context) (node *registry.ServiceInstance, done func(context.Context, DoneInfo), err error) |
||||
// Update nodes when nodes removed or added
|
||||
Update(nodes []*registry.ServiceInstance) |
||||
} |
||||
@ -1,43 +0,0 @@ |
||||
package random |
||||
|
||||
import ( |
||||
"context" |
||||
"fmt" |
||||
"math/rand" |
||||
"sync" |
||||
|
||||
"github.com/go-kratos/kratos/v2/registry" |
||||
"github.com/go-kratos/kratos/v2/transport/http/balancer" |
||||
) |
||||
|
||||
var _ balancer.Balancer = &Balancer{} |
||||
|
||||
type Balancer struct { |
||||
lock sync.RWMutex |
||||
nodes []*registry.ServiceInstance |
||||
} |
||||
|
||||
func New() *Balancer { |
||||
return &Balancer{} |
||||
} |
||||
|
||||
func (b *Balancer) Pick(ctx context.Context) (node *registry.ServiceInstance, done func(context.Context, balancer.DoneInfo), err error) { |
||||
b.lock.RLock() |
||||
nodes := b.nodes |
||||
b.lock.RUnlock() |
||||
|
||||
if len(nodes) == 0 { |
||||
return nil, nil, fmt.Errorf("no instances available") |
||||
} |
||||
if len(nodes) == 1 { |
||||
return nodes[0], func(context.Context, balancer.DoneInfo) {}, nil |
||||
} |
||||
idx := rand.Intn(len(nodes)) |
||||
return nodes[idx], func(context.Context, balancer.DoneInfo) {}, nil |
||||
} |
||||
|
||||
func (b *Balancer) Update(nodes []*registry.ServiceInstance) { |
||||
b.lock.Lock() |
||||
defer b.lock.Unlock() |
||||
b.nodes = nodes |
||||
} |
||||
Loading…
Reference in new issue