drpcpool: close connections returned to a closed pool

Pool.Close only sees the connections currently idle in the cache.
Connections that are checked out serving an in-flight Invoke or
NewStream live outside the cache and are returned later via Put. Because
the pool had no closed state, such a late Put re-inserted the connection
and armed a fresh expiration timer, resurrecting the pool and leaking the
connection (and its manager goroutines) until the timer fired.

This is reachable on shutdown: long-lived streams (e.g. rangefeeds) hold
their connection checked out, so Close cannot reach them. When the stream
context is later canceled, the connection is handed back via Put and
lingers well past the point the pool was closed.

Mark the pool closed under the lock in Close, and have Put close the
connection immediately instead of caching it once the pool is closed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: shubhamdhama

drpcpool/pool.go

@@ -52,6 +52,7 @@ type Options struct {
 type Pool[K comparable, V Conn] struct {
 	opts    Options
 	mu      sync.Mutex
+	closed  bool
 	entries map[K]*list[K, V]
 	order   list[K, V]
 }
@@ -119,10 +120,20 @@ func (p *Pool[K, V]) log(what string, cb func() string) {
 
 // Close evicts all entries from the Pool's cache, closing them and returning all
 // of the combined errors from closing.
+//
+// Close also marks the pool as closed so that any connection subsequently
+// returned via Put is closed immediately rather than cached. This matters
+// because Close can only see the connections currently idle in the cache:
+// connections that are checked out (serving an in-flight Invoke or NewStream)
+// live outside the cache and are returned later via Put. Without the closed
+// flag those late returns would resurrect the pool, re-arming an expiration
+// timer and leaking the connection (and its manager goroutines) until it fires.
 func (p *Pool[K, V]) Close() (err error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	p.closed = true
+
 	var eg errs.Group
 	for ent := p.order.head; ent != nil; ent = ent.global.next {
 		eg.Add(p.closeEntry(ent))
@@ -233,6 +244,13 @@ func (p *Pool[K, V]) Put(key K, val V) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	// If the pool has been closed, don't cache the connection (which would
+	// resurrect the pool and leak the connection); close it instead.
+	if p.closed {
+		_ = val.Close()
+		return
+	}
+
 	local := p.entries[key]
 	if local == nil {
 		local = new(list[K, V])

drpcpool/pool_test.go

@@ -672,3 +672,41 @@ func BenchmarkPool(b *testing.B) {
 		invoke(ctx, conn)
 	}
 }
+
+// TestPoolPutAfterClose verifies that a connection returned to the pool after
+// the pool is closed is closed immediately rather than cached. This models a
+// long-lived stream whose connection is checked out (and therefore invisible to
+// Close) when the pool is closed, and only handed back via Put once the stream
+// finally ends. Caching it would re-arm the expiration timer and leak the
+// connection until it fires.
+func TestPoolPutAfterClose(t *testing.T) {
+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	// Use a non-zero expiration so a cached connection would linger (and arm a
+	// timer) rather than be evicted by capacity limits.
+	pool := New[string, Conn](Options{Expiration: time.Hour})
+
+	closed := make(chan string, 1)
+	conn := &callbackConn{CloseFn: func() error { closed <- "key"; return nil }}
+
+	// Stage the connection in the pool, then check it out, modeling an in-flight
+	// stream. While checked out the connection lives outside the cache.
+	pool.Put("key", conn)
+	got, ok := pool.Take("key")
+	assert.True(t, ok)
+	assert.Equal(t, got, conn)
+	assert.Equal(t, len(closed), 0)
+
+	// Closing the pool cannot see the checked-out connection.
+	assert.NoError(t, pool.Close())
+	assert.Equal(t, len(closed), 0)
+
+	// Returning it now must close it instead of resurrecting the pool.
+	pool.Put("key", conn)
+	assert.Equal(t, <-closed, "key")
+
+	// And it must not have been retained.
+	_, ok = pool.Take("key")
+	assert.That(t, !ok)
+}