diff --git a/metrics/repl.lag/lag_test.go b/metrics/repl.lag/lag_test.go index 5baddb9..f2656d8 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,28 @@ 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)) + // 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")) 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 +218,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 { + readers := test.NewGoroutines() + readers.Label() + _, 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 +275,16 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) { level.Collect[DOMAIN].Options[OPT_REPORT_NO_HEARTBEAT] = report } } + readers := test.NewGoroutines() + readers.Label() cleanup, err := c.Prepare(context.Background(), plan) - stopReader(t, cleanup) + stopReader(t, cleanup) // backstop if the test fails before it stops the reader 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 { @@ -266,6 +298,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 cc4be4b..4103e66 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,20 +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) - if err != nil { - if test.Build { - t.Skip("MySQL test fixture not running") - } - t.Fatal(err) - } - defer db.Close() + // Engine.Prepare only pings the database, and the repl.lag blip writer + // doesn't query it, so MySQL isn't needed + db := mock.PingDB() + 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 @@ -46,39 +36,35 @@ 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() + 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) } - 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) - if err != nil { - if test.Build { - t.Skip("MySQL test fixture not running") - } - t.Fatal(err) - } + // Engine.Prepare only pings the database, and the mock collectors don't + // query it, so MySQL isn't needed + db := mock.PingDB() 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/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/goroutine.go b/test/goroutine.go index a827d3a..95f3aed 100644 --- a/test/goroutine.go +++ b/test/goroutine.go @@ -4,34 +4,148 @@ 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 { - var buf bytes.Buffer - pprof.Lookup("goroutine").WriteTo(&buf, 2) // 2 = one stack per goroutine - n := 0 - for _, stack := range strings.Split(buf.String(), "\n\n") { - if strings.Contains(stack, fn) { - n++ - } - } - return n +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" + +// 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: 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 } -// 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. +// NewGoroutines returns a group that has no goroutines until Label is called. +func NewGoroutines() *Goroutines { + return &Goroutines{id: strconv.FormatUint(lastGoroutinesId.Add(1), 10)} +} + +// 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. 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 := Goroutines(fn) + 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 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) + _, profile, _ := strings.Cut(buf.String(), "\n") // skip "goroutine profile: total N" + 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 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, " ") // " @ ..." + c, err := strconv.Atoi(count) + if err != nil { + panic(fmt.Sprintf("invalid goroutine profile: %s: %s", err, stack)) + } + n += c + } + 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())) + } + 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 +}