This commit is contained in:
2022-03-17 15:55:27 +08:00
commit bd5a9fad97
92 changed files with 13861 additions and 0 deletions

View 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
View 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
View 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
}