-
Notifications
You must be signed in to change notification settings - Fork 249
/
balancer.go
85 lines (70 loc) · 1.94 KB
/
balancer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package memberlist
import (
"github.com/unionj-cloud/go-doudou/v2/toolkit/zlogger"
"sync"
"google.golang.org/grpc/balancer"
balancerbase "google.golang.org/grpc/balancer/base"
)
const Name = "memberlist_weight_balancer"
func newBuilder() balancer.Builder {
return balancerbase.NewBalancerBuilder(Name, &wPickerBuilder{}, balancerbase.Config{HealthCheck: true})
}
func init() {
balancer.Register(newBuilder())
}
type wPickerBuilder struct{}
func (*wPickerBuilder) Build(info balancerbase.PickerBuildInfo) balancer.Picker {
zlogger.Debug().Msgf("[go-doudou] memberlist_weight_balancer Picker: Build called with info: %v", info)
if len(info.ReadySCs) == 0 {
return balancerbase.NewErrPicker(balancer.ErrNoSubConnAvailable)
}
scs := make([]*conn, 0, len(info.ReadySCs))
for sc, v := range info.ReadySCs {
weight := v.Address.BalancerAttributes.Value(WeightAttributeKey{}).(WeightAddrInfo).Weight
scs = append(scs, &conn{sc: sc, weight: weight})
}
return &wPicker{
subConns: scs,
}
}
type wPicker struct {
subConns []*conn
mu sync.Mutex
}
func (p *wPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
p.mu.Lock()
sc := newChooser(p.subConns).pick().sc
p.mu.Unlock()
return balancer.PickResult{SubConn: sc}, nil
}
type WeightAttributeKey struct{}
type WeightAddrInfo struct {
Weight int
}
type conn struct {
sc balancer.SubConn
weight int
currentWeight int
}
// Chooser from naming_client package in nacos-sdk-go
type Chooser struct {
data []*conn
}
// NewChooser initializes a new Chooser for picking from the provided Choices.
func newChooser(cs []*conn) Chooser {
return Chooser{data: cs}
}
func (chs Chooser) pick() conn {
var selected *conn
total := 0
for i := 0; i < len(chs.data); i++ {
s := chs.data[i]
s.currentWeight += s.weight
total += s.weight
if selected == nil || s.currentWeight > selected.currentWeight {
selected = s
}
}
selected.currentWeight -= total
return *selected
}