pubsub fixes

Author: ndyakov

async_handoff_integration_test.go

@@ -71,20 +71,27 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 			t.Fatalf("Failed to get connection: %v", err)
 		}
 
-		// Set initialization function
+		// Set initialization function with a small delay to ensure handoff is pending
 		initConnCalled := false
 		initConnFunc := func(ctx context.Context, cn *pool.Conn) error {
+			time.Sleep(50 * time.Millisecond) // Add delay to keep handoff pending
 			initConnCalled = true
 			return nil
 		}
 		conn.SetInitConnFunc(initConnFunc)
 
 		// Mark connection for handoff
-		conn.MarkForHandoff("new-endpoint:6379", 12345)
+		err = conn.MarkForHandoff("new-endpoint:6379", 12345)
+		if err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		// Return connection to pool - this should queue handoff
 		testPool.Put(ctx, conn)
 
+		// Give the on-demand worker a moment to start processing
+		time.Sleep(10 * time.Millisecond)
+
 		// Verify handoff was queued
 		if !processor.IsHandoffPending(conn) {
 			t.Error("Handoff should be queued in pending map")
@@ -303,14 +310,19 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
 		}
 
-		// Set a mock initialization function
+		// Set a mock initialization function with delay to ensure handoff is pending
 		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			time.Sleep(50 * time.Millisecond) // Add delay to keep handoff pending
 			return nil
 		})
 
 		testPool.Put(ctx, conn)
 
-		// Verify handoff was queued
+		// Give the on-demand worker a moment to start and begin processing
+		// The handoff should be pending because the slowDialer takes 100ms
+		time.Sleep(10 * time.Millisecond)
+
+		// Verify handoff was queued and is being processed
 		if !processor.IsHandoffPending(conn) {
 			t.Error("Handoff should be queued in pending map")
 		}

internal/pool/pool.go

@@ -3,7 +3,6 @@ package pool
 import (
 	"context"
 	"errors"
-	"fmt"
 	"log"
 	"net"
 	"sync"
@@ -156,7 +155,6 @@ func (p *ConnPool) AddPoolHook(hook PoolHook) {
 		p.initializeHooks()
 	}
 	p.hookManager.AddHook(hook)
-	p.hookManager = nil
 }
 
 // RemovePoolHook removes a pool hook from the pool.
@@ -300,7 +298,6 @@ func (p *ConnPool) dialConn(ctx context.Context, pooled bool) (*Conn, error) {
 
 	cn := NewConnWithBufferSize(netConn, p.cfg.ReadBufferSize, p.cfg.WriteBufferSize)
 	cn.pooled = pooled
-	fmt.Printf("New conn %d, pooled: %v\n", cn.GetID(), cn.pooled)
 	if p.cfg.ConnMaxLifetime > 0 {
 		cn.expiresAt = time.Now().Add(p.cfg.ConnMaxLifetime)
 	} else {
@@ -596,7 +593,6 @@ func (p *ConnPool) Put(ctx context.Context, cn *Conn) {
 }
 
 func (p *ConnPool) Remove(_ context.Context, cn *Conn, reason error) {
-	internal.Logger.Printf(context.Background(), "Removing connection %d from pool: %v", cn.GetID(), reason)
 	p.removeConnWithLock(cn)
 	p.freeTurn()
 	_ = p.closeConn(cn)

osscluster.go

@@ -1845,7 +1845,6 @@ func (c *ClusterClient) pubSub() *PubSub {
 	var node *clusterNode
 	pubsub := &PubSub{
 		opt: c.opt.clientOptions(),
-
 		newConn: func(ctx context.Context, addr string, channels []string) (*pool.Conn, error) {
 			if node != nil {
 				panic("node != nil")
@@ -1861,14 +1860,19 @@ func (c *ClusterClient) pubSub() *PubSub {
 			if err != nil {
 				return nil, err
 			}
-
-			cn, err := node.Client.newConn(context.TODO())
+			cn, err := node.Client.pubSubPool.NewConn(ctx, node.Client.opt.Network, node.Client.opt.Addr, channels)
 			if err != nil {
 				node = nil
-
 				return nil, err
 			}
-
+			// will return nil if already initialized
+			err = node.Client.initConn(ctx, cn)
+			if err != nil {
+				_ = cn.Close()
+				node = nil
+				return nil, err
+			}
+			node.Client.pubSubPool.TrackConn(cn)
 			return cn, nil
 		},
 		closeConn: func(cn *pool.Conn) error {

sentinel.go

@@ -475,6 +475,7 @@ func NewFailoverClient(failoverOpt *FailoverOptions) *Client {
 	rdb.pushProcessor = initializePushProcessor(opt)
 
 	rdb.connPool = newConnPool(opt, rdb.dialHook)
+	rdb.pubSubPool = newPubSubPool(opt, rdb.dialHook)
 
 	rdb.onClose = rdb.wrappedOnClose(failover.Close)
 
@@ -551,6 +552,7 @@ func NewSentinelClient(opt *Options) *SentinelClient {
 		process: c.baseClient.process,
 	})
 	c.connPool = newConnPool(opt, c.dialHook)
+	c.pubSubPool = newPubSubPool(opt, c.dialHook)
 
 	return c
 }