diff --git a/internal/pool/conn.go b/internal/pool/conn.go index 28751d76..aaf1e5ec 100644 --- a/internal/pool/conn.go +++ b/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 - // (IN_USE means it's usable but currently acquired by someone) - return state == StateIdle || state == StateInUse + // CREATED, IDLE, and IN_USE states are considered usable + // CREATED: new connection, not yet initialized (will be initialized by client) + // 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() { diff --git a/internal/pool/pool.go b/internal/pool/pool.go index b4118267..3fe8bfa5 100644 --- a/internal/pool/pool.go +++ b/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.) - // Put it back in the pool + // Connection is not in a valid state (might be UNUSABLE for handoff/re-auth, INITIALIZING, etc.) + // 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) diff --git a/maintnotifications/errors.go b/maintnotifications/errors.go index 5d335a2c..94a18927 100644 --- a/maintnotifications/errors.go +++ b/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 diff --git a/maintnotifications/pool_hook.go b/maintnotifications/pool_hook.go index 9fd24b4a..9ea0558b 100644 --- a/maintnotifications/pool_hook.go +++ b/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 - // in a handoff state at the moment. - - // Check if connection is usable (not in a handoff state) - // 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 + // This prevents using connections that have received MOVING notifications + if conn.ShouldHandoff() { + return false, ErrConnectionMarkedForHandoffWithState } - // Check if connection is marked for handoff, which means it will be queued for handoff on put. - if conn.ShouldHandoff() { + // Check if connection is usable (not in UNUSABLE or CLOSED state) + // This ensures we don't return connections that are currently being handed off or re-authenticated. + if !conn.IsUsable() { return false, ErrConnectionMarkedForHandoff } diff --git a/maintnotifications/pool_hook_test.go b/maintnotifications/pool_hook_test.go index a4c36873..c21f4221 100644 --- a/maintnotifications/pool_hook_test.go +++ b/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) - if conn.IsUsable() { - t.Error("New connection should not be usable before initialization") + // New connections in CREATED state are usable (they pass OnGet() before initialization) + // The initialization happens AFTER OnGet() in the client code + if !conn.IsUsable() { + t.Error("New connection should be usable (CREATED state is usable)") } - // Simulate initialization by setting usable to true - conn.SetUsable(true) + // Simulate initialization by transitioning to IDLE + 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) }