revert some optimizations

Author: ndyakov

hitless/notification_handler.go

@@ -91,7 +91,11 @@ func (snh *NotificationHandler) handleMoving(ctx context.Context, handlerCtx pus
 	}
 
 	// Mark the connection for handoff
-	poolConn.MarkForHandoff(newEndpoint, seqID)
+	if err := poolConn.MarkForHandoff(newEndpoint, seqID); err != nil {
+		// Connection is already marked for handoff, which is acceptable
+		// This can happen if multiple MOVING notifications are received for the same connection
+		return nil
+	}
 
 	// Optionally track in hitless manager for monitoring/debugging
 	if snh.manager != nil {

hitless/redis_connection_processor.go

@@ -517,7 +517,6 @@ func (rcp *RedisConnectionProcessor) performConnectionHandoffWithPool(ctx contex
 			}
 		} else {
 			cn.Close()
-			newNetConn.Close()
 			internal.Logger.Printf(ctx,
 				"hitless: no pool provided for connection %d, cannot remove due to handoff initialization failure: %v",
 				cn.GetID(), err)

hitless/redis_connection_processor_test.go

@@ -111,7 +111,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		defer processor.Shutdown(context.Background())
 
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("new-endpoint:6379", 12345)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 12345); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		// Set a mock initialization function
 		initConnCalled := false
@@ -182,7 +184,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 	t.Run("EmptyEndpoint", func(t *testing.T) {
 		processor := NewRedisConnectionProcessor(3, baseDialer, nil, nil)
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("", 12345) // Empty endpoint
+		if err := conn.MarkForHandoff("", 12345); err != nil { // Empty endpoint
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		ctx := context.Background()
 		shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
@@ -214,7 +218,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		defer processor.Shutdown(context.Background())
 
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("new-endpoint:6379", 12345)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 12345); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		ctx := context.Background()
 		shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
@@ -359,7 +365,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		connections := make([]*pool.Conn, 5)
 		for i := 0; i < 5; i++ {
 			connections[i] = createMockPoolConnection()
-			connections[i].MarkForHandoff("new-endpoint:6379", int64(i))
+			if err := connections[i].MarkForHandoff("new-endpoint:6379", int64(i)); err != nil {
+				t.Fatalf("Failed to mark connection %d for handoff: %v", i, err)
+			}
 		}
 
 		ctx := context.Background()
@@ -438,7 +446,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 
 		// Create a connection and trigger handoff
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("new-endpoint:6379", 1)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 1); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		// Process the connection to trigger handoff
 		shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
@@ -516,7 +526,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		}
 
 		// Mark connection for handoff
-		conn.MarkForHandoff("new-endpoint:6379", 1)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 1); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		// Connection should no longer be usable
 		if conn.IsUsable() {
@@ -583,7 +595,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		// Fill part of the queue
 		for i := 0; i < 10; i++ {
 			conn := createMockPoolConnection()
-			conn.MarkForHandoff("new-endpoint:6379", int64(i+1))
+			if err := conn.MarkForHandoff("new-endpoint:6379", int64(i+1)); err != nil {
+				t.Fatalf("Failed to mark connection %d for handoff: %v", i, err)
+			}
 
 			shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
 			if err != nil {
@@ -634,7 +648,9 @@ func TestRedisConnectionProcessor(t *testing.T) {
 
 		// Create a connection and mark it for handoff
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("new-endpoint:6379", 1)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 1); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
 
 		// Set a failing initialization function
 		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
@@ -712,4 +728,31 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Error("Relaxed timeout should be automatically cleared after post-handoff duration")
 		}
 	})
+
+	t.Run("MarkForHandoff returns error when already marked", func(t *testing.T) {
+		conn := createMockPoolConnection()
+
+		// First mark should succeed
+		if err := conn.MarkForHandoff("new-endpoint:6379", 1); err != nil {
+			t.Fatalf("First MarkForHandoff should succeed: %v", err)
+		}
+
+		// Second mark should fail
+		if err := conn.MarkForHandoff("another-endpoint:6379", 2); err == nil {
+			t.Fatal("Second MarkForHandoff should return error")
+		} else if err.Error() != "connection is already marked for handoff" {
+			t.Fatalf("Expected specific error message, got: %v", err)
+		}
+
+		// Verify original handoff data is preserved
+		if !conn.ShouldHandoff() {
+			t.Fatal("Connection should still be marked for handoff")
+		}
+		if conn.GetHandoffEndpoint() != "new-endpoint:6379" {
+			t.Fatalf("Expected original endpoint, got: %s", conn.GetHandoffEndpoint())
+		}
+		if conn.GetMovingSeqID() != 1 {
+			t.Fatalf("Expected original sequence ID, got: %d", conn.GetMovingSeqID())
+		}
+	})
 }

internal/pool/conn.go

@@ -3,6 +3,7 @@ package pool
 import (
 	"bufio"
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"sync/atomic"
@@ -162,11 +163,6 @@ func (cn *Conn) setNewEndpoint(endpoint string) {
 	cn.newEndpointAtomic.Store(endpoint)
 }
 
-// getHandoffRetries returns the retry count atomically (lock-free).
-func (cn *Conn) getHandoffRetries() int {
-	return int(atomic.LoadInt32(&cn.handoffRetriesAtomic))
-}
-
 // setHandoffRetries sets the retry count atomically (lock-free).
 func (cn *Conn) setHandoffRetries(retries int) {
 	atomic.StoreInt32(&cn.handoffRetriesAtomic, int32(retries))
@@ -343,11 +339,18 @@ func (cn *Conn) SetNetConnWithInitConn(ctx context.Context, netConn net.Conn) er
 }
 
 // MarkForHandoff marks the connection for handoff due to MOVING notification (lock-free).
-func (cn *Conn) MarkForHandoff(newEndpoint string, seqID int64) {
+// Returns an error if the connection is already marked for handoff.
+func (cn *Conn) MarkForHandoff(newEndpoint string, seqID int64) error {
+	// Check if connection is already marked for handoff
+	if cn.shouldHandoff() {
+		return errors.New("connection is already marked for handoff")
+	}
+
 	cn.setShouldHandoff(true)
 	cn.setNewEndpoint(newEndpoint)
 	cn.setMovingSeqID(seqID)
 	cn.setUsable(false) // Connection is not safe to use until handoff completes
+	return nil
 }
 
 // ShouldHandoff returns true if the connection needs to be handed off (lock-free).

internal/pool/pool.go

@@ -108,9 +108,6 @@ type ConnPool struct {
 	stats          Stats
 	waitDurationNs atomic.Int64
 
-	// Atomic flag to prevent concurrent minimum idle connection checks
-	checkingMinIdle uint32 // atomic: 1 if checking is in progress, 0 otherwise
-
 	_closed uint32 // atomic
 }
 
@@ -128,38 +125,14 @@ func NewConnPool(opt *Options) *ConnPool {
 	// Only create MinIdleConns if explicitly requested (> 0)
 	// This avoids creating connections during pool initialization for tests
 	if opt.MinIdleConns > 0 {
-		p.checkMinIdleConnsAsync()
-	}
-
-	return p
-}
-
-// checkMinIdleConnsAsync asynchronously checks and maintains minimum idle connections.
-// Uses an atomic flag to prevent concurrent checks - only creates a goroutine when needed.
-func (p *ConnPool) checkMinIdleConnsAsync() {
-	// Fast path: if MinIdleConns is 0, no need to check
-	if p.cfg.MinIdleConns == 0 {
-		return
-	}
-
-	// Check if idle connections are already being checked
-	if !atomic.CompareAndSwapUint32(&p.checkingMinIdle, 0, 1) {
-		// Already checking, return early to avoid duplicate work
-		return
-	}
-
-	// Start checking in a goroutine to avoid blocking the caller
-	go func() {
-		defer atomic.StoreUint32(&p.checkingMinIdle, 0) // Reset the flag when done
-
-		// Check and maintain minimum idle connections
 		p.connsMu.Lock()
 		p.checkMinIdleConns()
 		p.connsMu.Unlock()
-	}()
+	}
+
+	return p
 }
 
-// checkMinIdleConns is not thread-safe. Should be called with connsMu locked.
 func (p *ConnPool) checkMinIdleConns() {
 	if p.cfg.MinIdleConns == 0 {
 		return
@@ -172,19 +145,23 @@ func (p *ConnPool) checkMinIdleConns() {
 		case p.queue <- struct{}{}:
 			p.poolSize++
 			p.idleConnsLen++
-			err := p.addIdleConn()
-			if err != nil && err != ErrClosed {
-				p.poolSize--
-				p.idleConnsLen--
-			}
-			p.freeTurn()
+			go func() {
+				err := p.addIdleConn()
+				if err != nil && err != ErrClosed {
+					p.connsMu.Lock()
+					p.poolSize--
+					p.idleConnsLen--
+					p.connsMu.Unlock()
+				}
+
+				p.freeTurn()
+			}()
 		default:
 			return
 		}
 	}
 }
 
-// addIdleConn is not thread-safe. Should be called with connsMu locked.
 func (p *ConnPool) addIdleConn() error {
 	ctx, cancel := context.WithTimeout(context.Background(), p.cfg.DialTimeout)
 	defer cancel()
@@ -194,6 +171,9 @@ func (p *ConnPool) addIdleConn() error {
 		return err
 	}
 
+	p.connsMu.Lock()
+	defer p.connsMu.Unlock()
+
 	// It is not allowed to add new connections to the closed connection pool.
 	if p.closed() {
 		_ = cn.Close()
@@ -334,6 +314,7 @@ func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
 	}
 
 	tries := 0
+	now := time.Now()
 	for {
 		if tries > 10 {
 			log.Printf("redis: connection pool: failed to get a connection after %d tries", tries)
@@ -353,7 +334,7 @@ func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
 			break
 		}
 
-		if !p.isHealthyConn(cn) {
+		if !p.isHealthyConn(cn, now) {
 			_ = p.CloseConn(cn)
 			continue
 		}
@@ -475,12 +456,7 @@ func (p *ConnPool) popIdle() (*Conn, error) {
 		return nil, nil
 	}
 
-	// Asynchronously check minimum idle connections (only if MinIdleConns > 0)
-	// This avoids blocking the Get() operation while still maintaining pool health
-	if p.cfg.MinIdleConns > 0 {
-		p.checkMinIdleConnsAsync()
-	}
-
+	p.checkMinIdleConns()
 	return cn, nil
 }
 
@@ -550,18 +526,9 @@ func (p *ConnPool) CloseConn(cn *Conn) error {
 }
 
 func (p *ConnPool) removeConnWithLock(cn *Conn) {
-	var shouldCheckMinIdle bool
-
 	p.connsMu.Lock()
-	oldPoolSize := p.poolSize
+	defer p.connsMu.Unlock()
 	p.removeConn(cn)
-	shouldCheckMinIdle = cn.pooled && p.poolSize < oldPoolSize // Connection was removed from pool
-	p.connsMu.Unlock()
-
-	// Check minimum idle connections asynchronously after releasing the lock
-	if shouldCheckMinIdle {
-		p.checkMinIdleConnsAsync()
-	}
 }
 
 func (p *ConnPool) removeConn(cn *Conn) {
@@ -570,6 +537,8 @@ func (p *ConnPool) removeConn(cn *Conn) {
 			p.conns = append(p.conns[:i], p.conns[i+1:]...)
 			if cn.pooled {
 				p.poolSize--
+				// Immediately check for minimum idle connections when a pooled connection is removed
+				p.checkMinIdleConns()
 			}
 			break
 		}
@@ -653,9 +622,7 @@ func (p *ConnPool) Close() error {
 	return firstErr
 }
 
-func (p *ConnPool) isHealthyConn(cn *Conn) bool {
-	now := time.Now()
-
+func (p *ConnPool) isHealthyConn(cn *Conn, now time.Time) bool {
 	// slight optimization, check expiresAt first.
 	if cn.expiresAt.Before(now) {
 		return false
@@ -667,7 +634,7 @@ func (p *ConnPool) isHealthyConn(cn *Conn) bool {
 
 	// Check basic connection health
 	// Use GetNetConn() to safely access netConn and avoid data races
-	if err := connCheck(cn.GetNetConn()); err != nil {
+	if err := connCheck(cn.getNetConn()); err != nil {
 		return false
 	}
 

internal/pool/pool.test

@@ -390,3 +390,21 @@ func BenchmarkPoolStatsCollection(b *testing.B) {
 		_ = stats.Hits + stats.Misses + stats.Timeouts // Use the stats to prevent optimization
 	}
 }
+
+// BenchmarkPoolHighContention tests pool performance under high contention
+func BenchmarkPoolHighContention(b *testing.B) {
+	ctx := context.Background()
+	client := benchmarkClient(32)
+	defer client.Close()
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			// High contention Get/Put operations
+			pubsub := client.Subscribe(ctx, "test-channel")
+			pubsub.Close()
+		}
+	})
+}