From b02cef84485211e0943e73199f46c98da16c7d6b Mon Sep 17 00:00:00 2001
From: Maxime Ginters <maxime.ginters@shopify.com>
Date: Fri, 25 Oct 2024 10:04:01 -0400
Subject: [PATCH] Properly close zk.Conn while dereferencing it

---
 go/vt/topo/zk2topo/zk_conn.go | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/go/vt/topo/zk2topo/zk_conn.go b/go/vt/topo/zk2topo/zk_conn.go
index a0eec8b4340..fae2cb42c59 100644
--- a/go/vt/topo/zk2topo/zk_conn.go
+++ b/go/vt/topo/zk2topo/zk_conn.go
@@ -277,6 +277,7 @@ func (c *ZkConn) withRetry(ctx context.Context, action func(conn *zk.Conn) error
 			c.conn = nil
 		}
 		c.mu.Unlock()
+		conn.Close()
 	}
 	return
 }
@@ -327,13 +328,9 @@ func (c *ZkConn) maybeAddAuth(ctx context.Context) {
 // clears out the connection record.
 func (c *ZkConn) handleSessionEvents(conn *zk.Conn, session <-chan zk.Event) {
 	for event := range session {
-		closeRequired := false
 
 		switch event.State {
-		case zk.StateExpired, zk.StateConnecting:
-			closeRequired = true
-			fallthrough
-		case zk.StateDisconnected:
+		case zk.StateDisconnected, zk.StateExpired, zk.StateConnecting:
 			c.mu.Lock()
 			if c.conn == conn {
 				// The ZkConn still references this
@@ -341,9 +338,7 @@ func (c *ZkConn) handleSessionEvents(conn *zk.Conn, session <-chan zk.Event) {
 				c.conn = nil
 			}
 			c.mu.Unlock()
-			if closeRequired {
-				conn.Close()
-			}
+			conn.Close()
 			log.Infof("zk conn: session for addr %v ended: %v", c.addr, event)
 			return
 		}