From 93161e9ca87601e05614f217d7cbd17376c192e7 Mon Sep 17 00:00:00 2001 From: fasterrt Date: Thu, 27 Aug 2026 18:08:00 -0700 Subject: [PATCH 1/2] smartconnpool: do not hand connections to expired waiters A returned connection could be handed to a waiter whose acquisition context had already expired. The waiter had no way to refuse it, so the connection was consumed by a request that could no longer use it. Semantic backport of vitessio/vitess#20308, adapted to the v21 semaphore-based waitlist: - tryReturnConnSlow evicts expired waiters during the return scan and continues past interior expired waiters to find a live one - maybeStarvingCount no longer counts expired waiters as starving - go/list gains O(1) RemoveIfPresent so an evicted waiter that is concurrently timing out cannot double-remove - WaitCount/WaitTime remain successful-acquisition-only: waitForConn can now return (nil, nil), so recordWait is guarded on conn == nil No waiter cap, no timeout changes, no unrelated fixes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1f56db40-28ec-4335-97a1-b60deee78b07 --- go/list/list.go | 12 + go/pools/smartconnpool/expired_waiter_test.go | 181 +++++++++++++++ go/pools/smartconnpool/pool.go | 4 +- .../post_deadline_handoff_test.go | 214 ++++++++++++++++++ go/pools/smartconnpool/waitlist.go | 71 +++--- 5 files changed, 453 insertions(+), 29 deletions(-) create mode 100644 go/pools/smartconnpool/expired_waiter_test.go create mode 100644 go/pools/smartconnpool/post_deadline_handoff_test.go diff --git a/go/list/list.go b/go/list/list.go index 2ad837b7c64..a417408d29c 100644 --- a/go/list/list.go +++ b/go/list/list.go @@ -142,6 +142,18 @@ func (l *List[T]) Remove(e *Element[T]) { l.remove(e) } +// RemoveIfPresent removes e from l if e is currently an element of l, and +// reports whether it was removed. Unlike Remove, it does not panic when e does +// not belong to l, so callers that may race with another goroutine removing e +// can use the return value to decide who owns the element. +func (l *List[T]) RemoveIfPresent(e *Element[T]) bool { + if e.list != l { + return false + } + l.remove(e) + return true +} + // PushFront inserts a new element e with value v at the front of list l and returns e. func (l *List[T]) PushFront(v T) *Element[T] { return l.insertValue(v, &l.root) diff --git a/go/pools/smartconnpool/expired_waiter_test.go b/go/pools/smartconnpool/expired_waiter_test.go new file mode 100644 index 00000000000..c84917354f4 --- /dev/null +++ b/go/pools/smartconnpool/expired_waiter_test.go @@ -0,0 +1,181 @@ +/* +Copyright 2025 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package smartconnpool + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// lateExpiryCtx is valid at the moment its waiter enqueues and expires while +// that waiter is parked on the waitlist. Its Done channel is never closed, so +// the waiter does not remove itself on expiry. This deterministically produces +// a waitlist containing already-expired waiters, which is the shape that arises +// in production when clients time out faster than the pool hands out +// connections. +type lateExpiryCtx struct { + context.Context + deadline time.Time + done chan struct{} +} + +func (c *lateExpiryCtx) Deadline() (time.Time, bool) { return c.deadline, true } +func (c *lateExpiryCtx) Done() <-chan struct{} { return c.done } +func (c *lateExpiryCtx) Err() error { + if time.Now().After(c.deadline) { + return context.DeadlineExceeded + } + return nil +} + +func newLateExpiryCtx(deadline time.Time) *lateExpiryCtx { + return &lateExpiryCtx{ + Context: context.Background(), + deadline: deadline, + done: make(chan struct{}), + } +} + +// TestExpiredWaitersDoNotReceiveConnections asserts that a connection returned +// to the pool is never handed to a waiter whose context has already expired, +// including when expired waiters sit *behind* a live waiter in the list. +// +// Handing a connection to an expired waiter wastes it: the caller can no longer +// use it and, in the transaction pool, the failed Begin closes the connection +// and forces a fresh dial. Under load this starves the live waiters. +func TestExpiredWaitersDoNotReceiveConnections(t *testing.T) { + var state TestState + + p := NewPool(&Config[*TestConn]{ + Capacity: 1, + LogWait: state.LogWait, + }).Open(newConnector(&state), nil) + defer p.Close() + + // Occupy the only slot so that every subsequent Get parks on the waitlist. + held, err := p.Get(context.Background(), nil) + require.NoError(t, err) + + // A single shared deadline keeps expiry deterministic across all waiters. + deadline := time.Now().Add(2 * time.Second) + + // Expired (true) and live (false) waiters, interleaved so that expired + // waiters sit both ahead of and behind live waiters. + kinds := []bool{true, true, true, false, true, true, true, false} + + type outcome struct { + expired bool + gotConn bool + } + results := make([]outcome, len(kinds)) + + var wg sync.WaitGroup + for i, isExpired := range kinds { + results[i].expired = isExpired + + ctx := context.Context(context.Background()) + if isExpired { + ctx = newLateExpiryCtx(deadline) + } + + before := p.wait.waiting() + wg.Add(1) + go func() { + defer wg.Done() + conn, err := p.Get(ctx, nil) + if err == nil && conn != nil { + results[i].gotConn = true + p.put(conn) + } + }() + + // Wait until this waiter is on the list before starting the next one, + // so that the resulting list order is exactly `kinds`. + require.Eventuallyf(t, func() bool { + return p.wait.waiting() == before+1 + }, 5*time.Second, time.Millisecond, "waiter %d never enqueued", i) + } + + // Let the lateExpiryCtx waiters expire while they are parked. + time.Sleep(time.Until(deadline) + 100*time.Millisecond) + + // Releasing the held connection starts the handoff chain. + p.put(held) + + waited := make(chan struct{}) + go func() { wg.Wait(); close(waited) }() + select { + case <-waited: + case <-time.After(30 * time.Second): + t.Fatal("waiters did not finish; connection handoff stalled") + } + + for i, r := range results { + if r.expired { + require.Falsef(t, r.gotConn, + "expired waiter %d received a connection; it can never use it", i) + } else { + require.Truef(t, r.gotConn, + "live waiter %d was starved of a connection", i) + } + } +} + +// TestExpiredWaiterEvictionIsRaceSafe exercises the eviction path concurrently +// with waiters removing themselves on genuine context cancellation, to confirm +// that exactly one party owns each waitlist element. +func TestExpiredWaiterEvictionIsRaceSafe(t *testing.T) { + var state TestState + + p := NewPool(&Config[*TestConn]{ + Capacity: 2, + LogWait: state.LogWait, + }).Open(newConnector(&state), nil) + defer p.Close() + + var wg sync.WaitGroup + for i := 0; i < 200; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + var ctx context.Context + var cancel context.CancelFunc + switch i % 3 { + case 0: + ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond) + case 1: + ctx, cancel = context.WithCancel(context.Background()) + go func() { time.Sleep(time.Millisecond); cancel() }() + default: + ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) + } + defer cancel() + + conn, err := p.Get(ctx, nil) + if err == nil && conn != nil { + p.put(conn) + } + }(i) + } + wg.Wait() + + require.Zero(t, p.wait.waiting(), "waitlist should be drained") +} diff --git a/go/pools/smartconnpool/pool.go b/go/pools/smartconnpool/pool.go index 4eb6c4fa57f..76219c3d4d7 100644 --- a/go/pools/smartconnpool/pool.go +++ b/go/pools/smartconnpool/pool.go @@ -571,7 +571,7 @@ func (pool *ConnPool[C]) get(ctx context.Context) (*Pooled[C], error) { } conn, err = pool.wait.waitForConn(ctx, nil, *closeChan) - if err != nil { + if err != nil || conn == nil { return nil, ErrTimeout } pool.recordWait(start) @@ -634,7 +634,7 @@ func (pool *ConnPool[C]) getWithSetting(ctx context.Context, setting *Setting) ( } conn, err = pool.wait.waitForConn(ctx, setting, *closeChan) - if err != nil { + if err != nil || conn == nil { return nil, ErrTimeout } pool.recordWait(start) diff --git a/go/pools/smartconnpool/post_deadline_handoff_test.go b/go/pools/smartconnpool/post_deadline_handoff_test.go new file mode 100644 index 00000000000..e448093caa4 --- /dev/null +++ b/go/pools/smartconnpool/post_deadline_handoff_test.go @@ -0,0 +1,214 @@ +/* +Copyright 2025 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package smartconnpool + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// expiredCtx becomes expired (Err() != nil) once its deadline passes, but its +// Done() channel never fires. +// +// This deterministically pins the exact state that tryReturnConnSlow observes in +// production: a waiter whose acquisition deadline has already passed but which is +// still linked in the waitlist because its own goroutine has not yet executed, or +// has not yet won, its self-removal under wl.mu. +// +// Using a real context.WithTimeout would reach the same state only by winning a +// race against the returner. Suppressing Done() removes the race without changing +// anything the returner can observe: tryReturnConnSlow only ever reads +// waiter.ctx, never waiter's Done channel. +type expiredCtx struct{ deadline time.Time } + +func (c *expiredCtx) Deadline() (time.Time, bool) { return c.deadline, true } +func (c *expiredCtx) Done() <-chan struct{} { return nil } +func (c *expiredCtx) Value(any) any { return nil } +func (c *expiredCtx) Err() error { + if time.Now().After(c.deadline) { + return context.DeadlineExceeded + } + return nil +} + +// TestPostDeadlineHandoff asserts the invariant that upstream vitessio/vitess +// #20308 restores: the pool must never hand a connection to a waiter whose +// acquisition context has already expired. +// +// On pristine v21 this test FAILS, reproducing the production signature: +// - Get() returns success after its acquisition deadline +// - the success-only WaitTime/WaitCount metric therefore exceeds the deadline +// - and no connection is closed or re-dialled while it happens +func TestPostDeadlineHandoff(t *testing.T) { + // scaled equivalent of the production --queryserver-config-txpool-timeout=1s + const acquireTimeout = 100 * time.Millisecond + + var state TestState + p := NewPool(&Config[*TestConn]{ + Capacity: 1, + IdleTimeout: time.Minute, + LogWait: state.LogWait, + }).Open(newConnector(&state), nil) + defer p.Close() + + // Occupy the only connection so the next caller must enqueue as a waiter. + held, err := p.Get(context.Background(), nil) + require.NoError(t, err) + dialsAfterWarmup := state.open.Load() + + type result struct { + conn *Pooled[*TestConn] + err error + elapsed time.Duration + ctxErr error + } + resc := make(chan result, 1) + + ctx := &expiredCtx{deadline: time.Now().Add(acquireTimeout)} + go func() { + start := time.Now() + c, err := p.Get(ctx, nil) + resc <- result{conn: c, err: err, elapsed: time.Since(start), ctxErr: ctx.Err()} + }() + + // Wait until the caller is actually parked on the waitlist. + deadline := time.Now().Add(5 * time.Second) + for p.wait.waiting() == 0 { + require.True(t, time.Now().Before(deadline), "waiter never enqueued") + time.Sleep(time.Millisecond) + } + + // Let the acquisition deadline pass while the waiter is parked. + time.Sleep(2 * acquireTimeout) + require.Error(t, ctx.Err(), "precondition: waiter ctx must be expired before the return") + require.Equal(t, 1, p.wait.waiting(), "precondition: expired waiter is still on the waitlist") + + // The returner hands the connection back. Front of the list is the expired waiter. + waitsBefore := p.Metrics.WaitCount() + held.Recycle() + + r := <-resc + + t.Logf("Get -> err=%v elapsed=%v ctxErr=%v | WaitCount=%d WaitTime=%v | dials=%d closes=%d", + r.err, r.elapsed, r.ctxErr, + p.Metrics.WaitCount(), p.Metrics.WaitTime(), + state.open.Load(), state.close.Load()) + + // (1) The pool must not report success to a waiter that had already expired. + if r.err == nil && r.conn != nil { + mean := time.Duration(0) + if n := p.Metrics.WaitCount(); n > 0 { + mean = time.Duration(int64(p.Metrics.WaitTime()) / n) + } + r.conn.Recycle() + t.Fatalf("POST-DEADLINE HANDOFF: pool delivered a connection to a waiter whose "+ + "acquisition context expired %v earlier.\n"+ + " Get() returned SUCCESS after %v (acquisition deadline %v)\n"+ + " success-only metric mean wait = %v > deadline %v <-- production signature\n"+ + " dials during handoff = %d (no redial) closes = %d", + r.elapsed-acquireTimeout, r.elapsed, acquireTimeout, + mean, acquireTimeout, + state.open.Load()-dialsAfterWarmup, state.close.Load()) + } + + // (2) With the fix, the expired waiter is evicted and the acquisition fails. + require.Error(t, r.err, "expired waiter must not receive a connection") + require.Nil(t, r.conn) + + // (3) Evicting the expired waiter must NOT pollute the success-only wait + // metric. WaitCount/WaitTime have always counted successful acquisitions + // only; a failed eviction must leave them untouched. + require.Equal(t, waitsBefore, p.Metrics.WaitCount(), + "evicted expired waiter must not increment WaitCount (success-only metric)") + + // (3) Either way, no connection is closed or re-dialled by this path. This is + // why the mechanism is invisible in MySQL's Connections counter. + require.Equal(t, dialsAfterWarmup, state.open.Load(), "no redial must occur") + require.Zero(t, state.close.Load(), "no connection close must occur") +} + +// TestInteriorExpiredWaiterNotServed covers the case that a prefix-only +// eviction would miss. Production deadlines are heterogeneous, so an expired +// waiter can sit BEHIND a live one. tryReturnConnSlow picks its target by +// `age > maxAge || setting == connSetting`, so an interior expired waiter whose +// Setting matches the returned connection can be selected ahead of the live +// waiter at the front of the list. +// +// Upstream #20308 scans the whole list rather than stopping at the first live +// waiter. This test asserts that behaviour: the interior expired waiter must be +// evicted, and the live waiter must receive the connection. +func TestInteriorExpiredWaiterNotServed(t *testing.T) { + settingA := NewSetting("set workload='olap'", "set workload='oltp'") + + state := &TestState{} + p := NewPool(&Config[*TestConn]{Capacity: 1, IdleTimeout: time.Hour}). + Open(newConnector(state), nil) + defer p.Close() + + // Hold the only connection and stamp settingA on it, so the connection that + // is later returned carries settingA. + held, err := p.Get(context.Background(), settingA) + require.NoError(t, err) + + var liveServed, expiredServed atomic.Int64 + var wg sync.WaitGroup + + // (1) LIVE waiter, no Setting -> occupies the FRONT of the waitlist. + wg.Add(1) + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if c, err := p.Get(ctx, nil); err == nil && c != nil { + liveServed.Add(1) + c.Recycle() + } + }() + for p.wait.waiting() < 1 { + time.Sleep(200 * time.Microsecond) + } + + // (2) EXPIRED waiter, Setting == settingA -> sits BEHIND the live waiter, + // and is the better `setting == connSetting` match for the returned conn. + wg.Add(1) + go func() { + defer wg.Done() + ctx := &expiredCtx{deadline: time.Now().Add(50 * time.Millisecond)} + if c, err := p.Get(ctx, settingA); err == nil && c != nil { + expiredServed.Add(1) + c.Recycle() + } + }() + for p.wait.waiting() < 2 { + time.Sleep(200 * time.Microsecond) + } + time.Sleep(100 * time.Millisecond) // waiter 2 is now expired, still linked + + held.Recycle() + wg.Wait() + + t.Logf("live served=%d expired served=%d", liveServed.Load(), expiredServed.Load()) + require.Zero(t, expiredServed.Load(), + "interior expired waiter must not receive a connection") + require.Equal(t, int64(1), liveServed.Load(), + "live waiter must receive the connection instead") +} diff --git a/go/pools/smartconnpool/waitlist.go b/go/pools/smartconnpool/waitlist.go index 40c924da327..6850a546329 100644 --- a/go/pools/smartconnpool/waitlist.go +++ b/go/pools/smartconnpool/waitlist.go @@ -70,17 +70,10 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC select { case <-closeChan: // Pool was closed while we were waiting. - removed := false - + // Try to remove ourselves from the list. If we lose the race against + // tryReturnConnSlow, it owns the element and will notify our semaphore. wl.mu.Lock() - // Try to find and remove ourselves from the list. - for e := wl.list.Front(); e != nil; e = e.Next() { - if e == elem { - wl.list.Remove(elem) - removed = true - break - } - } + removed := wl.list.RemoveIfPresent(elem) wl.mu.Unlock() // If we removed ourselves from the waitlist, we need to notify our semaphore @@ -100,17 +93,10 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC case <-ctx.Done(): // Context expired. We need to try to remove ourselves from the waitlist to // prevent another goroutine from trying to hand us a connection later on. - removed := false - + // If we lose the race against tryReturnConnSlow, it owns the element and + // will notify our semaphore. wl.mu.Lock() - // Try to find and remove ourselves from the list. - for e := wl.list.Front(); e != nil; e = e.Next() { - if e == elem { - wl.list.Remove(elem) - removed = true - break - } - } + removed := wl.list.RemoveIfPresent(elem) wl.mu.Unlock() // If we removed ourselves from the waitlist, we need to notify our semaphore @@ -140,9 +126,13 @@ func (wl *waitlist[C]) maybeStarvingCount() (maybeStarving int) { wl.mu.Lock() defer wl.mu.Unlock() - // iterate the waitlist looking for waiters with an expired Context, - // or remove everything if force is true + // iterate the waitlist looking for waiters that are still live and have not + // been skipped over yet. Waiters whose context has already expired are not + // counted: opening a new connection on their behalf would be wasted work. for e := wl.list.Front(); e != nil; e = e.Next() { + if e.Value.ctx != nil && e.Value.ctx.Err() != nil { + continue + } if e.Value.age == 0 { maybeStarving++ } @@ -165,15 +155,33 @@ func (wl *waitlist[D]) tryReturnConnSlow(conn *Pooled[D]) bool { const maxAge = 8 var ( target *list.Element[waiter[D]] + expired []*list.Element[waiter[D]] connSetting = conn.Conn.Setting() ) wl.mu.Lock() - target = wl.list.Front() - // iterate through the waitlist looking for either waiters that have been - // here too long, or a waiter that is looking exactly for the same Setting - // as the one we have in our connection. - for e := target; e != nil; e = e.Next() { + var next *list.Element[waiter[D]] + for e := wl.list.Front(); e != nil; e = next { + // capture the successor before a possible Remove unlinks e + next = e.Next() + + // Never hand a connection over to a waiter whose context has already + // expired: that waiter can no longer use it, and for the transaction + // pool the subsequent failed Begin closes the connection and forces a + // fresh dial. Evict the expired waiter instead so the connection stays + // available for a waiter that can still use it. Expired waiters can sit + // anywhere in the list, not just at the front, because deadlines are + // heterogeneous, so the whole scan has to check. + if e.Value.ctx != nil && e.Value.ctx.Err() != nil { + wl.list.Remove(e) + expired = append(expired, e) + continue + } + + if target == nil { + // front-most live waiter, used unless a better match is found below + target = e + } if e.Value.age > maxAge || e.Value.setting == connSetting { target = e break @@ -190,6 +198,15 @@ func (wl *waitlist[D]) tryReturnConnSlow(conn *Pooled[D]) bool { } wl.mu.Unlock() + // Wake the evicted waiters. Their conn stays nil, which waitForConn hands + // back to the caller and which Get maps to a timeout. This is safe to do + // after releasing the lock: an evicted waiter is blocked on its semaphore + // and cannot return (nor recycle its list element) until we notify it, and + // because we removed it from the list under the lock nobody else can. + for _, e := range expired { + e.Value.sema.notify(false) + } + // maybe there isn't anybody to hand over the connection to, because we've // raced with another client returning another connection if target == nil { From be4736c7fbcf12e841b11cef81976033fcadd00d Mon Sep 17 00:00:00 2001 From: fasterrt Date: Fri, 28 Aug 2026 10:03:31 -0700 Subject: [PATCH 2/2] pool_test: restore Go 1.23 compatibility release-21.0-github builds with go 1.23.10 (go.mod), but the backport of #19004 brought in test code written against the release-22.0 toolchain: - t.Context() (testing.T.Context, Go 1.24) - wg.Go(...) (sync.WaitGroup.Go, Go 1.25) - b.Loop() (testing.B.Loop, Go 1.24) This makes go/pools/smartconnpool fail to typecheck, which fails 'Static Code Checks Etc', every 'Unit Test' matrix job and 'Code Coverage' on this branch, and hides the results of the new expired-waiter tests added by this PR. Rewrite the four call sites with their Go 1.23 equivalents. No test behaviour changes: the contexts already carry their own timeout and cancel, and b.ResetTimer() was already being called before the loop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- go/pools/smartconnpool/pool_test.go | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/go/pools/smartconnpool/pool_test.go b/go/pools/smartconnpool/pool_test.go index 29785d771ff..d18aef9bce1 100644 --- a/go/pools/smartconnpool/pool_test.go +++ b/go/pools/smartconnpool/pool_test.go @@ -1294,7 +1294,7 @@ func TestIdleTimeoutConnectionLeak(t *testing.T) { LogWait: state.LogWait, }).Open(newConnector(&state), nil) - getCtx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + getCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) defer cancel() // Get and return two connections @@ -1328,15 +1328,18 @@ func TestIdleTimeoutConnectionLeak(t *testing.T) { wg := sync.WaitGroup{} for i := 0; i < 2; i++ { - wg.Go(func() { - getCtx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond) + wg.Add(1) + go func() { + defer wg.Done() + + getCtx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) defer cancel() conn, err := p.Get(getCtx, nil) require.NoError(t, err) p.put(conn) - }) + }() } wg.Wait() @@ -1356,7 +1359,7 @@ func TestIdleTimeoutConnectionLeak(t *testing.T) { assert.Equal(t, int64(2), p.Metrics.IdleClosed()) // Try to close the pool - if there are leaked connections, this will timeout - closeCtx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + closeCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) defer cancel() err = p.CloseWithContext(closeCtx) @@ -1446,7 +1449,7 @@ func BenchmarkPoolCleanupIdleConnectionsPerformanceNoIdleConnections(b *testing. b.ResetTimer() - for b.Loop() { + for i := 0; i < b.N; i++ { p.closeIdleResources(time.Now()) } }