-
Notifications
You must be signed in to change notification settings - Fork 5
/
builder.go
90 lines (78 loc) · 2.09 KB
/
builder.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
86
87
88
89
90
package balancer
import (
"github.com/zehuamama/balancer/balancer"
grpcBalancer "google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
)
var (
BuilderName = "easycarBalancer"
_ base.PickerBuilder = (*Builder)(nil)
_ grpcBalancer.Picker = (*Picker)(nil)
)
type Option func(options *Options)
type Options struct {
algorithm Algorithm
}
func WithAlgorithm(algorithm Algorithm) Option {
return func(options *Options) {
options.algorithm = algorithm
}
}
func Register(fns ...Option) {
// RandomBalancer default algorithm
initOption := &Options{algorithm: RandomBalancer}
for _, fn := range fns {
fn(initOption)
}
builder := &Builder{options: initOption}
grpcBalancer.Register(base.NewBalancerBuilder(BuilderName, builder, base.Config{HealthCheck: true}))
}
type Builder struct {
options *Options
}
func (b *Builder) Build(info base.PickerBuildInfo) grpcBalancer.Picker {
if len(info.ReadySCs) == 0 {
return base.NewErrPicker(grpcBalancer.ErrNoSubConnAvailable)
}
var (
hosts []string
)
picker := &Picker{
readyAddrConn: make(map[string]grpcBalancer.SubConn),
}
for subConn, connInfo := range info.ReadySCs {
hosts = append(hosts, connInfo.Address.Addr)
picker.readyAddrConn[connInfo.Address.Addr] = subConn
}
var (
err error
)
picker.balancer, err = balancer.Build(b.options.algorithm.Name(), hosts)
if err != nil {
return base.NewErrPicker(err)
}
return picker
}
type Picker struct {
balancer balancer.Balancer
readyAddrConn map[string]grpcBalancer.SubConn
}
func (p *Picker) Pick(info grpcBalancer.PickInfo) (grpcBalancer.PickResult, error) {
// todo get client ip
// todo Add middleware
addr, err := p.balancer.Balance(info.FullMethodName)
if err != nil {
return grpcBalancer.PickResult{}, grpcBalancer.ErrNoSubConnAvailable
}
subConn, ok := p.readyAddrConn[addr]
if !ok {
return grpcBalancer.PickResult{}, grpcBalancer.ErrNoSubConnAvailable
}
p.balancer.Inc(addr)
return grpcBalancer.PickResult{
SubConn: subConn,
Done: func(info grpcBalancer.DoneInfo) {
p.balancer.Done(addr)
},
}, nil
}