async create conn

Author: cyningsun

internal/pool/pool.go

@@ -72,6 +72,8 @@ type Options struct {
 	ConnMaxIdleTime time.Duration
 	ConnMaxLifetime time.Duration
 
+	MaxConcurrentDials int
+
 	ReadBufferSize  int
 	WriteBufferSize int
 }
@@ -80,13 +82,43 @@ type lastDialErrorWrap struct {
 	err error
 }
 
+type wantConn struct {
+	mu        sync.Mutex      // protects ctx, done and sending of the result
+	ctx       context.Context // context for dial, cleared after delivered or canceled
+	cancelCtx context.CancelFunc
+	done      bool                // true after delivered or canceled
+	result    chan wantConnResult // channel to deliver connection or error
+}
+
+func (w *wantConn) tryDeliver(cn *Conn, err error) bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.done {
+		return false
+	}
+
+	w.done = true
+	w.ctx = nil
+
+	w.result <- wantConnResult{cn: cn, err: err}
+	close(w.result)
+
+	return true
+}
+
+type wantConnResult struct {
+	cn  *Conn
+	err error
+}
+
 type ConnPool struct {
 	cfg *Options
 
 	dialErrorsNum uint32 // atomic
 	lastDialError atomic.Value
 
-	queue chan struct{}
+	queue           chan struct{}
+	dialsInProgress chan struct{}
 
 	connsMu   sync.Mutex
 	conns     []*Conn
@@ -107,7 +139,9 @@ func NewConnPool(opt *Options) *ConnPool {
 	p := &ConnPool{
 		cfg: opt,
 
-		queue:     make(chan struct{}, opt.PoolSize),
+		queue:           make(chan struct{}, opt.PoolSize),
+		dialsInProgress: make(chan struct{}, opt.MaxConcurrentDials),
+
 		conns:     make([]*Conn, 0, opt.PoolSize),
 		idleConns: make([]*Conn, 0, opt.PoolSize),
 	}
@@ -316,13 +350,46 @@ func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
 
 	atomic.AddUint32(&p.stats.Misses, 1)
 
-	newcn, err := p.newConn(ctx, true)
-	if err != nil {
+	return p.asyncNewConn(ctx)
+}
+
+func (p *ConnPool) asyncNewConn(ctx context.Context) (*Conn, error) {
+	// First try to acquire permission to create a connection
+	select {
+	case p.dialsInProgress <- struct{}{}:
+		// Got permission, proceed to create connection
+	case <-ctx.Done():
 		p.freeTurn()
-		return nil, err
+		return nil, ctx.Err()
+	}
+
+	dialCtx, cancel := context.WithTimeout(ctx, p.cfg.DialTimeout)
+
+	w := &wantConn{
+		ctx:       dialCtx,
+		cancelCtx: cancel,
+		result:    make(chan wantConnResult, 1),
 	}
 
-	return newcn, nil
+	go func(w *wantConn) {
+		defer w.cancelCtx()
+		defer func() { <-p.dialsInProgress }() // Release connection creation permission
+
+		cn, err := p.newConn(w.ctx, true)
+		delivered := w.tryDeliver(cn, err)
+		if err == nil && !delivered {
+			p.Put(w.ctx, cn)
+		} else {
+			p.freeTurn()
+		}
+	}(w)
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case result := <-w.result:
+		return result.cn, result.err
+	}
 }
 
 func (p *ConnPool) waitTurn(ctx context.Context) error {

internal/pool/pool_test.go

@@ -2,6 +2,7 @@ package pool_test
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"sync"
 	"testing"
@@ -377,8 +378,8 @@ var _ = Describe("race", func() {
 			state := p.Stats()
 			return state.TotalConns == 0 && state.IdleConns == 0 && p.QueueLen() == 0
 		}, "3s", "50ms").Should(BeTrue())
-  })
-  
+	})
+
 	It("wait", func() {
 		opt := &pool.Options{
 			Dialer: func(ctx context.Context) (net.Conn, error) {
@@ -435,3 +436,80 @@ var _ = Describe("race", func() {
 		Expect(stats.Timeouts).To(Equal(uint32(1)))
 	})
 })
+
+var _ = Describe("asyncNewConn", func() {
+	ctx := context.Background()
+
+	It("should successfully create connection when pool is exhausted", func() {
+		testPool := pool.NewConnPool(&pool.Options{
+			Dialer:             dummyDialer,
+			PoolSize:           1,
+			MaxConcurrentDials: 2,
+			DialTimeout:        1 * time.Second,
+			PoolTimeout:        1 * time.Second,
+		})
+		defer testPool.Close()
+
+		// Fill the pool
+		conn1, err := testPool.Get(ctx)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(conn1).NotTo(BeNil())
+
+		// This should trigger asyncNewConn since pool is exhausted
+		conn2, err := testPool.Get(ctx)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(conn2).NotTo(BeNil())
+
+		// Basic validation - we got two different connections
+		Expect(conn1).NotTo(Equal(conn2))
+
+		// Close connections without Put to avoid queue issues
+		_ = conn1.Close()
+		_ = conn2.Close()
+	})
+
+	It("should handle context cancellation", func() {
+		testPool := pool.NewConnPool(&pool.Options{
+			Dialer:             dummyDialer,
+			PoolSize:           1,
+			MaxConcurrentDials: 1,
+			DialTimeout:        1 * time.Second,
+			PoolTimeout:        1 * time.Second,
+		})
+		defer testPool.Close()
+
+		// Get the only connection
+		conn1, err := testPool.Get(ctx)
+		Expect(err).NotTo(HaveOccurred())
+
+		// Create cancelled context
+		cancelledCtx, cancel := context.WithCancel(ctx)
+		cancel()
+
+		// This should fail immediately with context cancelled
+		_, err = testPool.Get(cancelledCtx)
+		Expect(err).To(Equal(context.Canceled))
+
+		_ = conn1.Close()
+	})
+
+	It("should handle dial failures gracefully", func() {
+		alwaysFailDialer := func(ctx context.Context) (net.Conn, error) {
+			return nil, fmt.Errorf("dial failed")
+		}
+
+		testPool := pool.NewConnPool(&pool.Options{
+			Dialer:             alwaysFailDialer,
+			PoolSize:           1,
+			MaxConcurrentDials: 1,
+			DialTimeout:        1 * time.Second,
+			PoolTimeout:        1 * time.Second,
+		})
+		defer testPool.Close()
+
+		// This call should fail
+		_, err := testPool.Get(ctx)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("dial failed"))
+	})
+})

options.go

@@ -127,6 +127,10 @@ type Options struct {
 	// default: 3 seconds
 	WriteTimeout time.Duration
 
+	// MaxConcurrentDials is the maximum number of concurrent connection creation goroutines.
+	// If 0, defaults to PoolSize/4+1. If negative, unlimited goroutines (not recommended).
+	MaxConcurrentDials int
+
 	// ContextTimeoutEnabled controls whether the client respects context timeouts and deadlines.
 	// See https://redis.uptrace.dev/guide/go-redis-debugging.html#timeouts
 	ContextTimeoutEnabled bool
@@ -261,6 +265,16 @@ func (opt *Options) init() {
 	if opt.PoolSize == 0 {
 		opt.PoolSize = 10 * runtime.GOMAXPROCS(0)
 	}
+	if opt.MaxConcurrentDials == 0 {
+		// Default to PoolSize/4+1, at least 1, at most PoolSize
+		opt.MaxConcurrentDials = opt.PoolSize/4 + 1
+		if opt.MaxConcurrentDials > opt.PoolSize {
+			opt.MaxConcurrentDials = opt.PoolSize
+		}
+	} else if opt.MaxConcurrentDials < 0 {
+		// Negative value means unlimited, set to PoolSize
+		opt.MaxConcurrentDials = opt.PoolSize
+	}
 	if opt.ReadBufferSize == 0 {
 		opt.ReadBufferSize = proto.DefaultBufferSize
 	}
@@ -565,6 +579,7 @@ func setupConnParams(u *url.URL, o *Options) (*Options, error) {
 	o.MinIdleConns = q.int("min_idle_conns")
 	o.MaxIdleConns = q.int("max_idle_conns")
 	o.MaxActiveConns = q.int("max_active_conns")
+	o.MaxConcurrentDials = q.int("max_concurrent_dials")
 	if q.has("conn_max_idle_time") {
 		o.ConnMaxIdleTime = q.duration("conn_max_idle_time")
 	} else {
@@ -609,16 +624,17 @@ func newConnPool(
 		Dialer: func(ctx context.Context) (net.Conn, error) {
 			return dialer(ctx, opt.Network, opt.Addr)
 		},
-		PoolFIFO:        opt.PoolFIFO,
-		PoolSize:        opt.PoolSize,
-		PoolTimeout:     opt.PoolTimeout,
-		DialTimeout:     opt.DialTimeout,
-		MinIdleConns:    opt.MinIdleConns,
-		MaxIdleConns:    opt.MaxIdleConns,
-		MaxActiveConns:  opt.MaxActiveConns,
-		ConnMaxIdleTime: opt.ConnMaxIdleTime,
-		ConnMaxLifetime: opt.ConnMaxLifetime,
-		ReadBufferSize:  opt.ReadBufferSize,
-		WriteBufferSize: opt.WriteBufferSize,
+		PoolFIFO:           opt.PoolFIFO,
+		PoolSize:           opt.PoolSize,
+		MaxConcurrentDials: opt.MaxConcurrentDials,
+		PoolTimeout:        opt.PoolTimeout,
+		DialTimeout:        opt.DialTimeout,
+		MinIdleConns:       opt.MinIdleConns,
+		MaxIdleConns:       opt.MaxIdleConns,
+		MaxActiveConns:     opt.MaxActiveConns,
+		ConnMaxIdleTime:    opt.ConnMaxIdleTime,
+		ConnMaxLifetime:    opt.ConnMaxLifetime,
+		ReadBufferSize:     opt.ReadBufferSize,
+		WriteBufferSize:    opt.WriteBufferSize,
 	})
 }