fix bug: close conn

discoveryregistry/zookeeper/discover.go

@@ -35,9 +35,12 @@ var (
 	ErrConnIsNilButLocalNotNil = errors.New("conn is nil, but local is not nil")
 )
 
-func (s *ZkClient) watch() {
-	for {
-		event := <-s.eventChan
+func (s *ZkClient) watch(ctx context.Context) {
+	select {
+	case <-ctx.Done():
+		log.ZDebug(context.Background(), "zk watch ctx done")
+		return
+	case event := <-s.eventChan:
 		log.ZDebug(context.Background(), "zk recv event", "event", event)
 		switch event.Type {
 		case zk.EventSession:

discoveryregistry/zookeeper/zk.go

@@ -15,6 +15,7 @@
 package zookeeper
 
 import (
+	"context"
 	"net"
 	"strconv"
 	"sync"
@@ -56,6 +57,7 @@ type ZkClient struct {
 
 	resolvers  map[string]*Resolver
 	localConns map[string][]grpc.ClientConnInterface
+	cancel     context.CancelFunc
 
 	balancerName string
 
@@ -111,6 +113,8 @@ func NewClient(zkServers []string, zkRoot string, options ...ZkOption) (*ZkClien
 		resolvers:  make(map[string]*Resolver),
 		lock:       &sync.Mutex{},
 	}
+	baseCtx, cancel := context.WithCancel(context.Background())
+	client.cancel = cancel
 	client.ticker = time.NewTicker(defaultFreq)
 	for _, option := range options {
 		option(client)
@@ -138,12 +142,14 @@ func NewClient(zkServers []string, zkRoot string, options ...ZkOption) (*ZkClien
 	}
 	resolver.Register(client)
 	go client.refresh()
-	go client.watch()
+	go client.watch(baseCtx)
 	time.Sleep(time.Millisecond * 50)
 	return client, nil
 }
 
 func (s *ZkClient) CloseZK() {
+	s.cancel()
+	s.ticker.Stop()
 	s.conn.Close()
 }