quic: don't consider goroutines running when tests start as leaked

neild · neild · commit a130fcc1c15a · 2024-04-04T20:10:37.000Z
Change-Id: I138e284ee74c9402402f564d25e50ab753c51e9e Reviewed-on: https://go-review.googlesource.com/c/net/+/576536 Reviewed-by: Chressie Himpel <chressie@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/quic/main_test.go b/quic/main_test.go
@@ -16,6 +16,20 @@ import (
 )
 
 func TestMain(m *testing.M) {
+	// Add all goroutines running at the start of the test to the set
+	// of not-leaked goroutines. This includes TestMain, and anything else
+	// that might have been started by test infrastructure.
+	skip := [][]byte{
+		[]byte("created by os/signal.Notify"),
+		[]byte("gotraceback_test.go"),
+	}
+	buf := make([]byte, 2<<20)
+	buf = buf[:runtime.Stack(buf, true)]
+	for _, g := range bytes.Split(buf, []byte("\n\n")) {
+		id, _, _ := bytes.Cut(g, []byte("["))
+		skip = append(skip, id)
+	}
+
 	defer os.Exit(m.Run())
 
 	// Look for leaked goroutines.
@@ -34,12 +48,13 @@ func TestMain(m *testing.M) {
 		buf = buf[:runtime.Stack(buf, true)]
 		leaked := false
 		for _, g := range bytes.Split(buf, []byte("\n\n")) {
-			if bytes.Contains(g, []byte("quic.TestMain")) ||
-				bytes.Contains(g, []byte("created by os/signal.Notify")) ||
-				bytes.Contains(g, []byte("gotraceback_test.go")) {
-				continue
-			}
 			leaked = true
+			for _, s := range skip {
+				if bytes.Contains(g, s) {
+					leaked = false
+					break
+				}
+			}
 		}
 		if !leaked {
 			break
