...
1
18
19
20
21
22 package roundrobin
23
24 import (
25 "sync/atomic"
26
27 "google.golang.org/grpc/balancer"
28 "google.golang.org/grpc/balancer/base"
29 "google.golang.org/grpc/grpclog"
30 "google.golang.org/grpc/internal/grpcrand"
31 )
32
33
34 const Name = "round_robin"
35
36 var logger = grpclog.Component("roundrobin")
37
38
39 func newBuilder() balancer.Builder {
40 return base.NewBalancerBuilder(Name, &rrPickerBuilder{}, base.Config{HealthCheck: true})
41 }
42
43 func init() {
44 balancer.Register(newBuilder())
45 }
46
47 type rrPickerBuilder struct{}
48
49 func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
50 logger.Infof("roundrobinPicker: Build called with info: %v", info)
51 if len(info.ReadySCs) == 0 {
52 return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
53 }
54 scs := make([]balancer.SubConn, 0, len(info.ReadySCs))
55 for sc := range info.ReadySCs {
56 scs = append(scs, sc)
57 }
58 return &rrPicker{
59 subConns: scs,
60
61
62
63 next: uint32(grpcrand.Intn(len(scs))),
64 }
65 }
66
67 type rrPicker struct {
68
69
70
71 subConns []balancer.SubConn
72 next uint32
73 }
74
75 func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
76 subConnsLen := uint32(len(p.subConns))
77 nextIndex := atomic.AddUint32(&p.next, 1)
78
79 sc := p.subConns[nextIndex%subConnsLen]
80 return balancer.PickResult{SubConn: sc}, nil
81 }
82
View as plain text