smartconnpool: export waiter and expired-handoff counters

go/pools/smartconnpool/bench_return_test.go

@@ -0,0 +1,43 @@
+package smartconnpool
+
+// Measures the cost of the tryReturnConnSlow critical section, so the exact
+// ExpiredHandoffs reconstruction can be shown not to widen it materially.
+// The scan is forced to traverse every waiter (no Setting match, no ageing
+// match), which is the worst case for the added accounting.
+
+import (
+	"context"
+	"fmt"
+	"testing"
+)
+
+func benchWaitlist(n int) *waitlist[*TestConn] {
+	wl := &waitlist[*TestConn]{}
+	wl.init()
+	for i := 0; i < n; i++ {
+		wl.list.PushBack(waiter[*TestConn]{
+			setting: &Setting{}, // never equal to the returned conn's nil Setting
+			ctx:     context.Background(),
+		})
+	}
+	return wl
+}
+
+func BenchmarkTryReturnConnSlow(b *testing.B) {
+	for _, n := range []int{1, 8, 64, 512} {
+		b.Run(fmt.Sprintf("waiters=%d", n), func(b *testing.B) {
+			conn := &Pooled[*TestConn]{Conn: &TestConn{}}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				// Rebuild outside the timer: tryReturnConnSlow ages the waiters
+				// it skips, so a reused list would trip age > maxAge after a few
+				// iterations and break out of the scan immediately, silently
+				// measuring a 1-element scan instead of an n-element one.
+				b.StopTimer()
+				wl := benchWaitlist(n)
+				b.StartTimer()
+				wl.tryReturnConnSlow(conn)
+			}
+		})
+	}
+}

go/pools/smartconnpool/expired_handoff_exact_test.go

@@ -0,0 +1,188 @@
+package smartconnpool
+
+// Exactness tests for ExpiredHandoffs.
+//
+// The counter must satisfy exactly one property:
+//
+//	ExpiredHandoffs increments on a return iff the waiter that the pristine v21
+//	selection rule would have chosen was expired.
+//
+// The pristine rule (865f789e15, waitlist.go tryReturnConnSlow) is:
+//
+//	target = list.Front()
+//	for e := target; e != nil; e = e.Next() {
+//	    if e.age > maxAge || e.setting == connSetting { target = e; break }
+//	    e.age++
+//	}
+//
+// i.e. the first waiter matching (age > maxAge || setting == connSetting),
+// otherwise the front element. It had no notion of contexts, so an expired
+// waiter could become the target through any of those three routes.
+//
+// The two named tests below are the counterexamples that falsified the previous
+// "idx == 0 || setting == connSetting" heuristic.
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// legacyMaxAge mirrors the const in tryReturnConnSlow.
+const legacyMaxAge = 8
+
+type shape struct {
+	expired bool
+	age     uint32
+	setting *Setting
+}
+
+// referencePristineTargetExpired is an INDEPENDENT implementation of the
+// pristine rule, written directly from the 865f789e15 source. It deliberately
+// shares no code with the production reconstruction.
+func referencePristineTargetExpired(shapes []shape, connSetting *Setting) bool {
+	if len(shapes) == 0 {
+		return false
+	}
+	target := 0 // list.Front()
+	for i := range shapes {
+		if shapes[i].age > legacyMaxAge || shapes[i].setting == connSetting {
+			target = i
+			break
+		}
+	}
+	return shapes[target].expired
+}
+
+// runShape drives one return against a synthetic waitlist and reports whether
+// ExpiredHandoffs incremented.
+func runShape(t *testing.T, shapes []shape, connSetting *Setting) bool {
+	t.Helper()
+	wl := &waitlist[*TestConn]{}
+	wl.init()
+	for _, s := range shapes {
+		var ctx context.Context = context.Background()
+		if s.expired {
+			c, cancel := context.WithCancel(context.Background())
+			cancel()
+			ctx = c
+		}
+		wl.list.PushBack(waiter[*TestConn]{setting: s.setting, ctx: ctx, age: s.age})
+	}
+
+	conn := &Pooled[*TestConn]{Conn: &TestConn{setting: connSetting}}
+
+	before := wl.preventedHandoffs.Load()
+	wl.tryReturnConnSlow(conn)
+	return wl.preventedHandoffs.Load()-before == 1
+}
+
+// TestExpiredHandoffUndercountAgeBranch is Agent 3 counterexample 1.
+//
+// A live waiter sits at the front, and an expired waiter behind it has
+// age > maxAge. Under pristine the age branch makes the EXPIRED waiter the
+// target, so this return would have handed a connection to an expired waiter.
+// The previous heuristic missed it because the expired waiter was neither at
+// idx 0 nor a Setting match.
+func TestExpiredHandoffUndercountAgeBranch(t *testing.T) {
+	s1, s2 := &Setting{}, &Setting{}
+	shapes := []shape{
+		{expired: false, age: 0, setting: s2},               // front, live, no match
+		{expired: true, age: legacyMaxAge + 1, setting: s2}, // expired, wins on age
+	}
+	require.True(t, referencePristineTargetExpired(shapes, s1),
+		"reference: pristine target must be the expired waiter")
+	require.True(t, runShape(t, shapes, s1),
+		"ExpiredHandoffs must count the age>maxAge route (undercount counterexample)")
+}
+
+// TestExpiredHandoffOvercountFrontNotTarget is Agent 3 counterexample 2.
+//
+// An expired waiter sits at the front but matches nothing, and a LIVE waiter
+// behind it matches the returned connection's Setting. Under pristine the live
+// waiter wins, so no bad handoff would have occurred. The previous heuristic
+// incremented purely because the expired waiter was at idx 0.
+func TestExpiredHandoffOvercountFrontNotTarget(t *testing.T) {
+	s1, s2 := &Setting{}, &Setting{}
+	shapes := []shape{
+		{expired: true, age: 0, setting: s2},  // front, expired, but NOT the target
+		{expired: false, age: 0, setting: s1}, // live, wins on Setting match
+	}
+	require.False(t, referencePristineTargetExpired(shapes, s1),
+		"reference: pristine target must be the live Setting match")
+	require.False(t, runShape(t, shapes, s1),
+		"ExpiredHandoffs must not count when a live waiter would have won (overcount counterexample)")
+}
+
+// TestExpiredHandoffExhaustive proves exactness over every waitlist shape up to
+// length 4, against the independent reference implementation.
+func TestExpiredHandoffExhaustive(t *testing.T) {
+	s1, s2 := &Setting{}, &Setting{}
+	settings := []*Setting{nil, s1, s2}
+	ages := []uint32{0, legacyMaxAge, legacyMaxAge + 1}
+	connSettings := []*Setting{nil, s1, s2}
+
+	var atoms []shape
+	for _, exp := range []bool{false, true} {
+		for _, a := range ages {
+			for _, st := range settings {
+				atoms = append(atoms, shape{expired: exp, age: a, setting: st})
+			}
+		}
+	}
+
+	var checked, mismatches int
+	var build func(cur []shape)
+	build = func(cur []shape) {
+		if len(cur) > 0 {
+			for _, cs := range connSettings {
+				want := referencePristineTargetExpired(cur, cs)
+				got := runShape(t, cur, cs)
+				checked++
+				if want != got {
+					mismatches++
+					if mismatches <= 5 {
+						t.Errorf("shape %v connSetting=%v: want ExpiredHandoffs increment=%v, got %v",
+							describe(cur), name(cs, s1, s2), want, got)
+					}
+				}
+			}
+		}
+		if len(cur) == 4 {
+			return
+		}
+		for _, a := range atoms {
+			build(append(cur, a))
+		}
+	}
+	build(nil)
+
+	require.Zero(t, mismatches, "ExpiredHandoffs must be exact for every shape")
+	t.Logf("exhaustive exactness: %d (shape, connSetting) cases, 0 mismatches", checked)
+}
+
+func name(s, s1, s2 *Setting) string {
+	switch s {
+	case nil:
+		return "nil"
+	case s1:
+		return "S1"
+	case s2:
+		return "S2"
+	}
+	return "?"
+}
+
+func describe(shapes []shape) string {
+	out := ""
+	for _, s := range shapes {
+		st := "nil"
+		if s.setting != nil {
+			st = "S"
+		}
+		out += fmt.Sprintf("[exp=%v age=%d set=%s]", s.expired, s.age, st)
+	}
+	return out
+}

go/pools/smartconnpool/pool.go

@@ -341,6 +341,27 @@ func (pool *ConnPool[C]) Active() int64 {
 	return pool.active.Load()
 }
 
+// Waiters returns the number of clients currently blocked waiting to acquire a
+// connection. Note that Available() is capacity minus borrowed and therefore
+// says nothing about contention; Waiters is the metric that does.
+func (pool *ConnPool[C]) Waiters() int64 {
+	return int64(pool.wait.waiting())
+}
+
+// ExpiredHandoffs returns the number of connection returns in which a waiter
+// whose acquisition context had already expired would have received the
+// connection under the pre-fix selection rules. A non-zero value is direct
+// evidence that the pool was attempting post-deadline handoffs.
+func (pool *ConnPool[C]) ExpiredHandoffs() int64 {
+	return pool.wait.preventedHandoffs.Load()
+}
+
+// WaitersEvicted returns the total number of waiters removed from the waitlist
+// because their acquisition context had already expired.
+func (pool *ConnPool[C]) WaitersEvicted() int64 {
+	return pool.wait.evicted.Load()
+}
+
 func (pool *ConnPool[D]) IdleTimeout() time.Duration {
 	return time.Duration(pool.config.idleTimeout.Load())
 }
@@ -854,6 +875,15 @@ func (pool *ConnPool[C]) RegisterStats(stats *servenv.Exporter, name string) {
 	stats.NewCounterFunc(name+"WaitCount", "Tablet server conn pool wait count", func() int64 {
 		return pool.Metrics.WaitCount()
 	})
+	stats.NewGaugeFunc(name+"Waiters", "Tablet server conn pool waiters currently blocked on acquisition", func() int64 {
+		return pool.Waiters()
+	})
+	stats.NewCounterFunc(name+"ExpiredHandoffs", "Number of connection returns in which an expired waiter would have received the connection", func() int64 {
+		return pool.ExpiredHandoffs()
+	})
+	stats.NewCounterFunc(name+"WaitersEvicted", "Number of waiters removed from the waitlist because their acquisition context had expired", func() int64 {
+		return pool.WaitersEvicted()
+	})
 	stats.NewCounterDurationFunc(name+"WaitTime", "Tablet server wait time", func() time.Duration {
 		return pool.Metrics.WaitTime()
 	})

go/pools/smartconnpool/waiter_observability_test.go

@@ -0,0 +1,60 @@
+package smartconnpool
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestWaiterObservability asserts that the new counters answer the question we
+// could not answer during the incident: "did a returner attempt to hand a
+// connection to an expired waiter?" -- without inferring it from aggregate
+// acquire latency.
+func TestWaiterObservability(t *testing.T) {
+	const acquireTimeout = 100 * time.Millisecond
+
+	state := &TestState{}
+	p := NewPool(&Config[*TestConn]{Capacity: 1, IdleTimeout: time.Hour}).
+		Open(newConnector(state), nil)
+	defer p.Close()
+
+	held, err := p.Get(context.Background(), nil)
+	require.NoError(t, err)
+
+	require.Zero(t, p.Waiters(), "no waiters before the test starts")
+	require.Zero(t, p.ExpiredHandoffs())
+	require.Zero(t, p.WaitersEvicted())
+
+	type result struct{ err error }
+	resc := make(chan result, 1)
+	go func() {
+		ctx := &expiredCtx{deadline: time.Now().Add(acquireTimeout)}
+		_, err := p.Get(ctx, nil)
+		resc <- result{err: err}
+	}()
+
+	for p.wait.waiting() < 1 {
+		time.Sleep(200 * time.Microsecond)
+	}
+	require.EqualValues(t, 1, p.Waiters(), "Waiters gauge must observe the blocked client")
+
+	time.Sleep(2 * acquireTimeout) // waiter is now expired but still linked
+
+	waitsBefore := p.Metrics.WaitCount()
+	held.Recycle()
+	r := <-resc
+
+	require.Error(t, r.err, "expired waiter must not receive a connection")
+	require.EqualValues(t, 1, p.ExpiredHandoffs(),
+		"a returner attempted to hand a connection to an expired waiter")
+	require.EqualValues(t, 1, p.WaitersEvicted(),
+		"the expired waiter was evicted from the waitlist")
+	require.Zero(t, p.Waiters(), "gauge returns to zero after eviction")
+	require.Equal(t, waitsBefore, p.Metrics.WaitCount(),
+		"success-only wait metric semantics must be preserved")
+
+	t.Logf("Waiters=%d ExpiredHandoffs=%d WaitersEvicted=%d WaitCount=%d",
+		p.Waiters(), p.ExpiredHandoffs(), p.WaitersEvicted(), p.Metrics.WaitCount())
+}