/
rrpicker.go
67 lines (55 loc) · 1.6 KB
/
rrpicker.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
package picker
import (
"math/rand"
"sync"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
"google.golang.org/grpc/resolver"
)
type RRPickerBuilder struct {
}
func (pb *RRPickerBuilder) Build(info PickerBuildInfo) balancer.Picker {
if len(info.ReadySCs) == 0 {
return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
}
scs := []balancer.SubConn{}
scToAddr := make(map[balancer.SubConn]resolver.Address)
for sc, scInfo := range info.ReadySCs {
scs = append(scs, sc)
scToAddr[sc] = scInfo.Address
}
return &rrPicker{
subConns: scs,
scToAddr: scToAddr,
// Start at a random index, as the same RR balancer rebuilds a new
// picker when SubConn states change, and we don't want to apply excess
// load to the first server in the list.
next: rand.Intn(len(scs)),
}
}
type rrPicker struct {
// subConns is the snapshot of the roundrobin balancer when this picker was
// created. The slice is immutable. Each Get() will do a round robin
// selection from it and return the selected SubConn.
subConns []balancer.SubConn
scToAddr map[balancer.SubConn]resolver.Address
mu sync.Mutex
next int
}
func (p *rrPicker) Pick(opts balancer.PickInfo) (balancer.PickResult, error) {
p.mu.Lock()
n := len(p.subConns)
p.mu.Unlock()
if n == 0 {
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
}
p.mu.Lock()
sc := p.subConns[p.next]
// picked := p.scToAddr[sc]
p.next = (p.next + 1) % len(p.subConns)
p.mu.Unlock()
done := func(info balancer.DoneInfo) {
// log.Println("rrpicker done", picked.Addr)
}
return balancer.PickResult{SubConn: sc, Done: done}, nil
}