fix bug: zk client recreate node when reconn

pkg/discoveryregistry/zookeeper/discover.go

@@ -40,7 +40,17 @@ func (s *ZkClient) watch() {
 		event := <-s.eventChan
 		switch event.Type {
 		case zk.EventSession:
+			if event.State == zk.StateHasSession && s.isRegistered {
+				s.logger.Printf("zk session event stateHasSession: %+v, client prepare to create new temp node", event)
+				node, err := s.CreateTempNode(s.rpcRegisterName, s.rpcRegisterAddr)
+				if err != nil {
+					s.logger.Printf("zk session event stateHasSession: %+v, create temp node error: %v", event, err)
+				} else {
+					s.node = node
+				}
+			} else {
 				s.logger.Printf("zk session event: %+v", event)
+			}
 		case zk.EventNodeChildrenChanged:
 			s.logger.Printf("zk event: %s", event.Path)
 			l := strings.Split(event.Path, "/")

pkg/discoveryregistry/zookeeper/register.go

@@ -30,6 +30,14 @@ func (s *ZkClient) CreateRpcRootNodes(serviceNames []string) error {
 	return nil
 }
 
+func (s *ZkClient) CreateTempNode(rpcRegisterName, addr string) (node string, err error) {
+	return s.conn.CreateProtectedEphemeralSequential(
+		s.getPath(rpcRegisterName)+"/"+addr+"_",
+		[]byte(addr),
+		zk.WorldACL(zk.PermAll),
+	)
+}
+
 func (s *ZkClient) Register(rpcRegisterName, host string, port int, opts ...grpc.DialOption) error {
 	if err := s.ensureName(rpcRegisterName); err != nil {
 		return err
@@ -39,15 +47,14 @@ func (s *ZkClient) Register(rpcRegisterName, host string, port int, opts ...grpc
 	if err != nil {
 		return err
 	}
-	node, err := s.conn.CreateProtectedEphemeralSequential(
+	node, err := s.CreateTempNode(rpcRegisterName, addr)
-		s.getPath(rpcRegisterName)+"/"+addr+"_",
-		[]byte(addr),
-		zk.WorldACL(zk.PermAll),
-	)
 	if err != nil {
 		return err
 	}
+	s.rpcRegisterName = rpcRegisterName
+	s.rpcRegisterAddr = addr
 	s.node = node
+	s.isRegistered = true
 	return nil
 }
 
@@ -60,6 +67,9 @@ func (s *ZkClient) UnRegister() error {
 	}
 	time.Sleep(time.Second)
 	s.node = ""
+	s.rpcRegisterName = ""
+	s.rpcRegisterAddr = ""
+	s.isRegistered = false
 	s.localConns = make(map[string][]grpc.ClientConnInterface)
 	s.resolvers = make(map[string]*Resolver)
 	return nil

pkg/discoveryregistry/zookeeper/zk.go

@@ -40,6 +40,9 @@ type ZkClient struct {
 	userName  string
 	password  string
 
+	rpcRegisterName string
+	rpcRegisterAddr string
+	isRegistered    bool
 	scheme          string
 
 	timeout   int
@@ -136,6 +139,7 @@ func NewClient(zkServers []string, zkRoot string, options ...ZkOption) (*ZkClien
 	resolver.Register(client)
 	go client.refresh()
 	go client.watch()
+	time.Sleep(time.Millisecond * 50)
 	return client, nil
 }