feat(pool): Improve success rate of new connections (#3518)

* async create conn

* update default values and testcase

* fix comments

* fix data race

* remove context.WithoutCancel, which is a function introduced in Go 1.21

* fix TestDialerRetryConfiguration/DefaultDialerRetries, because tryDial are likely done in async flow

* change to share failed to delivery connection to other waiting

* remove chinese comment

* fix: optimize WantConnQueue benchmarks to prevent memory exhaustion

- Fix BenchmarkWantConnQueue_Dequeue timeout issue by limiting pre-population
- Use object pooling in BenchmarkWantConnQueue_Enqueue to reduce allocations
- Optimize BenchmarkWantConnQueue_EnqueueDequeue with reusable wantConn pool
- Prevent GitHub Actions benchmark failures due to excessive memory usage

Before: BenchmarkWantConnQueue_Dequeue ran for 11+ minutes and was killed
After: All benchmarks complete in ~8 seconds with consistent performance

* format

* fix turn leaks

---------

Co-authored-by: Nedyalko Dyakov <[email protected]>
Co-authored-by: Hristo Temelski <[email protected]>

Author: 3 people

async_handoff_integration_test.go

@@ -53,8 +53,9 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 			Dialer: func(ctx context.Context) (net.Conn, error) {
 				return &mockNetConn{addr: "original:6379"}, nil
 			},
-			PoolSize:    int32(5),
-			PoolTimeout: time.Second,
+			PoolSize:           int32(5),
+			MaxConcurrentDials: 5,
+			PoolTimeout:        time.Second,
 		})
 
 		// Add the hook to the pool after creation
@@ -153,8 +154,9 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 				return &mockNetConn{addr: "original:6379"}, nil
 			},
 
-			PoolSize:    int32(10),
-			PoolTimeout: time.Second,
+			PoolSize:           int32(10),
+			MaxConcurrentDials: 10,
+			PoolTimeout:        time.Second,
 		})
 		defer testPool.Close()
 
@@ -225,8 +227,9 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 				return &mockNetConn{addr: "original:6379"}, nil
 			},
 
-			PoolSize:    int32(3),
-			PoolTimeout: time.Second,
+			PoolSize:           int32(3),
+			MaxConcurrentDials: 3,
+			PoolTimeout:        time.Second,
 		})
 		defer testPool.Close()
 
@@ -288,8 +291,9 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 				return &mockNetConn{addr: "original:6379"}, nil
 			},
 
-			PoolSize:    int32(2),
-			PoolTimeout: time.Second,
+			PoolSize:           int32(2),
+			MaxConcurrentDials: 2,
+			PoolTimeout:        time.Second,
 		})
 		defer testPool.Close()
 

internal/pool/bench_test.go

@@ -31,11 +31,12 @@ func BenchmarkPoolGetPut(b *testing.B) {
 	for _, bm := range benchmarks {
 		b.Run(bm.String(), func(b *testing.B) {
 			connPool := pool.NewConnPool(&pool.Options{
-				Dialer:          dummyDialer,
-				PoolSize:        int32(bm.poolSize),
-				PoolTimeout:     time.Second,
-				DialTimeout:     1 * time.Second,
-				ConnMaxIdleTime: time.Hour,
+				Dialer:             dummyDialer,
+				PoolSize:           int32(bm.poolSize),
+				MaxConcurrentDials: bm.poolSize,
+				PoolTimeout:        time.Second,
+				DialTimeout:        1 * time.Second,
+				ConnMaxIdleTime:    time.Hour,
 			})
 
 			b.ResetTimer()
@@ -75,11 +76,12 @@ func BenchmarkPoolGetRemove(b *testing.B) {
 	for _, bm := range benchmarks {
 		b.Run(bm.String(), func(b *testing.B) {
 			connPool := pool.NewConnPool(&pool.Options{
-				Dialer:          dummyDialer,
-				PoolSize:        int32(bm.poolSize),
-				PoolTimeout:     time.Second,
-				DialTimeout:     1 * time.Second,
-				ConnMaxIdleTime: time.Hour,
+				Dialer:             dummyDialer,
+				PoolSize:           int32(bm.poolSize),
+				MaxConcurrentDials: bm.poolSize,
+				PoolTimeout:        time.Second,
+				DialTimeout:        1 * time.Second,
+				ConnMaxIdleTime:    time.Hour,
 			})
 
 			b.ResetTimer()

internal/pool/buffer_size_test.go

@@ -24,9 +24,10 @@ var _ = Describe("Buffer Size Configuration", func() {
 
 	It("should use default buffer sizes when not specified", func() {
 		connPool = pool.NewConnPool(&pool.Options{
-			Dialer:      dummyDialer,
-			PoolSize:    int32(1),
-			PoolTimeout: 1000,
+			Dialer:             dummyDialer,
+			PoolSize:           int32(1),
+			MaxConcurrentDials: 1,
+			PoolTimeout:        1000,
 		})
 
 		cn, err := connPool.NewConn(ctx)
@@ -46,11 +47,12 @@ var _ = Describe("Buffer Size Configuration", func() {
 		customWriteSize := 64 * 1024 // 64KB
 
 		connPool = pool.NewConnPool(&pool.Options{
-			Dialer:          dummyDialer,
-			PoolSize:        int32(1),
-			PoolTimeout:     1000,
-			ReadBufferSize:  customReadSize,
-			WriteBufferSize: customWriteSize,
+			Dialer:             dummyDialer,
+			PoolSize:           int32(1),
+			MaxConcurrentDials: 1,
+			PoolTimeout:        1000,
+			ReadBufferSize:     customReadSize,
+			WriteBufferSize:    customWriteSize,
 		})
 
 		cn, err := connPool.NewConn(ctx)
@@ -67,11 +69,12 @@ var _ = Describe("Buffer Size Configuration", func() {
 
 	It("should handle zero buffer sizes by using defaults", func() {
 		connPool = pool.NewConnPool(&pool.Options{
-			Dialer:          dummyDialer,
-			PoolSize:        int32(1),
-			PoolTimeout:     1000,
-			ReadBufferSize:  0, // Should use default
-			WriteBufferSize: 0, // Should use default
+			Dialer:             dummyDialer,
+			PoolSize:           int32(1),
+			MaxConcurrentDials: 1,
+			PoolTimeout:        1000,
+			ReadBufferSize:     0, // Should use default
+			WriteBufferSize:    0, // Should use default
 		})
 
 		cn, err := connPool.NewConn(ctx)
@@ -103,9 +106,10 @@ var _ = Describe("Buffer Size Configuration", func() {
 		// Test the scenario where someone creates a pool directly (like in tests)
 		// without setting ReadBufferSize and WriteBufferSize
 		connPool = pool.NewConnPool(&pool.Options{
-			Dialer:      dummyDialer,
-			PoolSize:    int32(1),
-			PoolTimeout: 1000,
+			Dialer:             dummyDialer,
+			PoolSize:           int32(1),
+			MaxConcurrentDials: 1,
+			PoolTimeout:        1000,
 			// ReadBufferSize and WriteBufferSize are not set (will be 0)
 		})
 

internal/pool/hooks_test.go

@@ -191,8 +191,9 @@ func TestPoolWithHooks(t *testing.T) {
 		Dialer: func(ctx context.Context) (net.Conn, error) {
 			return &net.TCPConn{}, nil // Mock connection
 		},
-		PoolSize:    1,
-		DialTimeout: time.Second,
+		PoolSize:           1,
+		MaxConcurrentDials: 1,
+		DialTimeout:        time.Second,
 	}
 
 	pool := NewConnPool(opt)

internal/pool/pool.go

@@ -98,6 +98,7 @@ type Options struct {
 
 	PoolFIFO                 bool
 	PoolSize                 int32
+	MaxConcurrentDials       int
 	DialTimeout              time.Duration
 	PoolTimeout              time.Duration
 	MinIdleConns             int32
@@ -126,7 +127,9 @@ type ConnPool struct {
 	dialErrorsNum uint32 // atomic
 	lastDialError atomic.Value
 
-	queue chan struct{}
+	queue           chan struct{}
+	dialsInProgress chan struct{}
+	dialsQueue      *wantConnQueue
 
 	connsMu   sync.Mutex
 	conns     map[uint64]*Conn
@@ -152,9 +155,11 @@ func NewConnPool(opt *Options) *ConnPool {
 	p := &ConnPool{
 		cfg: opt,
 
-		queue:     make(chan struct{}, opt.PoolSize),
-		conns:     make(map[uint64]*Conn),
-		idleConns: make([]*Conn, 0, opt.PoolSize),
+		queue:           make(chan struct{}, opt.PoolSize),
+		conns:           make(map[uint64]*Conn),
+		dialsInProgress: make(chan struct{}, opt.MaxConcurrentDials),
+		dialsQueue:      newWantConnQueue(),
+		idleConns:       make([]*Conn, 0, opt.PoolSize),
 	}
 
 	// Only create MinIdleConns if explicitly requested (> 0)
@@ -233,6 +238,7 @@ func (p *ConnPool) checkMinIdleConns() {
 			return
 		}
 	}
+
 }
 
 func (p *ConnPool) addIdleConn() error {
@@ -491,9 +497,8 @@ func (p *ConnPool) getConn(ctx context.Context) (*Conn, error) {
 
 	atomic.AddUint32(&p.stats.Misses, 1)
 
-	newcn, err := p.newConn(ctx, true)
+	newcn, err := p.queuedNewConn(ctx)
 	if err != nil {
-		p.freeTurn()
 		return nil, err
 	}
 
@@ -512,6 +517,99 @@ func (p *ConnPool) getConn(ctx context.Context) (*Conn, error) {
 	return newcn, nil
 }
 
+func (p *ConnPool) queuedNewConn(ctx context.Context) (*Conn, error) {
+	select {
+	case p.dialsInProgress <- struct{}{}:
+		// Got permission, proceed to create connection
+	case <-ctx.Done():
+		p.freeTurn()
+		return nil, ctx.Err()
+	}
+
+	dialCtx, cancel := context.WithTimeout(context.Background(), p.cfg.DialTimeout)
+
+	w := &wantConn{
+		ctx:       dialCtx,
+		cancelCtx: cancel,
+		result:    make(chan wantConnResult, 1),
+	}
+	var err error
+	defer func() {
+		if err != nil {
+			if cn := w.cancel(); cn != nil {
+				p.putIdleConn(ctx, cn)
+				p.freeTurn()
+			}
+		}
+	}()
+
+	p.dialsQueue.enqueue(w)
+
+	go func(w *wantConn) {
+		var freeTurnCalled bool
+		defer func() {
+			if err := recover(); err != nil {
+				if !freeTurnCalled {
+					p.freeTurn()
+				}
+				internal.Logger.Printf(context.Background(), "queuedNewConn panic: %+v", err)
+			}
+		}()
+
+		defer w.cancelCtx()
+		defer func() { <-p.dialsInProgress }() // Release connection creation permission
+
+		dialCtx := w.getCtxForDial()
+		cn, cnErr := p.newConn(dialCtx, true)
+		delivered := w.tryDeliver(cn, cnErr)
+		if cnErr == nil && delivered {
+			return
+		} else if cnErr == nil && !delivered {
+			p.putIdleConn(dialCtx, cn)
+			p.freeTurn()
+			freeTurnCalled = true
+		} else {
+			p.freeTurn()
+			freeTurnCalled = true
+		}
+	}(w)
+
+	select {
+	case <-ctx.Done():
+		err = ctx.Err()
+		return nil, err
+	case result := <-w.result:
+		err = result.err
+		return result.cn, err
+	}
+}
+
+func (p *ConnPool) putIdleConn(ctx context.Context, cn *Conn) {
+	for {
+		w, ok := p.dialsQueue.dequeue()
+		if !ok {
+			break
+		}
+		if w.tryDeliver(cn, nil) {
+			return
+		}
+	}
+
+	cn.SetUsable(true)
+
+	p.connsMu.Lock()
+	defer p.connsMu.Unlock()
+
+	if p.closed() {
+		_ = cn.Close()
+		return
+	}
+
+	// poolSize is increased in newConn
+	p.idleConns = append(p.idleConns, cn)
+	p.idleConnsLen.Add(1)
+}
+
 func (p *ConnPool) waitTurn(ctx context.Context) error {
 	select {
 	case <-ctx.Done():