/
discovery.go
125 lines (113 loc) · 2.82 KB
/
discovery.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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package discovery
import (
"context"
"encoding/json"
"github.com/hwholiday/learning_tools/hlb-grpc/register"
"go.etcd.io/etcd/api/v3/mvccpb"
clientv3 "go.etcd.io/etcd/client/v3"
"google.golang.org/grpc/resolver"
"log"
"sync"
"time"
)
type NodeArray struct {
Node []register.Options `json:"node"`
}
type Discovery struct {
etcdCli *clientv3.Client
cc resolver.ClientConn
Node sync.Map
opts *Options
}
func NewDiscovery(opt ...ClientOptions) (resolver.Builder, error) {
s := &Discovery{
opts: newOptions(opt...),
}
if s.opts.LoadBalancingPolicy == CustomizeLB {
newCustomizeBuilder(s.opts)
} else {
return nil, ErrLoadBalancingPolicy
}
etcdCli, err := clientv3.New(s.opts.EtcdConf)
if err != nil {
return nil, err
}
s.etcdCli = etcdCli
return s, nil
}
// Build 当调用`grpc.Dial()`时执行
func (d *Discovery) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
d.cc = cc
var ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
res, err := d.etcdCli.Get(ctx, d.opts.SrvName, clientv3.WithPrefix())
if err != nil {
return nil, err
}
for _, v := range res.Kvs {
if err = d.AddNode(v.Key, v.Value); err != nil {
log.Println(err)
continue
}
}
if len(res.Kvs) == 0 {
log.Printf("no %s service found , waiting for the service to join \n", d.opts.SrvName)
}
go func(dd *Discovery) {
dd.watcher()
}(d)
return d, err
}
func (d *Discovery) AddNode(key, val []byte) error {
var data = new(register.Options)
err := json.Unmarshal(val, data)
if err != nil {
return err
}
addr := resolver.Address{Addr: data.Node.Address}
addr = SetNodeInfo(addr, data)
d.Node.Store(string(key), addr)
d.cc.UpdateState(resolver.State{Addresses: d.GetAddress()})
return nil
}
func (d *Discovery) DelNode(key []byte) error {
keyStr := string(key)
d.Node.Delete(keyStr)
d.cc.UpdateState(resolver.State{Addresses: d.GetAddress()})
return nil
}
func (d *Discovery) GetAddress() []resolver.Address {
var addr []resolver.Address
d.Node.Range(func(key, value interface{}) bool {
addr = append(addr, value.(resolver.Address))
return true
})
return addr
}
func (d *Discovery) Scheme() string {
return "discovery"
}
//watcher 监听前缀
func (d *Discovery) watcher() {
rch := d.etcdCli.Watch(context.Background(), d.opts.SrvName, clientv3.WithPrefix())
for res := range rch {
for _, ev := range res.Events {
switch ev.Type {
case mvccpb.PUT: //新增或修改
if err := d.AddNode(ev.Kv.Key, ev.Kv.Value); err != nil {
log.Println(err)
}
case mvccpb.DELETE: //删除
if err := d.DelNode(ev.Kv.Key); err != nil {
log.Println(err)
}
}
}
}
}
func (s *Discovery) ResolveNow(rn resolver.ResolveNowOptions) {
//log.Println("ResolveNow")
}
func (s *Discovery) Close() {
s.etcdCli.Close()
}