better errors for tests, hook should work now

internal/pool/conn.go

@@ -187,15 +187,19 @@ func (cn *Conn) CompareAndSwapUsable(old, new bool) bool {
 //
 // A connection is "usable" when it's in a stable state and can be returned to clients.
 // It becomes unusable during:
-//   - Initialization (before first use)
 //   - Handoff operations (network connection replacement)
 //   - Re-authentication (credential updates)
 //   - Other background operations that need exclusive access
+//
+// Note: CREATED state is considered usable because new connections need to pass OnGet() hook
+// before initialization. The initialization happens after OnGet() in the client code.
 func (cn *Conn) IsUsable() bool {
 	state := cn.stateMachine.GetState()
-	// IDLE and IN_USE states are considered usable
+	// CREATED, IDLE, and IN_USE states are considered usable
-	// (IN_USE means it's usable but currently acquired by someone)
+	// CREATED: new connection, not yet initialized (will be initialized by client)
-	return state == StateIdle || state == StateInUse
+	// IDLE: initialized and ready to be acquired
+	// IN_USE: usable but currently acquired by someone
+	return state == StateCreated || state == StateIdle || state == StateInUse
 }
 
 // SetUsable sets the usable flag for the connection (lock-free).
@@ -566,6 +570,8 @@ func (cn *Conn) SetNetConnAndInitConn(ctx context.Context, netConn net.Conn) err
 
 // MarkForHandoff marks the connection for handoff due to MOVING notification.
 // Returns an error if the connection is already marked for handoff.
+// Note: This only sets metadata - the connection state is not changed until OnPut.
+// This allows the current user to finish using the connection before handoff.
 func (cn *Conn) MarkForHandoff(newEndpoint string, seqID int64) error {
 	// Check if already marked for handoff
 	if cn.ShouldHandoff() {

internal/pool/pool.go

@@ -593,8 +593,8 @@ func (p *ConnPool) popIdle() (*Conn, error) {
 			break
 		}
 
-		// Connection is not in a valid state (might be UNUSABLE for re-auth, INITIALIZING, etc.)
+		// Connection is not in a valid state (might be UNUSABLE for handoff/re-auth, INITIALIZING, etc.)
-		// Put it back in the pool
+		// Put it back in the pool and try the next one
 		if p.cfg.PoolFIFO {
 			// FIFO: put at end (will be picked up last since we pop from front)
 			p.idleConns = append(p.idleConns, cn)

maintnotifications/errors.go

@@ -40,9 +40,13 @@ var (
 var (
 	// ErrConnectionMarkedForHandoff is returned when a connection is marked for handoff
 	// and should not be used until the handoff is complete
-	ErrConnectionMarkedForHandoff = errors.New("" + logs.ConnectionMarkedForHandoffErrorMessage)
+	ErrConnectionMarkedForHandoff = errors.New(logs.ConnectionMarkedForHandoffErrorMessage)
+
+	// ErrConnectionMarkedForHandoffWithState is returned when a connection is marked for handoff
+	// and should not be used until the handoff is complete
+	ErrConnectionMarkedForHandoffWithState = errors.New(logs.ConnectionMarkedForHandoffErrorMessage + " with state.")
 	// ErrConnectionInvalidHandoffState is returned when a connection is in an invalid state for handoff
-	ErrConnectionInvalidHandoffState = errors.New("" + logs.ConnectionInvalidHandoffStateErrorMessage)
+	ErrConnectionInvalidHandoffState = errors.New(logs.ConnectionInvalidHandoffStateErrorMessage)
 )
 
 // general errors

maintnotifications/pool_hook.go

@@ -117,17 +117,15 @@ func (ph *PoolHook) ResetCircuitBreakers() {
 
 // OnGet is called when a connection is retrieved from the pool
 func (ph *PoolHook) OnGet(_ context.Context, conn *pool.Conn, _ bool) (accept bool, err error) {
-	// NOTE: There are two conditions to make sure we don't return a connection that should be handed off or is
+	// Check if connection is marked for handoff
-	// in a handoff state at the moment.
+	// This prevents using connections that have received MOVING notifications
-
+	if conn.ShouldHandoff() {
-	// Check if connection is usable (not in a handoff state)
+		return false, ErrConnectionMarkedForHandoffWithState
-	// Should not happen since the pool will not return a connection that is not usable.
-	if !conn.IsUsable() {
-		return false, ErrConnectionMarkedForHandoff
 	}
 
-	// Check if connection is marked for handoff, which means it will be queued for handoff on put.
+	// Check if connection is usable (not in UNUSABLE or CLOSED state)
-	if conn.ShouldHandoff() {
+	// This ensures we don't return connections that are currently being handed off or re-authenticated.
+	if !conn.IsUsable() {
 		return false, ErrConnectionMarkedForHandoff
 	}
 

maintnotifications/pool_hook_test.go

@@ -626,19 +626,20 @@ func TestConnectionHook(t *testing.T) {
 
 		ctx := context.Background()
 
-		// Create a new connection without setting it usable
+		// Create a new connection
 		mockNetConn := &mockNetConn{addr: "test:6379"}
 		conn := pool.NewConn(mockNetConn)
 
-		// Initially, connection should not be usable (not initialized)
+		// New connections in CREATED state are usable (they pass OnGet() before initialization)
-		if conn.IsUsable() {
+		// The initialization happens AFTER OnGet() in the client code
-			t.Error("New connection should not be usable before initialization")
+		if !conn.IsUsable() {
+			t.Error("New connection should be usable (CREATED state is usable)")
 		}
 
-		// Simulate initialization by setting usable to true
+		// Simulate initialization by transitioning to IDLE
-		conn.SetUsable(true)
+		conn.GetStateMachine().Transition(pool.StateIdle)
 		if !conn.IsUsable() {
-			t.Error("Connection should be usable after initialization")
+			t.Error("Connection should be usable after initialization (IDLE state)")
 		}
 
 		// OnGet should succeed for usable connection
@@ -669,12 +670,14 @@ func TestConnectionHook(t *testing.T) {
 			t.Error("Connection should be marked for handoff")
 		}
 
-		// OnGet should fail for connection marked for handoff
+		// OnGet should FAIL for connection marked for handoff
+		// Even though the connection is still in a usable state, the metadata indicates
+		// it should be handed off, so we reject it to prevent using a connection that
+		// will be moved to a different endpoint
 		acceptConn, err = processor.OnGet(ctx, conn, false)
 		if err == nil {
 			t.Error("OnGet should fail for connection marked for handoff")
 		}
-
 		if err != ErrConnectionMarkedForHandoff {
 			t.Errorf("Expected ErrConnectionMarkedForHandoff, got %v", err)
 		}