what if those conns are tracked and closed on Close?

Author: ndyakov

internal/pool/pool.go

@@ -68,6 +68,9 @@ type Pooler interface {
 	AddPoolHook(hook PoolHook)
 	RemovePoolHook(hook PoolHook)
 
+	TrackConn(*Conn)
+	UntrackConn(*Conn)
+
 	Close() error
 }
 
@@ -736,3 +739,22 @@ func (p *ConnPool) isHealthyConn(cn *Conn, now time.Time) bool {
 	}
 	return true
 }
+
+func (p *ConnPool) TrackConn(cn *Conn) {
+	p.connsMu.Lock()
+	p.poolSize.Add(1)
+	p.conns = append(p.conns, cn)
+	p.connsMu.Unlock()
+}
+
+func (p *ConnPool) UntrackConn(cn *Conn) {
+	p.connsMu.Lock()
+	for i, c := range p.conns {
+		if c == cn {
+			p.conns = append(p.conns[:i], p.conns[i+1:]...)
+			p.poolSize.Add(-1)
+			break
+		}
+	}
+	p.connsMu.Unlock()
+}

internal/pool/pool_single.go

@@ -62,3 +62,7 @@ func (p *SingleConnPool) Stats() *Stats {
 func (p *SingleConnPool) AddPoolHook(hook PoolHook) {}
 
 func (p *SingleConnPool) RemovePoolHook(hook PoolHook) {}
+
+func (p *SingleConnPool) TrackConn(cn *Conn) {}
+
+func (p *SingleConnPool) UntrackConn(cn *Conn) {}

internal/pool/pool_sticky.go

@@ -203,3 +203,7 @@ func (p *StickyConnPool) Stats() *Stats {
 func (p *StickyConnPool) AddPoolHook(hook PoolHook) {}
 
 func (p *StickyConnPool) RemovePoolHook(hook PoolHook) {}
+
+func (p *StickyConnPool) TrackConn(cn *Conn) {}
+
+func (p *StickyConnPool) UntrackConn(cn *Conn) {}

redis.go

@@ -1071,10 +1071,12 @@ func (c *Client) pubSub() *PubSub {
 				return nil, err
 			}
 
+			c.connPool.TrackConn(cn)
 			return cn, nil
 		},
 		closeConn: func(cn *pool.Conn) error {
 			_ = cn.Close()
+			c.connPool.UntrackConn(cn)
 			return nil
 		},
 		pushProcessor: c.pushProcessor,