From f924f970e2a677c1aa3bce1c1f8c846b5627a5b4 Mon Sep 17 00:00:00 2001
From: wangchuxiao <wangchuxiao97@outlook.com>
Date: Mon, 31 Jul 2023 15:29:33 +0800
Subject: [PATCH] fix: zk

---
 discoveryregistry/zookeeper/discover.go | 23 +++++++++++++++--------
 discoveryregistry/zookeeper/zk.go       | 10 +++++-----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/discoveryregistry/zookeeper/discover.go b/discoveryregistry/zookeeper/discover.go
index b08e417..e355f6b 100644
--- a/discoveryregistry/zookeeper/discover.go
+++ b/discoveryregistry/zookeeper/discover.go
@@ -44,15 +44,22 @@ func (s *ZkClient) watch(ctx context.Context) {
 		log.ZDebug(context.Background(), "zk recv event", "event", event)
 		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
+			switch event.State {
+			case zk.StateHasSession:
+				if s.isRegistered && !s.isStateDisconnected {
+					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 {
+			case zk.StateDisconnected:
+				s.isStateDisconnected = true
+			case zk.StateConnected:
+				s.isStateDisconnected = false
+			default:
 				s.logger.Printf("zk session event: %+v", event)
 			}
 		case zk.EventNodeChildrenChanged:
diff --git a/discoveryregistry/zookeeper/zk.go b/discoveryregistry/zookeeper/zk.go
index 341d514..b5eb58f 100644
--- a/discoveryregistry/zookeeper/zk.go
+++ b/discoveryregistry/zookeeper/zk.go
@@ -55,11 +55,11 @@ type ZkClient struct {
 	lock    sync.Locker
 	options []grpc.DialOption
 
-	resolvers  map[string]*Resolver
-	localConns map[string][]grpc.ClientConnInterface
-	cancel     context.CancelFunc
-
-	balancerName string
+	resolvers           map[string]*Resolver
+	localConns          map[string][]grpc.ClientConnInterface
+	cancel              context.CancelFunc
+	isStateDisconnected bool
+	balancerName        string
 
 	logger Logger
 }