init
This commit is contained in:
88
naming/balancer/key/key.go
Normal file
88
naming/balancer/key/key.go
Normal file
@@ -0,0 +1,88 @@
|
||||
package key
|
||||
|
||||
import (
|
||||
"github.com/rs/zerolog/log"
|
||||
"google.golang.org/grpc/balancer"
|
||||
"google.golang.org/grpc/balancer/base"
|
||||
"google.golang.org/grpc/grpclog"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const (
|
||||
Name = "picker_key"
|
||||
DefaultKey = "X-rand-string-balabala"
|
||||
)
|
||||
|
||||
func init() {
|
||||
balancer.Register(newBuilder(DefaultKey))
|
||||
}
|
||||
|
||||
func newBuilder(key string) balancer.Builder {
|
||||
return base.NewBalancerBuilder(
|
||||
Name,
|
||||
&keyPickerBuilder{key: key},
|
||||
base.Config{HealthCheck: false},
|
||||
)
|
||||
}
|
||||
|
||||
type keyPickerBuilder struct {
|
||||
key string
|
||||
}
|
||||
|
||||
func (b *keyPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
|
||||
grpclog.Infof("keyPickerBuilder: newPicker called with info: %v", info)
|
||||
if len(info.ReadySCs) == 0 {
|
||||
return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
|
||||
}
|
||||
|
||||
picker := &keyPicker{
|
||||
addr2sc: make(map[string]balancer.SubConn),
|
||||
key: b.key,
|
||||
}
|
||||
|
||||
for sc, conInfo := range info.ReadySCs {
|
||||
addr := conInfo.Address.Addr
|
||||
picker.setAddr2sc(addr, sc)
|
||||
}
|
||||
|
||||
return picker
|
||||
}
|
||||
|
||||
type keyPicker struct {
|
||||
addr2sc map[string]balancer.SubConn
|
||||
sync.RWMutex
|
||||
key string
|
||||
}
|
||||
|
||||
func (p *keyPicker) setAddr2sc(addr string, sc balancer.SubConn) {
|
||||
p.Lock()
|
||||
defer p.Unlock()
|
||||
p.addr2sc[addr] = sc
|
||||
log.Debug().Str("addr", addr).Msg("set addr 2 sc")
|
||||
}
|
||||
|
||||
func (p *keyPicker) getAddr2sc(addr string) (balancer.SubConn, error) {
|
||||
p.RLock()
|
||||
defer p.RUnlock()
|
||||
sc, ok := p.addr2sc[addr]
|
||||
if ok {
|
||||
return sc, nil
|
||||
}
|
||||
log.Warn().Str("addr", addr).Msg("server warn")
|
||||
// TODO 如果找不到就随便给一个
|
||||
for _, tsc := range p.addr2sc {
|
||||
|
||||
return tsc, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (p *keyPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
|
||||
var res balancer.PickResult
|
||||
|
||||
addr, _ := info.Ctx.Value(p.key).(string)
|
||||
|
||||
sc, _ := p.getAddr2sc(addr)
|
||||
res.SubConn = sc
|
||||
return res, nil
|
||||
}
|
||||
86
naming/register.go
Normal file
86
naming/register.go
Normal file
@@ -0,0 +1,86 @@
|
||||
package naming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.etcd.io/etcd/client/v3"
|
||||
)
|
||||
|
||||
func Register(etcdAddr, name, addr, schema string, ttl int64) error {
|
||||
var err error
|
||||
if cli == nil {
|
||||
cli, err = clientv3.New(clientv3.Config{
|
||||
Endpoints: strings.Split(etcdAddr, ";"),
|
||||
DialTimeout: 5 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("connect to etcd err:%s", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(time.Second * time.Duration(ttl))
|
||||
|
||||
go func() {
|
||||
for {
|
||||
getResp, err := cli.Get(context.Background(), "/"+schema+"/"+name+"/"+addr)
|
||||
if err != nil {
|
||||
fmt.Println("etcd get err:", err)
|
||||
} else if getResp.Count == 0 {
|
||||
err = withAlive(name, addr, schema, ttl)
|
||||
if err != nil {
|
||||
log.Printf("keep alive:%s", err)
|
||||
}
|
||||
}
|
||||
<-ticker.C
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func withAlive(name, addr, schema string, ttl int64) error {
|
||||
leaseResp, err := cli.Grant(context.Background(), ttl)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.Printf("key:%v\n", "/"+schema+"/"+name+"/"+addr)
|
||||
_, err = cli.Put(context.Background(), "/"+schema+"/"+name+"/"+addr, addr, clientv3.WithLease(leaseResp.ID))
|
||||
if err != nil {
|
||||
log.Printf("put etcd error:%s", err)
|
||||
return err
|
||||
}
|
||||
|
||||
ch, err := cli.KeepAlive(context.Background(), leaseResp.ID)
|
||||
if err != nil {
|
||||
log.Printf("keep alive error:%s", err)
|
||||
return err
|
||||
}
|
||||
|
||||
go func() {
|
||||
for range ch {
|
||||
}
|
||||
//for leaseKeepResp := range ch {
|
||||
// log.Println("续约成功", leaseKeepResp.ID)
|
||||
//}
|
||||
//
|
||||
//log.Println("关闭续租")
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func UnRegister(name, addr, schema string) error {
|
||||
if cli != nil {
|
||||
//fmt.Println("unregister...")
|
||||
_, err := cli.Delete(context.Background(), "/"+schema+"/"+name+"/"+addr)
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
112
naming/resolver.go
Normal file
112
naming/resolver.go
Normal file
@@ -0,0 +1,112 @@
|
||||
package naming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.etcd.io/etcd/api/v3/mvccpb"
|
||||
"go.etcd.io/etcd/client/v3"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
var schema string
|
||||
var cli *clientv3.Client
|
||||
|
||||
type etcdResolver struct {
|
||||
rawAddr string
|
||||
schema string
|
||||
cc resolver.ClientConn
|
||||
}
|
||||
|
||||
func NewResolver(etcdAddr, schema string) resolver.Builder {
|
||||
return &etcdResolver{rawAddr: etcdAddr, schema: schema}
|
||||
}
|
||||
|
||||
func (r *etcdResolver) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
|
||||
//fmt.Println("target:", target)
|
||||
var err error
|
||||
if cli == nil {
|
||||
cli, err = clientv3.New(clientv3.Config{
|
||||
Endpoints: strings.Split(r.rawAddr, ";"),
|
||||
DialTimeout: 15 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
r.cc = cc
|
||||
|
||||
go r.watch("/" + target.Scheme + "/" + target.Endpoint + "/")
|
||||
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (r etcdResolver) Scheme() string {
|
||||
return r.schema
|
||||
}
|
||||
|
||||
func (r etcdResolver) ResolveNow(rn resolver.ResolveNowOptions) {
|
||||
//log.Println("ResolveNow")
|
||||
}
|
||||
|
||||
func (r etcdResolver) Close() {
|
||||
//log.Println("Close")
|
||||
}
|
||||
|
||||
func (r *etcdResolver) watch(keyPrefix string) {
|
||||
var addrList []resolver.Address
|
||||
|
||||
getResp, err := cli.Get(context.Background(), keyPrefix, clientv3.WithPrefix())
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
} else {
|
||||
for i := range getResp.Kvs {
|
||||
addrList = append(addrList, resolver.Address{Addr: strings.TrimPrefix(string(getResp.Kvs[i].Key), keyPrefix)})
|
||||
}
|
||||
}
|
||||
|
||||
// 新版本etcd去除了NewAddress方法 以UpdateState代替
|
||||
r.cc.UpdateState(resolver.State{Addresses: addrList})
|
||||
|
||||
rch := cli.Watch(context.Background(), keyPrefix, clientv3.WithPrefix())
|
||||
for n := range rch {
|
||||
for _, ev := range n.Events {
|
||||
addr := strings.TrimPrefix(string(ev.Kv.Key), keyPrefix)
|
||||
switch ev.Type {
|
||||
case mvccpb.PUT:
|
||||
if !exist(addrList, addr) {
|
||||
addrList = append(addrList, resolver.Address{Addr: addr})
|
||||
r.cc.UpdateState(resolver.State{Addresses: addrList})
|
||||
}
|
||||
case mvccpb.DELETE:
|
||||
if s, ok := remove(addrList, addr); ok {
|
||||
addrList = s
|
||||
r.cc.UpdateState(resolver.State{Addresses: addrList})
|
||||
}
|
||||
}
|
||||
log.Printf("%s %q : %q\n", ev.Type, ev.Kv.Key, ev.Kv.Value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func exist(l []resolver.Address, addr string) bool {
|
||||
for i := range l {
|
||||
if l[i].Addr == addr {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func remove(s []resolver.Address, addr string) ([]resolver.Address, bool) {
|
||||
for i := range s {
|
||||
if s[i].Addr == addr {
|
||||
s[i] = s[len(s)-1]
|
||||
return s[:len(s)-1], true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
Reference in New Issue
Block a user