wangchuxiao
1 year ago
parent
4c45148eca
commit
4c17c64403
@ -0,0 +1,27 @@
|
||||
package openKeeper
|
||||
|
||||
import (
|
||||
"github.com/go-zookeeper/zk"
|
||||
)
|
||||
|
||||
func (s *ZkClient) RegisterConf2Registry(key string, conf []byte) error {
|
||||
exists, _, err := s.conn.Exists(s.getPath(key))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
if err := s.conn.Delete(s.getPath(key), 0); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
_, err = s.conn.Create(s.getPath(key), conf, 0, zk.WorldACL(zk.PermAll))
|
||||
if err != zk.ErrNodeExists {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetConfFromRegistry(key string) ([]byte, error) {
|
||||
bytes, _, err := s.conn.Get(s.getPath(key))
|
||||
return bytes, err
|
||||
}
|
||||
@ -0,0 +1,110 @@
|
||||
package openKeeper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/go-zookeeper/zk"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
var ErrConnIsNil = errors.New("conn is nil")
|
||||
var ErrConnIsNilButLocalNotNil = errors.New("conn is nil, but local is not nil")
|
||||
|
||||
func (s *ZkClient) watch(wg *sync.WaitGroup) {
|
||||
wg.Done()
|
||||
for {
|
||||
event := <-s.eventChan
|
||||
switch event.Type {
|
||||
case zk.EventSession:
|
||||
case zk.EventNodeCreated:
|
||||
case zk.EventNodeChildrenChanged:
|
||||
l := strings.Split(event.Path, "/")
|
||||
s.lock.Lock()
|
||||
if len(l) > 1 {
|
||||
rpcName := l[len(l)-1]
|
||||
s.flushResolver(rpcName)
|
||||
if len(s.localConns[rpcName]) != 0 {
|
||||
delete(s.localConns, rpcName)
|
||||
}
|
||||
}
|
||||
s.lock.Unlock()
|
||||
case zk.EventNodeDataChanged:
|
||||
case zk.EventNodeDeleted:
|
||||
case zk.EventNotWatching:
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetConnsRemote(serviceName string) (conns []resolver.Address, err error) {
|
||||
path := s.getPath(serviceName)
|
||||
childNodes, _, err := s.conn.Children(path)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get children error")
|
||||
}
|
||||
for _, child := range childNodes {
|
||||
fullPath := path + "/" + child
|
||||
data, _, err := s.conn.Get(fullPath)
|
||||
if err != nil {
|
||||
if err == zk.ErrNoNode {
|
||||
return nil, errors.Wrap(err, "this is zk ErrNoNode")
|
||||
}
|
||||
return nil, errors.Wrap(err, "get children error")
|
||||
}
|
||||
conns = append(conns, resolver.Address{Addr: string(data), ServerName: serviceName})
|
||||
}
|
||||
_, _, _, err = s.conn.ChildrenW(s.getPath(serviceName))
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "children watch error")
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
return nil, fmt.Errorf("no conn for service %s, grpc server may not exist, local conn is %v, please check zookeeper server %v, path: %s", serviceName, s.localConns, s.zkServers, s.zkRoot)
|
||||
}
|
||||
return conns, nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetConns(ctx context.Context, serviceName string, opts ...grpc.DialOption) ([]*grpc.ClientConn, error) {
|
||||
s.lock.Lock()
|
||||
defer s.lock.Unlock()
|
||||
opts = append(s.options, opts...)
|
||||
conns := s.localConns[serviceName]
|
||||
if len(conns) == 0 {
|
||||
var err error
|
||||
conns, err = s.GetConnsRemote(serviceName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.localConns[serviceName] = conns
|
||||
}
|
||||
var ret []*grpc.ClientConn
|
||||
for _, conn := range conns {
|
||||
c, err := grpc.DialContext(ctx, conn.Addr, append(s.options, opts...)...)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, fmt.Sprintf("conns dialContext error, conn: %s", conn.Addr))
|
||||
}
|
||||
ret = append(ret, c)
|
||||
}
|
||||
return ret, nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
|
||||
newOpts := append(s.options, grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"LoadBalancingPolicy": "%s"}`, s.balancerName)))
|
||||
return grpc.DialContext(ctx, fmt.Sprintf("%s:///%s", s.scheme, serviceName), append(newOpts, opts...)...)
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetFirstConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
|
||||
conns, err := s.GetConns(ctx, serviceName, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
return nil, ErrConnIsNil
|
||||
}
|
||||
return conns[0], nil
|
||||
}
|
||||
@ -0,0 +1,39 @@
|
||||
package openKeeper
|
||||
|
||||
import (
|
||||
"github.com/go-zookeeper/zk"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
func (s *ZkClient) Register(rpcRegisterName, host string, port int, opts ...grpc.DialOption) error {
|
||||
s.lock.Lock()
|
||||
defer s.lock.Unlock()
|
||||
if err := s.ensureName(rpcRegisterName); err != nil {
|
||||
return err
|
||||
}
|
||||
addr := s.getAddr(host, port)
|
||||
_, err := grpc.Dial(addr, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
node, err := s.conn.CreateProtectedEphemeralSequential(s.getPath(rpcRegisterName)+"/"+addr+"_", []byte(addr), zk.WorldACL(zk.PermAll))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.node = node
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) UnRegister() error {
|
||||
s.lock.Lock()
|
||||
defer s.lock.Unlock()
|
||||
err := s.conn.Delete(s.node, -1)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.node = ""
|
||||
s.localConns = make(map[string][]resolver.Address)
|
||||
s.resolvers = make(map[string]*Resolver)
|
||||
return nil
|
||||
}
|
||||
@ -0,0 +1,51 @@
|
||||
package openKeeper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/OpenIMSDK/Open-IM-Server/pkg/common/log"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
type Resolver struct {
|
||||
target resolver.Target
|
||||
cc resolver.ClientConn
|
||||
addrs []resolver.Address
|
||||
lock sync.RWMutex
|
||||
|
||||
getConnsRemote func(serviceName string) (conns []resolver.Address, err error)
|
||||
}
|
||||
|
||||
func (r *Resolver) ResolveNow(o resolver.ResolveNowOptions) {
|
||||
log.ZDebug(context.Background(), "start resolve now", "target", r.target)
|
||||
newConns, err := r.getConnsRemote(strings.TrimLeft(r.target.URL.Path, "/"))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
log.ZDebug(context.Background(), "resolve now", "newConns", newConns, "target", r.target)
|
||||
r.lock.Lock()
|
||||
r.addrs = newConns
|
||||
r.lock.Unlock()
|
||||
r.cc.UpdateState(resolver.State{Addresses: r.addrs})
|
||||
log.ZDebug(context.Background(), "resolve now ok", "newConns", newConns, "target", r.target)
|
||||
}
|
||||
|
||||
func (s *Resolver) Close() {}
|
||||
|
||||
func (s *ZkClient) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
|
||||
log.ZDebug(context.Background(), "build resolver", "target", target, "opts", opts, "cc", cc)
|
||||
r := &Resolver{}
|
||||
r.target = target
|
||||
r.cc = cc
|
||||
r.getConnsRemote = s.GetConnsRemote
|
||||
r.ResolveNow(resolver.ResolveNowOptions{})
|
||||
s.lock.Lock()
|
||||
defer s.lock.Unlock()
|
||||
s.resolvers[strings.TrimLeft(target.URL.Path, "/")] = r
|
||||
log.ZDebug(context.Background(), "build resolver ok", "target", target)
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) Scheme() string { return s.scheme }
|
||||
@ -0,0 +1,187 @@
|
||||
package openKeeper
|
||||
|
||||
import (
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-zookeeper/zk"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultFreq = time.Minute * 30
|
||||
timeout = 5
|
||||
)
|
||||
|
||||
type ZkClient struct {
|
||||
zkServers []string
|
||||
zkRoot string
|
||||
userName string
|
||||
password string
|
||||
|
||||
scheme string
|
||||
|
||||
timeout int
|
||||
conn *zk.Conn
|
||||
eventChan <-chan zk.Event
|
||||
node string
|
||||
ticker *time.Ticker
|
||||
|
||||
lock sync.RWMutex
|
||||
options []grpc.DialOption
|
||||
|
||||
resolvers map[string]*Resolver
|
||||
localConns map[string][]resolver.Address
|
||||
balancerName string
|
||||
}
|
||||
|
||||
type ZkOption func(*ZkClient)
|
||||
|
||||
func WithRoundRobin() ZkOption {
|
||||
return func(client *ZkClient) {
|
||||
client.balancerName = "round_robin"
|
||||
}
|
||||
}
|
||||
|
||||
func WithUserNameAndPassword(userName, password string) ZkOption {
|
||||
return func(client *ZkClient) {
|
||||
client.userName = userName
|
||||
client.password = password
|
||||
}
|
||||
}
|
||||
|
||||
func WithOptions(opts ...grpc.DialOption) ZkOption {
|
||||
return func(client *ZkClient) {
|
||||
client.options = opts
|
||||
}
|
||||
}
|
||||
|
||||
func WithFreq(freq time.Duration) ZkOption {
|
||||
return func(client *ZkClient) {
|
||||
client.ticker = time.NewTicker(freq)
|
||||
}
|
||||
}
|
||||
|
||||
func WithTimeout(timeout int) ZkOption {
|
||||
return func(client *ZkClient) {
|
||||
client.timeout = timeout
|
||||
}
|
||||
}
|
||||
|
||||
func NewClient(zkServers []string, zkRoot string, options ...ZkOption) (*ZkClient, error) {
|
||||
client := &ZkClient{
|
||||
zkServers: zkServers,
|
||||
zkRoot: "/",
|
||||
scheme: zkRoot,
|
||||
timeout: timeout,
|
||||
localConns: make(map[string][]resolver.Address),
|
||||
resolvers: make(map[string]*Resolver),
|
||||
}
|
||||
client.ticker = time.NewTicker(defaultFreq)
|
||||
for _, option := range options {
|
||||
option(client)
|
||||
}
|
||||
conn, eventChan, err := zk.Connect(zkServers, time.Duration(client.timeout)*time.Second, zk.WithLogInfo(false))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if client.userName != "" && client.password != "" {
|
||||
if err := conn.AddAuth("digest", []byte(client.userName+":"+client.password)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
client.zkRoot += zkRoot
|
||||
client.eventChan = eventChan
|
||||
client.conn = conn
|
||||
if err := client.ensureRoot(); err != nil {
|
||||
client.CloseZK()
|
||||
return nil, err
|
||||
}
|
||||
resolver.Register(client)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go client.refresh(&wg)
|
||||
go client.watch(&wg)
|
||||
wg.Wait()
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) CloseZK() {
|
||||
s.conn.Close()
|
||||
}
|
||||
|
||||
func (s *ZkClient) ensureAndCreate(node string) error {
|
||||
exists, _, err := s.conn.Exists(node)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
_, err := s.conn.Create(node, []byte(""), 0, zk.WorldACL(zk.PermAll))
|
||||
if err != nil && err != zk.ErrNodeExists {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ZkClient) refresh(wg *sync.WaitGroup) {
|
||||
wg.Done()
|
||||
for range s.ticker.C {
|
||||
s.lock.Lock()
|
||||
for rpcName := range s.resolvers {
|
||||
s.flushResolver(rpcName)
|
||||
}
|
||||
for rpcName := range s.localConns {
|
||||
delete(s.localConns, rpcName)
|
||||
}
|
||||
s.lock.Unlock()
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (s *ZkClient) flushResolver(serviceName string) {
|
||||
r, ok := s.resolvers[serviceName]
|
||||
if ok {
|
||||
r.ResolveNow(resolver.ResolveNowOptions{})
|
||||
s.resolvers[serviceName] = r
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetZkConn() *zk.Conn {
|
||||
return s.conn
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetRootPath() string {
|
||||
return s.zkRoot
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetNode() string {
|
||||
return s.node
|
||||
}
|
||||
|
||||
func (s *ZkClient) ensureRoot() error {
|
||||
return s.ensureAndCreate(s.zkRoot)
|
||||
}
|
||||
|
||||
func (s *ZkClient) ensureName(rpcRegisterName string) error {
|
||||
return s.ensureAndCreate(s.getPath(rpcRegisterName))
|
||||
}
|
||||
|
||||
func (s *ZkClient) getPath(rpcRegisterName string) string {
|
||||
return s.zkRoot + "/" + rpcRegisterName
|
||||
}
|
||||
|
||||
func (s *ZkClient) getAddr(host string, port int) string {
|
||||
return net.JoinHostPort(host, strconv.Itoa(port))
|
||||
}
|
||||
|
||||
func (s *ZkClient) AddOption(opts ...grpc.DialOption) {
|
||||
s.options = append(s.options, opts...)
|
||||
}
|
||||
|
||||
func (s *ZkClient) GetClientLocalConns() map[string][]resolver.Address {
|
||||
return s.localConns
|
||||
}
|
||||
Loading…
Reference in new issue