You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
92 lines
2.1 KiB
92 lines
2.1 KiB
package wrr
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/go-kratos/kratos/v2/selector"
|
|
"github.com/go-kratos/kratos/v2/selector/node/direct"
|
|
)
|
|
|
|
const (
|
|
// Name is wrr balancer name
|
|
Name = "wrr"
|
|
)
|
|
|
|
var _ selector.Balancer = &Balancer{} // Name is balancer name
|
|
|
|
// WithFilter with select filters
|
|
func WithFilter(filters ...selector.Filter) Option {
|
|
return func(o *options) {
|
|
o.filters = filters
|
|
}
|
|
}
|
|
|
|
// Option is random builder option.
|
|
type Option func(o *options)
|
|
|
|
// options is random builder options
|
|
type options struct {
|
|
filters []selector.Filter
|
|
}
|
|
|
|
// Balancer is a random balancer.
|
|
type Balancer struct {
|
|
mu sync.Mutex
|
|
currentWeight map[string]float64
|
|
}
|
|
|
|
// New random a selector.
|
|
func New(opts ...Option) selector.Selector {
|
|
return NewBuilder(opts...).Build()
|
|
}
|
|
|
|
// 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
|
|
}
|
|
var totalWeight float64
|
|
var selected selector.WeightedNode
|
|
var selectWeight float64
|
|
|
|
// nginx wrr load balancing algorithm: http://blog.csdn.net/zhangskd/article/details/50194069
|
|
p.mu.Lock()
|
|
for _, node := range nodes {
|
|
totalWeight += node.Weight()
|
|
cwt := p.currentWeight[node.Address()]
|
|
// current += effectiveWeight
|
|
cwt += node.Weight()
|
|
p.currentWeight[node.Address()] = cwt
|
|
if selected == nil || selectWeight < cwt {
|
|
selectWeight = cwt
|
|
selected = node
|
|
}
|
|
}
|
|
p.currentWeight[selected.Address()] = selectWeight - totalWeight
|
|
p.mu.Unlock()
|
|
|
|
d := selected.Pick()
|
|
return selected, d, nil
|
|
}
|
|
|
|
// NewBuilder returns a selector builder with wrr balancer
|
|
func NewBuilder(opts ...Option) selector.Builder {
|
|
var option options
|
|
for _, opt := range opts {
|
|
opt(&option)
|
|
}
|
|
return &selector.DefaultBuilder{
|
|
Filters: option.filters,
|
|
Balancer: &Builder{},
|
|
Node: &direct.Builder{},
|
|
}
|
|
}
|
|
|
|
// Builder is wrr builder
|
|
type Builder struct{}
|
|
|
|
// Build creates Balancer
|
|
func (b *Builder) Build() selector.Balancer {
|
|
return &Balancer{currentWeight: make(map[string]float64)}
|
|
}
|
|
|