Add pool generation numbers to allow a pool to be closed instantly.

Every connection that is opened by the pool is tracked with the pool's generation number. This allows us to close the pool immediately when `.Close()` is called.

* For connections that are idle, we close them directly when lowering the capacity.
* For connections that are returned to the pool, we check whether the generation numbers match and if they don't, we close the returned connection directly.

Signed-off-by: Arthur Schreiber <[email protected]>

Author: arthurschreiber

go/pools/smartconnpool/connection.go

@@ -36,6 +36,8 @@ type Pooled[C Connection] struct {
 	timeUsed    timestamp
 	pool        *ConnPool[C]
 
+	generation int64 // generation of the pool when this connection was created
+
 	Conn C
 }
 
@@ -48,16 +50,16 @@ func (dbc *Pooled[C]) Recycle() {
 	case dbc.pool == nil:
 		dbc.Conn.Close()
 	case dbc.Conn.IsClosed():
-		dbc.pool.put(nil)
+		dbc.pool.put(nil, dbc.generation)
 	default:
-		dbc.pool.put(dbc)
+		dbc.pool.put(dbc, dbc.generation)
 	}
 }
 
 func (dbc *Pooled[C]) Taint() {
 	if dbc.pool == nil {
 		return
 	}
-	dbc.pool.put(nil)
+	dbc.pool.put(nil, dbc.generation)
 	dbc.pool = nil
 }

go/pools/smartconnpool/pool.go

@@ -131,6 +131,9 @@ type ConnPool[C Connection] struct {
 	close      chan struct{}
 	capacityMu sync.Mutex
 
+	// generation of the pool, this is incremented every time the pool is closed
+	generation atomic.Int64
+
 	config struct {
 		// connect is the callback to create a new connection for the pool
 		connect Connector[C]
@@ -395,7 +398,17 @@ func (pool *ConnPool[C]) Get(ctx context.Context, setting *Setting) (*Pooled[C],
 
 // put returns a connection to the pool. This is a private API.
 // Return connections to the pool by calling Pooled.Recycle
-func (pool *ConnPool[C]) put(conn *Pooled[C]) {
+func (pool *ConnPool[C]) put(conn *Pooled[C], generation int64) {
+	// If the generations don't match, just close the connection and return.
+	// Don't update any metrics here.
+	if generation != pool.generation.Load() {
+		if conn != nil {
+			conn.Close()
+		}
+
+		return
+	}
+
 	pool.borrowed.Add(-1)
 
 	if conn == nil {
@@ -524,8 +537,9 @@ func (pool *ConnPool[C]) connNew(ctx context.Context) (*Pooled[C], error) {
 		return nil, err
 	}
 	pooled := &Pooled[C]{
-		pool: pool,
-		Conn: conn,
+		pool:       pool,
+		generation: pool.generation.Load(),
+		Conn:       conn,
 	}
 	now := monotonicNow()
 	pooled.timeUsed.set(now)
@@ -720,32 +734,46 @@ func (pool *ConnPool[C]) setCapacity(ctx context.Context, newcap int64) error {
 	// wait for connections to be returned to the pool if we're reducing the capacity.
 	defer pool.setIdleCount()
 
-	const delay = 10 * time.Millisecond
+	if newcap == 0 {
+		pool.generation.Add(1)
+		pool.active.Store(0)
+		pool.borrowed.Store(0)
 
-	// close connections until we're under capacity
-	for pool.active.Load() > newcap {
-		if err := ctx.Err(); err != nil {
-			return vterrors.Errorf(vtrpcpb.Code_ABORTED,
-				"timed out while waiting for connections to be returned to the pool (capacity=%d, active=%d, borrowed=%d)",
-				pool.capacity.Load(), pool.active.Load(), pool.borrowed.Load())
-		}
-		// if we're closing down the pool, make sure there's no clients waiting
-		// for connections because they won't be returned in the future
-		if newcap == 0 {
-			pool.wait.expire(true)
-		}
+		for {
+			conn := pool.getFromSettingsStack(nil)
+			if conn == nil {
+				conn = pool.pop(&pool.clean)
+			}
+			if conn == nil {
+				break
+			}
 
-		// try closing from connections which are currently idle in the stacks
-		conn := pool.getFromSettingsStack(nil)
-		if conn == nil {
-			conn = pool.pop(&pool.clean)
+			// Close the connection but don't update the metrics as we've already updated them
+			conn.Close()
 		}
-		if conn == nil {
-			time.Sleep(delay)
-			continue
+
+		pool.wait.expire(true)
+	} else {
+		const delay = 10 * time.Millisecond
+
+		// close connections until we're under capacity
+		for pool.active.Load() > newcap {
+			// try closing from connections which are currently idle in the stacks
+			conn := pool.getFromSettingsStack(nil)
+			if conn == nil {
+				conn = pool.pop(&pool.clean)
+			}
+
+			// Stop the loop if we don't have any more connections to close.
+			// Connections that were borrowed by clients and are over capacity
+			// will be closed when they are put back into the pool.
+			if conn == nil {
+				break
+			}
+
+			conn.Close()
+			pool.closedConn()
 		}
-		conn.Close()
-		pool.closedConn()
 	}
 
 	return nil