quic: avoid leaking tls goroutines in tests

neild · neild · commit 65efbad9474a · 2023-12-07T01:27:06.000Z
Change-Id: Iaf273294ba3245bfeb387a72e068c048d0fcf93a Reviewed-on: https://go-review.googlesource.com/c/net/+/547736 Reviewed-by: Jonathan Amsterdam <jba@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/internal/quic/conn_test.go b/internal/quic/conn_test.go
@@ -279,6 +279,9 @@ func newTestConnForConn(t *testing.T, endpoint *testEndpoint, conn *Conn) *testC
 	}
 	tc.peerTLSConn.SetTransportParameters(marshalTransportParameters(peerProvidedParams))
 	tc.peerTLSConn.Start(context.Background())
+	t.Cleanup(func() {
+		tc.peerTLSConn.Close()
+	})
 
 	return tc
 }
diff --git a/internal/quic/main_test.go b/internal/quic/main_test.go
@@ -0,0 +1,52 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.21
+
+package quic
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"runtime"
+	"testing"
+	"time"
+)
+
+func TestMain(m *testing.M) {
+	defer os.Exit(m.Run())
+
+	// Look for leaked goroutines.
+	//
+	// Checking after every test makes it easier to tell which test is the culprit,
+	// but checking once at the end is faster and less likely to miss something.
+	start := time.Now()
+	warned := false
+	for {
+		buf := make([]byte, 2<<20)
+		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
+		}
+		if !leaked {
+			break
+		}
+		if !warned && time.Since(start) > 1*time.Second {
+			// Print a warning quickly, in case this is an interactive session.
+			// Keep waiting until the test times out, in case this is a slow trybot.
+			fmt.Printf("Tests seem to have leaked some goroutines, still waiting.\n\n")
+			fmt.Print(string(buf))
+			warned = true
+		}
+		// Goroutines might still be shutting down.
+		time.Sleep(1 * time.Millisecond)
+	}
+}
diff --git a/internal/quic/retry_test.go b/internal/quic/retry_test.go
@@ -533,7 +533,9 @@ func initialClientCrypto(t *testing.T, e *testEndpoint, p transportParameters) [
 	tlsClient := tls.QUICClient(config)
 	tlsClient.SetTransportParameters(marshalTransportParameters(p))
 	tlsClient.Start(context.Background())
-	//defer tlsClient.Close()
+	t.Cleanup(func() {
+		tlsClient.Close()
+	})
 	e.peerTLSConn = tlsClient
 	var data []byte
 	for {

Original file line number	Diff line number	Diff line change
`@@ -279,6 +279,9 @@ func newTestConnForConn(t testing.T, endpoint testEndpoint, conn Conn) testC`
`279`	`279`	`}`
`280`	`280`	`tc.peerTLSConn.SetTransportParameters(marshalTransportParameters(peerProvidedParams))`
`281`	`281`	`tc.peerTLSConn.Start(context.Background())`
	`282`	`+ t.Cleanup(func() {`
	`283`	`+ tc.peerTLSConn.Close()`
	`284`	`+ })`
`282`	`285`
`283`	`286`	`return tc`
`284`	`287`	`}`