refactor: idle pool and fix non-deterministic teardown of idle cleanup goroutine

clickhouse.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/rand"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -34,6 +35,7 @@ var (
 	ErrBindMixedParamsFormats    = errors.New("clickhouse [bind]: mixed named, numeric or positional parameters")
 	ErrAcquireConnNoAddress      = errors.New("clickhouse: no valid address supplied")
 	ErrServerUnexpectedData      = errors.New("code: 101, message: Unexpected packet Data received from client")
+	ErrConnectionClosed          = errors.New("clickhouse: connection is closed")
 )
 
 type OpError struct {
@@ -67,12 +69,13 @@ func Open(opt *Options) (driver.Conn, error) {
 	o := opt.setDefaults()
 
 	conn := &clickhouse{
-		opt:  o,
-		idle: make(chan nativeTransport, o.MaxIdleConns),
-		open: make(chan struct{}, o.MaxOpenConns),
-		exit: make(chan struct{}),
+		opt:       o,
+		idle:      newConnPool(o.ConnMaxLifetime, o.MaxIdleConns),
+		open:      make(chan struct{}, o.MaxOpenConns),
+		closeOnce: &sync.Once{},
+		closed:    &atomic.Bool{},
 	}
-	go conn.startAutoCloseIdleConnections()
+
 	return conn, nil
 }
 
@@ -101,10 +104,13 @@ type nativeTransportRelease func(nativeTransport, error)
 
 type clickhouse struct {
 	opt    *Options
-	idle   chan nativeTransport
-	open   chan struct{}
-	exit   chan struct{}
 	connID int64
+
+	idle *connPool
+	open chan struct{}
+
+	closeOnce *sync.Once
+	closed    *atomic.Bool
 }
 
 func (clickhouse) Contributors() []string {
@@ -226,13 +232,18 @@ func (ch *clickhouse) Ping(ctx context.Context) (err error) {
 func (ch *clickhouse) Stats() driver.Stats {
 	return driver.Stats{
 		Open:         len(ch.open),
-		Idle:         len(ch.idle),
 		MaxOpenConns: cap(ch.open),
-		MaxIdleConns: cap(ch.idle),
+
+		Idle:         ch.idle.Len(),
+		MaxIdleConns: ch.idle.Cap(),
 	}
 }
 
 func (ch *clickhouse) dial(ctx context.Context) (conn nativeTransport, err error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
 	connID := int(atomic.AddInt64(&ch.connID, 1))
 
 	dialFunc := func(ctx context.Context, addr string, opt *Options) (DialResult, error) {
@@ -286,87 +297,46 @@ func DefaultDialStrategy(ctx context.Context, connID int, opt *Options, dial Dia
 }
 
 func (ch *clickhouse) acquire(ctx context.Context) (conn nativeTransport, err error) {
-	timer := time.NewTimer(ch.opt.DialTimeout)
-	defer timer.Stop()
-	select {
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	default:
+	if ch.closed.Load() {
+		return nil, ErrConnectionClosed
 	}
+
+	ctx, cancel := context.WithTimeoutCause(ctx, ch.opt.DialTimeout, ErrAcquireConnTimeout)
+	defer cancel()
+
 	select {
-	case <-timer.C:
-		return nil, ErrAcquireConnTimeout
-	case <-ctx.Done():
-		return nil, ctx.Err()
 	case ch.open <- struct{}{}:
+	case <-ctx.Done():
+		return nil, context.Cause(ctx)
 	}
-	select {
-	case <-timer.C:
-		select {
-		case <-ch.open:
-		default:
-		}
-		return nil, ErrAcquireConnTimeout
-	case conn := <-ch.idle:
-		if conn.isBad() {
-			conn.close()
-			if conn, err = ch.dial(ctx); err != nil {
-				select {
-				case <-ch.open:
-				default:
-				}
-				return nil, err
-			}
+
+	conn, err = ch.idle.Get(ctx)
+	if err != nil && !errors.Is(err, errQueueEmpty) {
+		return nil, err
+	}
+
+	if err == nil && conn != nil {
+		if !conn.isBad() {
+			conn.setReleased(false)
+			conn.debugf("[acquired from pool]")
+			return conn, nil
 		}
-		conn.setReleased(false)
-		conn.debugf("[acquired from pool]")
-		return conn, nil
-	default:
+
+		conn.close()
 	}
+
 	if conn, err = ch.dial(ctx); err != nil {
 		select {
 		case <-ch.open:
 		default:
 		}
+
 		return nil, err
 	}
+
 	conn.debugf("[acquired new]")
 	return conn, nil
-}
-
-func (ch *clickhouse) startAutoCloseIdleConnections() {
-	ticker := time.NewTicker(ch.opt.ConnMaxLifetime)
-	defer ticker.Stop()
 
-	for {
-		select {
-		case <-ticker.C:
-			ch.closeIdleExpired()
-		case <-ch.exit:
-			return
-		}
-	}
-}
-
-func (ch *clickhouse) closeIdleExpired() {
-	cutoff := time.Now().Add(-ch.opt.ConnMaxLifetime)
-	for {
-		select {
-		case conn := <-ch.idle:
-			if conn.connectedAtTime().Before(cutoff) {
-				conn.close()
-			} else {
-				select {
-				case ch.idle <- conn:
-				default:
-					conn.close()
-				}
-				return
-			}
-		default:
-			return
-		}
-	}
 }
 
 func (ch *clickhouse) release(conn nativeTransport, err error) {
@@ -401,28 +371,19 @@ func (ch *clickhouse) release(conn nativeTransport, err error) {
 		conn.freeBuffer()
 	}
 
-	select {
-	case ch.idle <- conn:
-	default:
-		conn.debugf("[close: idle pool full %d/%d]", len(ch.idle), cap(ch.idle))
+	if ch.closed.Load() {
 		conn.close()
+		return
 	}
+
+	ch.idle.Put(conn)
 }
 
-func (ch *clickhouse) Close() error {
-	for {
-		select {
-		case conn := <-ch.idle:
-			conn.debugf("[close: closing pool]")
-			conn.close()
-		default:
-			// In rare cases, close may be called multiple times, don't block
-			//TODO: add proper close flag to indicate this pool is unusable after Close
-			select {
-			case ch.exit <- struct{}{}:
-			default:
-			}
-			return nil
-		}
-	}
+func (ch *clickhouse) Close() (err error) {
+	ch.closeOnce.Do(func() {
+		err = ch.idle.Close()
+		ch.closed.Store(true)
+	})
+
+	return
 }