diff --git a/hitless/pool_hook.go b/hitless/pool_hook.go index 2a31a5d0..86ca9b31 100644 --- a/hitless/pool_hook.go +++ b/hitless/pool_hook.go @@ -138,40 +138,40 @@ func (ph *PoolHook) OnGet(ctx context.Context, conn *pool.Conn, isNewConn bool) // OnPut is called when a connection is returned to the pool func (ph *PoolHook) OnPut(ctx context.Context, conn *pool.Conn) (shouldPool bool, shouldRemove bool, err error) { // first check if we should handoff for faster rejection - if conn.ShouldHandoff() { - // check pending handoff to not queue the same connection twice - _, hasPendingHandoff := ph.pending.Load(conn.GetID()) - if !hasPendingHandoff { - // Check for empty endpoint first (synchronous check) - if conn.GetHandoffEndpoint() == "" { - conn.ClearHandoffState() - } else { - if err := ph.queueHandoff(conn); err != nil { - // Failed to queue handoff, remove the connection - internal.Logger.Printf(ctx, "Failed to queue handoff: %v", err) - return false, true, nil // Don't pool, remove connection, no error to caller - } - - // Check if handoff was already processed by a worker before we can mark it as queued - if !conn.ShouldHandoff() { - // Handoff was already processed - this is normal and the connection should be pooled - return true, false, nil - } - - if err := conn.MarkQueuedForHandoff(); err != nil { - // If marking fails, check if handoff was processed in the meantime - if !conn.ShouldHandoff() { - // Handoff was processed - this is normal, pool the connection - return true, false, nil - } - // Other error - remove the connection - return false, true, nil - } - return true, false, nil - } - } + if !conn.ShouldHandoff() { + // Default behavior (no handoff): pool the connection + return true, false, nil + } + + // check pending handoff to not queue the same connection twice + _, hasPendingHandoff := ph.pending.Load(conn.GetID()) + if hasPendingHandoff { + // Default behavior (pending handoff): pool the connection + return true, false, nil + } + + if err := ph.queueHandoff(conn); err != nil { + // Failed to queue handoff, remove the connection + internal.Logger.Printf(ctx, "Failed to queue handoff: %v", err) + // Don't pool, remove connection, no error to caller + return false, true, nil + } + + // Check if handoff was already processed by a worker before we can mark it as queued + if !conn.ShouldHandoff() { + // Handoff was already processed - this is normal and the connection should be pooled + return true, false, nil + } + + if err := conn.MarkQueuedForHandoff(); err != nil { + // If marking fails, check if handoff was processed in the meantime + if !conn.ShouldHandoff() { + // Handoff was processed - this is normal, pool the connection + return true, false, nil + } + // Other error - remove the connection + return false, true, nil } - // Default: pool the connection return true, false, nil } diff --git a/internal/pool/conn.go b/internal/pool/conn.go index aa2da01a..60875b04 100644 --- a/internal/pool/conn.go +++ b/internal/pool/conn.go @@ -409,14 +409,6 @@ func (cn *Conn) MarkQueuedForHandoff() error { return nil } -// RestoreHandoffState restores the handoff state after a failed handoff (lock-free). -func (cn *Conn) RestoreHandoffState() { - // Restore shouldHandoff flag for retry - cn.shouldHandoffAtomic.Store(true) - // Keep usable=false to prevent the connection from being used until handoff succeeds - cn.setUsable(false) -} - // ShouldHandoff returns true if the connection needs to be handed off (lock-free). func (cn *Conn) ShouldHandoff() bool { return cn.shouldHandoff()