refactor for readibility

hitless/pool_hook.go

@@ -138,18 +138,23 @@ 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() {
+	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 {
+	if hasPendingHandoff {
-			// Check for empty endpoint first (synchronous check)
+		// Default behavior (pending handoff): pool the connection
-			if conn.GetHandoffEndpoint() == "" {
+		return true, false, nil
-				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
+		// 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
@@ -169,11 +174,6 @@ func (ph *PoolHook) OnPut(ctx context.Context, conn *pool.Conn) (shouldPool bool
 	}
 	return true, false, nil
 }
-		}
-	}
-	// Default: pool the connection
-	return true, false, nil
-}
 
 // ensureWorkerAvailable ensures at least one worker is available to process requests
 // Creates a new worker if needed and under the max limit

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()