From cfcf37de4e056fbd73d30ea86a8ff06a1586e2da Mon Sep 17 00:00:00 2001
From: Nedyalko Dyakov <nedyalko.dyakov@gmail.com>
Date: Fri, 24 Oct 2025 16:54:23 +0300
Subject: [PATCH] empty endpoint handoff case

---
 maintnotifications/handoff_worker.go |  8 +++++++-
 maintnotifications/pool_hook_test.go | 29 ++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/maintnotifications/handoff_worker.go b/maintnotifications/handoff_worker.go
index e042e4c6..be15747c 100644
--- a/maintnotifications/handoff_worker.go
+++ b/maintnotifications/handoff_worker.go
@@ -349,7 +349,13 @@ func (hwm *handoffWorkerManager) performConnectionHandoff(ctx context.Context, c
 
 	newEndpoint := conn.GetHandoffEndpoint()
 	if newEndpoint == "" {
-		return false, ErrConnectionInvalidHandoffState
+		// Empty endpoint means handoff to current endpoint (reconnect)
+		// Use the current connection's remote address
+		if conn.RemoteAddr() != nil {
+			newEndpoint = conn.RemoteAddr().String()
+		} else {
+			return false, ErrConnectionInvalidHandoffState
+		}
 	}
 
 	// Use circuit breaker to protect against failing endpoints
diff --git a/maintnotifications/pool_hook_test.go b/maintnotifications/pool_hook_test.go
index c94bd67d..2a3abb73 100644
--- a/maintnotifications/pool_hook_test.go
+++ b/maintnotifications/pool_hook_test.go
@@ -235,29 +235,46 @@ func TestConnectionHook(t *testing.T) {
 	})
 
 	t.Run("EmptyEndpoint", func(t *testing.T) {
-		processor := NewPoolHook(baseDialer, "tcp", nil, nil)
+		config := &Config{
+			Mode:              ModeAuto,
+			EndpointType:      EndpointTypeAuto,
+			MaxWorkers:        1,
+			HandoffQueueSize:  10,
+			MaxHandoffRetries: 3,
+		}
+		processor := NewPoolHook(baseDialer, "tcp", config, nil)
+		defer processor.Shutdown(context.Background())
+
 		conn := createMockPoolConnection()
 		if err := conn.MarkForHandoff("", 12345); err != nil { // Empty endpoint
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
 		}
 
+		// Set a mock initialization function
+		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			return nil
+		})
+
 		ctx := context.Background()
 		shouldPool, shouldRemove, err := processor.OnPut(ctx, conn)
 		if err != nil {
 			t.Errorf("OnPut should not error with empty endpoint: %v", err)
 		}
 
-		// Should pool the connection (empty endpoint clears state)
+		// Should pool the connection (empty endpoint triggers handoff to current endpoint)
 		if !shouldPool {
-			t.Error("Connection should be pooled after clearing empty endpoint")
+			t.Error("Connection should be pooled when handoff is queued")
 		}
 		if shouldRemove {
-			t.Error("Connection should not be removed after clearing empty endpoint")
+			t.Error("Connection should not be removed when handoff is queued")
 		}
 
-		// State should be cleared
+		// Wait for handoff to complete
+		time.Sleep(100 * time.Millisecond)
+
+		// After handoff completes, state should be cleared
 		if conn.ShouldHandoff() {
-			t.Error("Connection should not be marked for handoff after clearing empty endpoint")
+			t.Error("Connection should not be marked for handoff after handoff completes")
 		}
 	})