From 663a60e47f3bc6fe6d29730635a322442fbada9f Mon Sep 17 00:00:00 2001
From: Nedyalko Dyakov <nedyalko.dyakov@gmail.com>
Date: Fri, 24 Oct 2025 14:13:13 +0300
Subject: [PATCH] correct handling OnPut

---
 internal/pool/conn.go                | 13 ++++++++++---
 internal/pool/pool.go                | 15 +++++++++++----
 maintnotifications/pool_hook_test.go |  6 ++++--
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/internal/pool/conn.go b/internal/pool/conn.go
index 6755eda3..28751d76 100644
--- a/internal/pool/conn.go
+++ b/internal/pool/conn.go
@@ -583,15 +583,22 @@ func (cn *Conn) MarkForHandoff(newEndpoint string, seqID int64) error {
 
 // MarkQueuedForHandoff marks the connection as queued for handoff processing.
 // This makes the connection unusable until handoff completes.
+// This is called from OnPut hook, where the connection is typically in IN_USE state.
+// The pool will preserve the UNUSABLE state and not overwrite it with IDLE.
 func (cn *Conn) MarkQueuedForHandoff() error {
 	// Check if marked for handoff
 	if !cn.ShouldHandoff() {
 		return errors.New("connection was not marked for handoff")
 	}
 
-	// Mark connection as unusable while queued for handoff
-	// Use state machine directly instead of deprecated SetUsable
-	cn.stateMachine.Transition(StateUnusable)
+	// Transition to UNUSABLE from either IN_USE (normal flow) or IDLE (edge cases/tests)
+	// The connection is typically in IN_USE state when OnPut is called (normal Put flow)
+	// But in some edge cases or tests, it might already be in IDLE state
+	// The pool will detect this state change and preserve it (not overwrite with IDLE)
+	_, err := cn.stateMachine.TryTransition([]ConnState{StateInUse, StateIdle}, StateUnusable)
+	if err != nil {
+		return fmt.Errorf("failed to mark connection as unusable: %w", err)
+	}
 	return nil
 }
 
diff --git a/internal/pool/pool.go b/internal/pool/pool.go
index 7ee02999..b4118267 100644
--- a/internal/pool/pool.go
+++ b/internal/pool/pool.go
@@ -663,10 +663,17 @@ func (p *ConnPool) Put(ctx context.Context, cn *Conn) {
 	var shouldCloseConn bool
 
 	if p.cfg.MaxIdleConns == 0 || p.idleConnsLen.Load() < p.cfg.MaxIdleConns {
-		// Transition to IDLE state BEFORE adding to pool
-		// This prevents race condition where another goroutine could acquire
-		// a connection that's still in IN_USE state
-		cn.GetStateMachine().Transition(StateIdle)
+		// Try to transition to IDLE state BEFORE adding to pool
+		// Only transition if connection is still IN_USE (hooks might have changed state)
+		// This prevents:
+		// 1. Race condition where another goroutine could acquire a connection that's still in IN_USE state
+		// 2. Overwriting state changes made by hooks (e.g., IN_USE → UNUSABLE for handoff)
+		currentState, err := cn.GetStateMachine().TryTransition([]ConnState{StateInUse}, StateIdle)
+		if err != nil {
+			// Hook changed the state (e.g., to UNUSABLE for handoff)
+			// Keep the state set by the hook and pool the connection anyway
+			internal.Logger.Printf(ctx, "Connection state changed by hook to %v, pooling as-is", currentState)
+		}
 
 		// unusable conns are expected to become usable at some point (background process is reconnecting them)
 		// put them at the opposite end of the queue
diff --git a/maintnotifications/pool_hook_test.go b/maintnotifications/pool_hook_test.go
index 51e73c3e..a4c36873 100644
--- a/maintnotifications/pool_hook_test.go
+++ b/maintnotifications/pool_hook_test.go
@@ -39,7 +39,9 @@ func (m *mockAddr) String() string  { return m.addr }
 func createMockPoolConnection() *pool.Conn {
 	mockNetConn := &mockNetConn{addr: "test:6379"}
 	conn := pool.NewConn(mockNetConn)
-	conn.SetUsable(true) // Make connection usable for testing
+	conn.SetUsable(true) // Make connection usable for testing (transitions to IDLE)
+	// Simulate real flow: connection is acquired (IDLE → IN_USE) before OnPut is called
+	conn.SetUsed(true) // Transition to IN_USE state
 	return conn
 }
 
@@ -686,7 +688,7 @@ func TestConnectionHook(t *testing.T) {
 			t.Errorf("OnPut should succeed: %v", err)
 		}
 		if !shouldPool || shouldRemove {
-			t.Error("Connection should be pooled after handoff")
+			t.Errorf("Connection should be pooled after handoff (shouldPool=%v, shouldRemove=%v)", shouldPool, shouldRemove)
 		}
 
 		// Wait for handoff to complete