/
discoverer.go
98 lines (85 loc) · 2.12 KB
/
discoverer.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
91
92
93
94
95
96
97
98
package etcdcenter
import (
"context"
"strings"
"go.etcd.io/etcd/api/v3/mvccpb"
clientv3 "go.etcd.io/etcd/client/v3"
"google.golang.org/grpc/resolver"
)
func newEtcdBuilder(cli *clientv3.Client) resolver.Builder {
return &EtcdBuilder{cli: cli}
}
type EtcdBuilder struct {
cli *clientv3.Client
}
func (s *EtcdBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
r := &etcdResolver{
cli: s.cli,
cc: cc,
}
prefix := strings.Replace(target.URL.Path, "/", "", 1)
r.search(prefix)
go r.watch(prefix)
return r, nil
}
func (s *EtcdBuilder) Scheme() string {
return "etcd"
}
type etcdResolver struct {
cli *clientv3.Client
cc resolver.ClientConn
addrs []resolver.Address
}
func (s *etcdResolver) search(prefix string) error {
rsp, err := s.cli.Get(context.TODO(), prefix, clientv3.WithPrefix())
if err != nil {
return err
}
for _, v := range rsp.Kvs {
s.addrs = append(s.addrs, resolver.Address{Addr: string(v.Value)})
}
s.cc.UpdateState(resolver.State{Addresses: s.addrs})
return nil
}
func (s *etcdResolver) watch(prefix string) {
wchan := s.cli.Watch(context.Background(), prefix, clientv3.WithPrefix())
for ch := range wchan {
for _, v := range ch.Events {
addr := string(v.Kv.Key)
switch v.Type {
case mvccpb.PUT:
if !s.existAddr(addr) {
s.addrs = append(s.addrs, resolver.Address{Addr: addr})
s.cc.UpdateState(resolver.State{Addresses: s.addrs})
}
case mvccpb.DELETE:
if s.removeAddr(addr) {
s.cc.UpdateState(resolver.State{Addresses: s.addrs})
}
}
}
}
}
func (s *etcdResolver) existAddr(addr string) bool {
for _, v := range s.addrs {
if v.Addr == addr {
return true
}
}
return false
}
func (s *etcdResolver) removeAddr(addr string) bool {
isRemove := false
addrs := make([]resolver.Address, 0, len(s.addrs))
for _, v := range s.addrs {
if v.Addr == addr {
isRemove = true
continue
}
addrs = append(addrs, v)
}
s.addrs = addrs
return isRemove
}
func (s *etcdResolver) ResolveNow(resolver.ResolveNowOptions) {}
func (s *etcdResolver) Close() {}