push use local conn

pkg/discoveryregistry/discovery_register.go

@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/resolver"
 )
 
 type Conn interface {
@@ -13,7 +12,7 @@ type Conn interface {
 	AddOption(opts ...grpc.DialOption)
 	CloseConn(conn grpc.ClientConnInterface)
 	// do not use this method for call rpc
-	GetClientLocalConns() map[string][]resolver.Address
+	GetClientLocalConns() map[string][]grpc.ClientConnInterface
 }
 
 type SvcDiscoveryRegistry interface {

pkg/discoveryregistry/zookeeper/discover.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/log"
+	"github.com/OpenIMSDK/Open-IM-Server/pkg/errs"
 	"github.com/pkg/errors"
 
 	"github.com/go-zookeeper/zk"
@@ -61,7 +62,7 @@ func (s *ZkClient) GetConnsRemote(serviceName string) (conns []resolver.Address,
 				}
 				return nil, errors.Wrap(err, "get children error")
 			}
-			log.ZDebug(context.Background(), "get conns from remote", "conn", string(data))
+			log.ZDebug(context.Background(), "get addrs from remote", "conn", string(data))
 			conns = append(conns, resolver.Address{Addr: string(data), ServerName: serviceName})
 		}
 	}
@@ -70,34 +71,31 @@ func (s *ZkClient) GetConnsRemote(serviceName string) (conns []resolver.Address,
 
 func (s *ZkClient) GetConns(ctx context.Context, serviceName string, opts ...grpc.DialOption) ([]grpc.ClientConnInterface, error) {
 	s.logger.Printf("get conns from client, serviceName: %s", serviceName)
-	s.lock.Lock()
 	opts = append(s.options, opts...)
+	s.lock.Lock()
+	defer s.lock.Unlock()
 	conns := s.localConns[serviceName]
 	if len(conns) == 0 {
 		var err error
 		s.logger.Printf("get conns from zk remote, serviceName: %s", serviceName)
-		conns, err = s.GetConnsRemote(serviceName)
+		addrs, err := s.GetConnsRemote(serviceName)
 		if err != nil {
-			s.lock.Unlock()
 			return nil, err
 		}
-		if len(conns) == 0 {
+		if len(addrs) == 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)
 		}
+		for _, addr := range addrs {
+			cc, err := grpc.DialContext(ctx, addr.Addr, append(s.options, opts...)...)
+			if err != nil {
+				log.ZError(context.Background(), "dialContext failed", err, "addr", addr.Addr, "opts", append(s.options, opts...))
+				return nil, errs.Wrap(err)
+			}
+			conns = append(conns, cc)
+		}
 		s.localConns[serviceName] = conns
 	}
-	s.lock.Unlock()
-	var ret []grpc.ClientConnInterface
-	s.logger.Printf("get conns from zk success, serviceName: %s", serviceName)
-	for _, conn := range conns {
-		cc, 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, cc)
-	}
-	s.logger.Printf("dial ctx success, serviceName: %s", serviceName)
-	return ret, nil
+	return conns, nil
 }
 
 func (s *ZkClient) GetConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (grpc.ClientConnInterface, error) {

pkg/discoveryregistry/zookeeper/register.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/go-zookeeper/zk"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/resolver"
 )
 
 func (s *ZkClient) CreateRpcRootNodes(serviceNames []string) error {
@@ -43,7 +42,7 @@ func (s *ZkClient) UnRegister() error {
 	}
 	time.Sleep(time.Second)
 	s.node = ""
-	s.localConns = make(map[string][]resolver.Address)
+	s.localConns = make(map[string][]grpc.ClientConnInterface)
 	s.resolvers = make(map[string]*Resolver)
 	return nil
 }

pkg/discoveryregistry/zookeeper/zk.go

@@ -37,8 +37,9 @@ type ZkClient struct {
 	lock    sync.Locker
 	options []grpc.DialOption
 
-	resolvers    map[string]*Resolver
-	localConns   map[string][]resolver.Address
+	resolvers  map[string]*Resolver
+	localConns map[string][]grpc.ClientConnInterface
+
 	balancerName string
 
 	logger Logger
@@ -89,7 +90,7 @@ func NewClient(zkServers []string, zkRoot string, options ...ZkOption) (*ZkClien
 		zkRoot:     "/",
 		scheme:     zkRoot,
 		timeout:    timeout,
-		localConns: make(map[string][]resolver.Address),
+		localConns: make(map[string][]grpc.ClientConnInterface),
 		resolvers:  make(map[string]*Resolver),
 		lock:       &sync.Mutex{},
 	}
@@ -197,6 +198,6 @@ func (s *ZkClient) AddOption(opts ...grpc.DialOption) {
 	s.options = append(s.options, opts...)
 }
 
-func (s *ZkClient) GetClientLocalConns() map[string][]resolver.Address {
+func (s *ZkClient) GetClientLocalConns() map[string][]grpc.ClientConnInterface {
 	return s.localConns
 }