From 986f565ebc19acbdb0f9b39d4309b0eaae63ee11 Mon Sep 17 00:00:00 2001 From: Ian Oberst Date: Thu, 22 Oct 2026 13:58:41 -0700 Subject: [PATCH 1/4] Run heartbeat leak tests without MySQL or test order The engine cleanup tests skipped in CI because they needed the MySQL fixture, but Engine.Prepare only pings the database. Add test.PingDB, a database/sql driver whose pings succeed and queries fail, and run both engine tests with it. The reader leak tests required no heartbeat readers in the whole process before they started, so they depended on earlier tests and on how fast stopped readers exit. Count only readers started by the test, using a pprof label that goroutines inherit (test.Goroutines). Error path tests also check that Prepare called Stop on the reader, because a leaked reader might not have started when the count is checked. Export the heartbeat reader function and timeout from the test package, instead of defining them in both test packages. Co-Authored-By: Claude Opus 5.5 --- metrics/repl.lag/lag_test.go | 73 +++++++++++++++++++++++----------- monitor/engine_cleanup_test.go | 39 +++++++----------- test/db.go | 38 ++++++++++++++++++ test/goroutine.go | 65 +++++++++++++++++++++++++----- 4 files changed, 157 insertions(+), 58 deletions(-) create mode 100644 test/db.go diff --git a/metrics/repl.lag/lag_test.go b/metrics/repl.lag/lag_test.go index 5baddb9..c23562b 100644 --- a/metrics/repl.lag/lag_test.go +++ b/metrics/repl.lag/lag_test.go @@ -7,8 +7,8 @@ import ( "database/sql" "errors" "strings" + "sync/atomic" "testing" - "time" "github.com/stretchr/testify/assert" @@ -121,25 +121,42 @@ func TestCollectWithAllSources(t *testing.T) { } } -const ( - readerFunc = "heartbeat.(*BlipReader).run" - readerTimeout = 5 * time.Second // > heartbeat.ReadTimeout: reader exits after an in-flight query -) - -// stopReader stops the heartbeat reader, if any, when the test ends. Tests -// that count reader goroutines require that earlier tests don't leak them. +// stopReader stops the heartbeat reader, if any, when the test ends. func stopReader(t *testing.T, cleanup func()) { if cleanup != nil { t.Cleanup(cleanup) } } -// requireNoReaders waits for readers stopped by earlier tests to exit, so the -// test can count only the readers that it starts. -func requireNoReaders(t *testing.T) { +// spyReader is a heartbeat reader that records whether it was stopped. +type spyReader struct { + heartbeat.Reader + stopped atomic.Bool +} + +func (r *spyReader) Stop() { + r.stopped.Store(true) + r.Reader.Stop() +} + +// spyOnReader makes c wrap its heartbeat reader in a spyReader. +func spyOnReader(c *Lag) { + c.newReader = func(args heartbeat.BlipReaderArgs) heartbeat.Reader { + return &spyReader{Reader: newBlipReader(args)} + } +} + +// requireReaderStopped requires that Prepare made a reader and stopped it. +// Counting goroutines isn't enough: the reader goroutine might not have +// started yet. +func requireReaderStopped(t *testing.T, c *Lag) { t.Helper() - if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 { - t.Fatalf("%d heartbeat readers running before test, expected 0 (leaked by an earlier test?)", got) + spy, ok := c.lagReader.(*spyReader) + if !ok { + t.Fatalf("Prepare didn't make a heartbeat reader: %T", c.lagReader) + } + if !spy.stopped.Load() { + t.Error("Prepare didn't stop the heartbeat reader") } } @@ -172,21 +189,22 @@ func TestPrepareBlipAtMultipleLevelsReturnsReaderCleanup(t *testing.T) { // Plans usually collect repl.lag at more than one level because levels // inherit domains from lower levels. Prepare starts one reader for all // levels, so it must return the cleanup func that stops that reader. - requireNoReaders(t) - c := NewLag(unreachableDB(t)) - cleanup, err := c.Prepare(context.Background(), blipPlan("fast", "medium", "slow")) + var cleanup func() + var err error + readers := test.NewGoroutines() + readers.Do(func() { cleanup, err = c.Prepare(context.Background(), blipPlan("fast", "medium", "slow")) }) if err != nil { t.Fatal(err) } - if got := test.WaitGoroutines(readerFunc, 1, readerTimeout); got != 1 { + if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 { t.Fatalf("%d heartbeat readers running after Prepare, expected 1", got) } if cleanup == nil { t.Fatal("Prepare returned nil cleanup func, so the heartbeat reader can't be stopped") } cleanup() - if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 { + if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 { t.Errorf("%d heartbeat readers running after cleanup, expected 0", got) } } @@ -194,15 +212,18 @@ func TestPrepareBlipAtMultipleLevelsReturnsReaderCleanup(t *testing.T) { func TestPrepareBlipErrorStopsReader(t *testing.T) { // The reader starts before the repl check is verified, which fails because // the database is unreachable. Prepare must stop the reader it started. - requireNoReaders(t) - c := NewLag(unreachableDB(t)) + spyOnReader(c) plan := blipPlan("kpi") plan.Levels["kpi"].Collect[DOMAIN].Options[OPT_REPL_CHECK] = "read_only" - if _, err := c.Prepare(context.Background(), plan); err == nil { + var err error + readers := test.NewGoroutines() + readers.Do(func() { _, err = c.Prepare(context.Background(), plan) }) + if err == nil { t.Fatal("Prepare returned nil error, expected repl check error") } - if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 { + requireReaderStopped(t, c) + if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 { t.Errorf("%d heartbeat readers running after Prepare error, expected 0", got) } } @@ -248,11 +269,17 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) { level.Collect[DOMAIN].Options[OPT_REPORT_NO_HEARTBEAT] = report } } - cleanup, err := c.Prepare(context.Background(), plan) + var cleanup func() + var err error + readers := test.NewGoroutines() + readers.Do(func() { cleanup, err = c.Prepare(context.Background(), plan) }) stopReader(t, cleanup) if err != nil { t.Fatal(err) } + if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 { + t.Fatalf("%d heartbeat readers running after Prepare, expected 1 for all levels", got) + } for level := range plan.Levels { metrics, err := c.Collect(context.Background(), level) if err != nil { diff --git a/monitor/engine_cleanup_test.go b/monitor/engine_cleanup_test.go index cc4be4b..3a7a834 100644 --- a/monitor/engine_cleanup_test.go +++ b/monitor/engine_cleanup_test.go @@ -7,7 +7,6 @@ import ( "errors" "sync" "testing" - "time" "github.com/cashapp/blip/v2" "github.com/cashapp/blip/v2/metrics" @@ -16,17 +15,11 @@ import ( "github.com/cashapp/blip/v2/test/mock" ) -const ( - heartbeatReaderFunc = "heartbeat.(*BlipReader).run" - heartbeatReaderTimeout = 5 * time.Second // > heartbeat.ReadTimeout: reader exits after an in-flight query -) - func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { - _, db, err := test.Connection(test.DefaultMySQLVersion) + // Engine.Prepare only pings the database, and the repl.lag blip writer + // doesn't query it, so MySQL isn't needed + db, err := test.PingDB() if err != nil { - if test.Build { - t.Skip("MySQL test fixture not running") - } t.Fatal(err) } defer db.Close() @@ -46,39 +39,37 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { }, } - // Count only readers started by this test: wait for readers stopped by - // earlier tests to exit - if got := test.WaitGoroutines(heartbeatReaderFunc, 0, heartbeatReaderTimeout); got != 0 { - t.Fatalf("%d heartbeat readers running before test, expected 0 (leaked by an earlier test?)", got) - } + // Count only readers started by this test + readers := test.NewGoroutines() for i := 0; i < 3; i++ { engine := NewEngine(blip.ConfigMonitor{MonitorId: plan.MonitorId}, db) - if err := engine.Prepare(context.Background(), plan, func() {}, func() {}); err != nil { + readers.Do(func() { err = engine.Prepare(context.Background(), plan, func() {}, func() {}) }) + if err != nil { t.Fatal(err) } - if got := test.WaitGoroutines(heartbeatReaderFunc, 1, heartbeatReaderTimeout); got != 1 { + if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 { t.Fatalf("cycle %d: %d heartbeat readers running after Prepare, expected 1", i, got) } engine.Stop() - if got := test.WaitGoroutines(heartbeatReaderFunc, 0, heartbeatReaderTimeout); got != 0 { + if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 { t.Fatalf("cycle %d: %d heartbeat readers running after Stop, expected 0", i, got) } } } func TestEnginePrepareErrorCleansUpPreparedCollectors(t *testing.T) { - _, db, err := test.Connection(test.DefaultMySQLVersion) + // Engine.Prepare only pings the database, and the mock collectors don't + // query it, so MySQL isn't needed + db, err := test.PingDB() if err != nil { - if test.Build { - t.Skip("MySQL test fixture not running") - } t.Fatal(err) } defer db.Close() // Two collectors. Domains are prepared in random order, so whichever is - // prepared second fails. Both return a cleanup func, which Prepare must - // call because it discards both collectors. + // prepared second fails. Prepare discards both collectors, so it must call + // the cleanup func of the first. The second returns a cleanup func with its + // error, which Prepare calls as a backstop (see blip.Collector.Prepare). domains := []string{"test.cleanup-a", "test.cleanup-b"} mux := &sync.Mutex{} prepared := 0 diff --git a/test/db.go b/test/db.go new file mode 100644 index 0000000..0072778 --- /dev/null +++ b/test/db.go @@ -0,0 +1,38 @@ +// Copyright 2026 Block, Inc. + +package test + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "sync" +) + +const pingDriverName = "blip-test-ping" + +var registerPingDriver sync.Once + +// PingDB returns a database handle that doesn't need MySQL: pings succeed and +// queries fail. It's for tests that need a database only to connect, like +// monitor.Engine.Prepare. +func PingDB() (*sql.DB, error) { + registerPingDriver.Do(func() { sql.Register(pingDriverName, pingDriver{}) }) + return sql.Open(pingDriverName, "") +} + +var errPingDB = errors.New("test.PingDB: queries not supported") + +type pingDriver struct{} + +func (pingDriver) Open(string) (driver.Conn, error) { return pingConn{}, nil } + +type pingConn struct{} + +var _ driver.Pinger = pingConn{} + +func (pingConn) Ping(context.Context) error { return nil } +func (pingConn) Prepare(string) (driver.Stmt, error) { return nil, errPingDB } +func (pingConn) Close() error { return nil } +func (pingConn) Begin() (driver.Tx, error) { return nil, errPingDB } diff --git a/test/goroutine.go b/test/goroutine.go index a827d3a..662210d 100644 --- a/test/goroutine.go +++ b/test/goroutine.go @@ -4,31 +4,74 @@ package test import ( "bytes" + "context" + "fmt" "runtime/pprof" + "strconv" "strings" + "sync/atomic" "time" ) -// Goroutines returns the number of goroutines with fn in their stack, like -// "heartbeat.(*BlipReader).run". Tests use it to detect leaked goroutines. -func Goroutines(fn string) int { +const ( + // HeartbeatReaderFunc is in the stack of a running heartbeat.BlipReader. + HeartbeatReaderFunc = "heartbeat.(*BlipReader).run" + + // HeartbeatReaderTimeout is how long to wait for a stopped heartbeat reader + // to exit. It's > heartbeat.ReadTimeout: the reader exits after an in-flight + // query. + HeartbeatReaderTimeout = 5 * time.Second +) + +const goroutinesLabel = "blip-test-goroutines" + +var lastGoroutinesId atomic.Uint64 + +// Goroutines is a group of goroutines started by Do, including goroutines +// that those goroutines start. Tests use it to detect leaked goroutines. +// Counts include only goroutines in the group, so they don't depend on test +// order or on how fast goroutines stopped by other tests exit. +type Goroutines struct { + id string // pprof label value +} + +func NewGoroutines() *Goroutines { + return &Goroutines{id: strconv.FormatUint(lastGoroutinesId.Add(1), 10)} +} + +// Do calls f. Goroutines that f starts are in the group. +func (g *Goroutines) Do(f func()) { + pprof.Do(context.Background(), pprof.Labels(goroutinesLabel, g.id), func(context.Context) { f() }) +} + +// Count returns the number of goroutines in the group with fn in their stack, +// like HeartbeatReaderFunc. +func (g *Goroutines) Count(fn string) int { var buf bytes.Buffer - pprof.Lookup("goroutine").WriteTo(&buf, 2) // 2 = one stack per goroutine + pprof.Lookup("goroutine").WriteTo(&buf, 1) // 1 = count goroutines by stack and labels + _, profile, _ := strings.Cut(buf.String(), "\n") // skip "goroutine profile: total N" + label := fmt.Sprintf("%q:%q", goroutinesLabel, g.id) n := 0 - for _, stack := range strings.Split(buf.String(), "\n\n") { - if strings.Contains(stack, fn) { - n++ + for _, stack := range strings.Split(profile, "\n\n") { + if !strings.Contains(stack, label) || !strings.Contains(stack, fn) { + continue + } + count, _, _ := strings.Cut(stack, " ") // " @ ..." + c, err := strconv.Atoi(count) + if err != nil { + panic(fmt.Sprintf("invalid goroutine profile: %s: %s", err, stack)) } + n += c } return n } -// WaitGoroutines waits up to timeout for Goroutines(fn) to equal n. It returns -// the last count, so the caller should check that it equals n. -func WaitGoroutines(fn string, n int, timeout time.Duration) int { +// Wait waits up to timeout for Count(fn) to equal n. It returns the last +// count, so the caller should check that it equals n. +func (g *Goroutines) Wait(fn string, n int, timeout time.Duration) int { deadline := time.Now().Add(timeout) for { - got := Goroutines(fn) + got := g.Count(fn) if got == n || time.Now().After(deadline) { return got } From d079a847238bc1cb4b20ee46410e0c348631b239 Mon Sep 17 00:00:00 2001 From: Ian Oberst Date: Wed, 28 Oct 2026 14:24:18 -0700 Subject: [PATCH 2/4] Catch heartbeat reader leaks across tests and guard goroutine counts Group-scoped reader counts miss a reader leaked by a test that doesn't count readers. Add a TestMain to monitor and metrics/repl.lag that waits after all tests for no heartbeat readers in the process, using the new test.CountGoroutines and test.WaitGoroutines. It runs once, so it doesn't depend on test order. Replace test.Goroutines.Do with Label, which labels the test goroutine, so the goroutines it starts afterward are in the group without a closure at each call site. Make goroutine counts fail loudly instead of returning 0: the count must find its own goroutine in the profile, with the group label, and panics on a stack deep enough that the profile might have truncated it. Check that the reader stops in TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat. Move PingDB to test/mock, next to ContextDB, built on sql.OpenDB with a connector instead of a registered driver. ContextDB also connects without MySQL, but its queries block until the read timeout, so each Engine.Stop in the reader leak test would wait 2s. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ian Oberst --- metrics/repl.lag/lag_test.go | 20 +++--- metrics/repl.lag/main_test.go | 23 +++++++ monitor/engine_cleanup_test.go | 14 ++-- monitor/main_test.go | 23 +++++++ test/db.go | 38 ----------- test/goroutine.go | 115 ++++++++++++++++++++++++++------- test/mock/db.go | 40 ++++++++++++ 7 files changed, 193 insertions(+), 80 deletions(-) create mode 100644 metrics/repl.lag/main_test.go create mode 100644 monitor/main_test.go delete mode 100644 test/db.go diff --git a/metrics/repl.lag/lag_test.go b/metrics/repl.lag/lag_test.go index c23562b..adc3ac6 100644 --- a/metrics/repl.lag/lag_test.go +++ b/metrics/repl.lag/lag_test.go @@ -190,10 +190,9 @@ func TestPrepareBlipAtMultipleLevelsReturnsReaderCleanup(t *testing.T) { // inherit domains from lower levels. Prepare starts one reader for all // levels, so it must return the cleanup func that stops that reader. c := NewLag(unreachableDB(t)) - var cleanup func() - var err error readers := test.NewGoroutines() - readers.Do(func() { cleanup, err = c.Prepare(context.Background(), blipPlan("fast", "medium", "slow")) }) + readers.Label() + cleanup, err := c.Prepare(context.Background(), blipPlan("fast", "medium", "slow")) if err != nil { t.Fatal(err) } @@ -216,9 +215,9 @@ func TestPrepareBlipErrorStopsReader(t *testing.T) { spyOnReader(c) plan := blipPlan("kpi") plan.Levels["kpi"].Collect[DOMAIN].Options[OPT_REPL_CHECK] = "read_only" - var err error readers := test.NewGoroutines() - readers.Do(func() { _, err = c.Prepare(context.Background(), plan) }) + readers.Label() + _, err := c.Prepare(context.Background(), plan) if err == nil { t.Fatal("Prepare returned nil error, expected repl check error") } @@ -269,11 +268,10 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) { level.Collect[DOMAIN].Options[OPT_REPORT_NO_HEARTBEAT] = report } } - var cleanup func() - var err error readers := test.NewGoroutines() - readers.Do(func() { cleanup, err = c.Prepare(context.Background(), plan) }) - stopReader(t, cleanup) + readers.Label() + cleanup, err := c.Prepare(context.Background(), plan) + stopReader(t, cleanup) // backstop if the test fails before it stops the reader if err != nil { t.Fatal(err) } @@ -293,6 +291,10 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) { t.Errorf("%s: got %+v, expected no metrics (dropped)", level, metrics) } } + cleanup() // stopReader calls it again, which is a no-op + if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 { + t.Errorf("%d heartbeat readers running after cleanup, expected 0", got) + } }) } } diff --git a/metrics/repl.lag/main_test.go b/metrics/repl.lag/main_test.go new file mode 100644 index 0000000..c0619a7 --- /dev/null +++ b/metrics/repl.lag/main_test.go @@ -0,0 +1,23 @@ +// Copyright 2026 Block, Inc. + +package repllag + +import ( + "fmt" + "os" + "testing" + + "github.com/cashapp/blip/v2/test" +) + +// TestMain fails the run if a heartbeat reader is still running after all +// tests. Tests count only the readers that they start (test.Goroutines), so +// this catches a reader leaked by a test that doesn't count readers. +func TestMain(m *testing.M) { + code := m.Run() + if n := test.WaitGoroutines(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); n != 0 { + fmt.Fprintf(os.Stderr, "FAIL: %d heartbeat readers running after all tests, expected 0: a test didn't stop the reader that it started\n", n) + code = 1 + } + os.Exit(code) +} diff --git a/monitor/engine_cleanup_test.go b/monitor/engine_cleanup_test.go index 3a7a834..45f5365 100644 --- a/monitor/engine_cleanup_test.go +++ b/monitor/engine_cleanup_test.go @@ -18,10 +18,7 @@ import ( func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { // Engine.Prepare only pings the database, and the repl.lag blip writer // doesn't query it, so MySQL isn't needed - db, err := test.PingDB() - if err != nil { - t.Fatal(err) - } + db := mock.PingDB() defer db.Close() // repl.lag at two levels, like a plan where the slow level inherits @@ -41,10 +38,10 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { // Count only readers started by this test readers := test.NewGoroutines() + readers.Label() for i := 0; i < 3; i++ { engine := NewEngine(blip.ConfigMonitor{MonitorId: plan.MonitorId}, db) - readers.Do(func() { err = engine.Prepare(context.Background(), plan, func() {}, func() {}) }) - if err != nil { + if err := engine.Prepare(context.Background(), plan, func() {}, func() {}); err != nil { t.Fatal(err) } if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 { @@ -60,10 +57,7 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { func TestEnginePrepareErrorCleansUpPreparedCollectors(t *testing.T) { // Engine.Prepare only pings the database, and the mock collectors don't // query it, so MySQL isn't needed - db, err := test.PingDB() - if err != nil { - t.Fatal(err) - } + db := mock.PingDB() defer db.Close() // Two collectors. Domains are prepared in random order, so whichever is diff --git a/monitor/main_test.go b/monitor/main_test.go new file mode 100644 index 0000000..fa304a1 --- /dev/null +++ b/monitor/main_test.go @@ -0,0 +1,23 @@ +// Copyright 2026 Block, Inc. + +package monitor + +import ( + "fmt" + "os" + "testing" + + "github.com/cashapp/blip/v2/test" +) + +// TestMain fails the run if a heartbeat reader is still running after all +// tests. Tests count only the readers that they start (test.Goroutines), so +// this catches a reader leaked by a test that doesn't count readers. +func TestMain(m *testing.M) { + code := m.Run() + if n := test.WaitGoroutines(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); n != 0 { + fmt.Fprintf(os.Stderr, "FAIL: %d heartbeat readers running after all tests, expected 0: a test didn't stop the reader that it started\n", n) + code = 1 + } + os.Exit(code) +} diff --git a/test/db.go b/test/db.go deleted file mode 100644 index 0072778..0000000 --- a/test/db.go +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright 2026 Block, Inc. - -package test - -import ( - "context" - "database/sql" - "database/sql/driver" - "errors" - "sync" -) - -const pingDriverName = "blip-test-ping" - -var registerPingDriver sync.Once - -// PingDB returns a database handle that doesn't need MySQL: pings succeed and -// queries fail. It's for tests that need a database only to connect, like -// monitor.Engine.Prepare. -func PingDB() (*sql.DB, error) { - registerPingDriver.Do(func() { sql.Register(pingDriverName, pingDriver{}) }) - return sql.Open(pingDriverName, "") -} - -var errPingDB = errors.New("test.PingDB: queries not supported") - -type pingDriver struct{} - -func (pingDriver) Open(string) (driver.Conn, error) { return pingConn{}, nil } - -type pingConn struct{} - -var _ driver.Pinger = pingConn{} - -func (pingConn) Ping(context.Context) error { return nil } -func (pingConn) Prepare(string) (driver.Stmt, error) { return nil, errPingDB } -func (pingConn) Close() error { return nil } -func (pingConn) Begin() (driver.Tx, error) { return nil, errPingDB } diff --git a/test/goroutine.go b/test/goroutine.go index 662210d..9f64be9 100644 --- a/test/goroutine.go +++ b/test/goroutine.go @@ -25,35 +25,111 @@ const ( const goroutinesLabel = "blip-test-goroutines" +// countFunc is in the stack of the goroutine that calls countGoroutines. +// countGoroutines checks that it finds itself in the profile, so a change in +// the profile format fails loudly instead of counting 0. +const countFunc = "test.countGoroutines" + +// maxFrames is less than the max stack depth of the goroutine profile (128 +// frames by default since Go 1.23). A stack with this many frames might be +// truncated, which drops its outer frames, like heartbeat.(*BlipReader).run. +const maxFrames = 100 + var lastGoroutinesId atomic.Uint64 -// Goroutines is a group of goroutines started by Do, including goroutines -// that those goroutines start. Tests use it to detect leaked goroutines. -// Counts include only goroutines in the group, so they don't depend on test -// order or on how fast goroutines stopped by other tests exit. +// Goroutines is a group of goroutines: the goroutine that calls Label and the +// goroutines that it starts afterward, including goroutines that those +// goroutines start. Tests use it to detect leaked goroutines. Counts include +// only goroutines in the group, so they don't depend on test order or on how +// fast goroutines stopped by other tests exit. type Goroutines struct { id string // pprof label value } +// NewGoroutines returns a group that has no goroutines until Label is called. func NewGoroutines() *Goroutines { return &Goroutines{id: strconv.FormatUint(lastGoroutinesId.Add(1), 10)} } -// Do calls f. Goroutines that f starts are in the group. -func (g *Goroutines) Do(f func()) { - pprof.Do(context.Background(), pprof.Labels(goroutinesLabel, g.id), func(context.Context) { f() }) +// Label adds the calling goroutine to the group. Goroutines that it starts +// afterward are in the group too. A test should call Label from its own +// goroutine (the one that calls the Test or t.Run func), which exits when the +// test ends, so the label doesn't apply to other tests. Subtests started +// afterward are in the group unless they call Label on another group. +func (g *Goroutines) Label() { + pprof.SetGoroutineLabels(pprof.WithLabels(context.Background(), pprof.Labels(goroutinesLabel, g.id))) } // Count returns the number of goroutines in the group with fn in their stack, -// like HeartbeatReaderFunc. +// like HeartbeatReaderFunc. It must be called from the goroutine that called +// Label. func (g *Goroutines) Count(fn string) int { + return countGoroutines(fn, g.id) +} + +// Wait waits up to timeout for Count(fn) to equal n. It returns the last +// count, so the caller should check that it equals n. +func (g *Goroutines) Wait(fn string, n int, timeout time.Duration) int { + return wait(func() int { return g.Count(fn) }, n, timeout) +} + +// CountGoroutines returns the number of goroutines in the process with fn in +// their stack, like HeartbeatReaderFunc. The count depends on other tests, so +// use it only after all tests have run, like in TestMain. +func CountGoroutines(fn string) int { + return countGoroutines(fn, "") +} + +// WaitGoroutines waits up to timeout for CountGoroutines(fn) to equal n. It +// returns the last count, so the caller should check that it equals n. +func WaitGoroutines(fn string, n int, timeout time.Duration) int { + return wait(func() int { return CountGoroutines(fn) }, n, timeout) +} + +func wait(count func() int, n int, timeout time.Duration) int { + deadline := time.Now().Add(timeout) + for { + got := count() + if got == n || time.Now().After(deadline) { + return got + } + time.Sleep(10 * time.Millisecond) + } +} + +// countGoroutines counts goroutines with fn in their stack and, if id isn't +// empty, the group label id. It parses the goroutine profile at debug=1, which +// counts goroutines by stack and labels: +// +// @ ... +// # labels: {"":""} +// # + : +// ... +// +// with a blank line between stacks. It panics if it doesn't find its own +// stack (with the label, if id is set) or if a counted stack might be +// truncated, because then it might count 0 for running goroutines. +func countGoroutines(fn, id string) int { var buf bytes.Buffer - pprof.Lookup("goroutine").WriteTo(&buf, 1) // 1 = count goroutines by stack and labels + pprof.Lookup("goroutine").WriteTo(&buf, 1) _, profile, _ := strings.Cut(buf.String(), "\n") // skip "goroutine profile: total N" - label := fmt.Sprintf("%q:%q", goroutinesLabel, g.id) + label := "" + if id != "" { + label = fmt.Sprintf("%q:%q", goroutinesLabel, id) // in the "# labels:" line + } n := 0 + foundSelf := false for _, stack := range strings.Split(profile, "\n\n") { - if !strings.Contains(stack, label) || !strings.Contains(stack, fn) { + if label != "" && !strings.Contains(stack, label) { + continue + } + if strings.Contains(stack, countFunc) { + foundSelf = true + } + if !strings.Contains(stack, fn) { + if frames := strings.Count(stack, "\n#\t"); frames >= maxFrames { + panic(fmt.Sprintf("goroutine stack might be truncated: %d frames:\n%s", frames, stack)) + } continue } count, _, _ := strings.Cut(stack, " ") // " @ ..." @@ -63,18 +139,11 @@ func (g *Goroutines) Count(fn string) int { } n += c } - return n -} - -// Wait waits up to timeout for Count(fn) to equal n. It returns the last -// count, so the caller should check that it equals n. -func (g *Goroutines) Wait(fn string, n int, timeout time.Duration) int { - deadline := time.Now().Add(timeout) - for { - got := g.Count(fn) - if got == n || time.Now().After(deadline) { - return got + if !foundSelf { + if id != "" { + panic(fmt.Sprintf("goroutine profile has no stack with %s and label %s: Count wasn't called from the goroutine that called Label, or the profile format changed:\n%s", countFunc, label, buf.String())) } - time.Sleep(10 * time.Millisecond) + panic(fmt.Sprintf("goroutine profile has no stack with %s: the profile format changed:\n%s", countFunc, buf.String())) } + return n } diff --git a/test/mock/db.go b/test/mock/db.go index d6982e7..4102149 100644 --- a/test/mock/db.go +++ b/test/mock/db.go @@ -78,3 +78,43 @@ func (r contextRows) Next([]driver.Value) error { <-r.ctx.Done() return r.ctx.Err() } + +// PingDB returns a *sql.DB for tests that need a database only to connect, +// like monitor.Engine.Prepare, without a running MySQL instance. Pings +// succeed and queries fail right away, so a heartbeat reader that uses it +// retries without blocking. +func PingDB() *sql.DB { + return sql.OpenDB(pingConnector{}) +} + +var errPingDB = errors.New("mock.PingDB: queries not supported") + +type pingConnector struct{} + +func (pingConnector) Connect(context.Context) (driver.Conn, error) { + return pingConn{}, nil +} + +func (pingConnector) Driver() driver.Driver { + return contextDriver{} +} + +type pingConn struct{} + +var _ driver.Pinger = pingConn{} + +func (pingConn) Ping(context.Context) error { + return nil +} + +func (pingConn) Prepare(string) (driver.Stmt, error) { + return nil, errPingDB +} + +func (pingConn) Close() error { + return nil +} + +func (pingConn) Begin() (driver.Tx, error) { + return nil, errPingDB +} From a8369b0b92425b234542a72f46e5c3762a706a78 Mon Sep 17 00:00:00 2001 From: Ian Oberst Date: Wed, 28 Oct 2026 15:29:07 -0700 Subject: [PATCH 3/4] Stop the engine if the heartbeat reader leak test fails When a check in the loop failed, the engine and its heartbeat reader kept running, so TestMain waited its full timeout and reported a second, redundant failure. Stop each engine in t.Cleanup; Stop is a no-op after the loop's own Stop. Close the database in t.Cleanup too, so it closes after the engines stop. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ian Oberst --- monitor/engine_cleanup_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monitor/engine_cleanup_test.go b/monitor/engine_cleanup_test.go index 45f5365..4103e66 100644 --- a/monitor/engine_cleanup_test.go +++ b/monitor/engine_cleanup_test.go @@ -19,7 +19,7 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { // Engine.Prepare only pings the database, and the repl.lag blip writer // doesn't query it, so MySQL isn't needed db := mock.PingDB() - defer db.Close() + t.Cleanup(func() { db.Close() }) // after the engines stop: cleanups run last-in, first-out // repl.lag at two levels, like a plan where the slow level inherits // domains from the fast level @@ -41,6 +41,7 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) { readers.Label() for i := 0; i < 3; i++ { engine := NewEngine(blip.ConfigMonitor{MonitorId: plan.MonitorId}, db) + t.Cleanup(engine.Stop) // if the test fails before Stop; a second Stop is a no-op if err := engine.Prepare(context.Background(), plan, func() {}, func() {}); err != nil { t.Fatal(err) } From 40cc872115e74bb0f5b59fdf1d31a2eda1d42cd8 Mon Sep 17 00:00:00 2001 From: Ian Oberst Date: Thu, 29 Oct 2026 06:56:47 -0700 Subject: [PATCH 4/4] Stop the reader if the reader cleanup test fails When a check failed before cleanup, the heartbeat reader kept running, so TestMain waited its full timeout and reported a second, redundant failure. Stop the reader in t.Cleanup. It stops c.lagReader directly, so it also covers the case the test checks for, where Prepare returns a nil cleanup func. Fix the countGoroutines doc comment: the truncation check applies to stacks that don't contain fn, because truncation drops the outer frames. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ian Oberst --- metrics/repl.lag/lag_test.go | 7 +++++++ test/goroutine.go | 4 +++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/metrics/repl.lag/lag_test.go b/metrics/repl.lag/lag_test.go index adc3ac6..f2656d8 100644 --- a/metrics/repl.lag/lag_test.go +++ b/metrics/repl.lag/lag_test.go @@ -190,6 +190,13 @@ func TestPrepareBlipAtMultipleLevelsReturnsReaderCleanup(t *testing.T) { // inherit domains from lower levels. Prepare starts one reader for all // levels, so it must return the cleanup func that stops that reader. c := NewLag(unreachableDB(t)) + // Stop the reader if the test fails before cleanup, even if Prepare + // returned a nil cleanup func. Stop can be called more than once. + t.Cleanup(func() { + if c.lagReader != nil { + c.lagReader.Stop() + } + }) readers := test.NewGoroutines() readers.Label() cleanup, err := c.Prepare(context.Background(), blipPlan("fast", "medium", "slow")) diff --git a/test/goroutine.go b/test/goroutine.go index 9f64be9..95f3aed 100644 --- a/test/goroutine.go +++ b/test/goroutine.go @@ -107,8 +107,10 @@ func wait(count func() int, n int, timeout time.Duration) int { // ... // // with a blank line between stacks. It panics if it doesn't find its own -// stack (with the label, if id is set) or if a counted stack might be +// stack (with the label, if id is set) or if a stack without fn might be // truncated, because then it might count 0 for running goroutines. +// Truncation drops the outer frames, like heartbeat.(*BlipReader).run, so a +// truncated stack doesn't contain fn. func countGoroutines(fn, id string) int { var buf bytes.Buffer pprof.Lookup("goroutine").WriteTo(&buf, 1)