bootstrap: Make bootstrapper stoppable to plug Goroutine leaks

The bootstrapper is currently leaking multiple Goroutines in the
upgrader because we never make sure to clean it up after tests have
finished. Add its existing `Stop()` function to the upgrader interface
and call is as required to plug these leaks.

2 files changed, 34 insertions, 11 deletions

diff --git a/internal/bootstrap/bootstrap.go b/internal/bootstrap/bootstrap.go
index ef414184d..95e43aa10 100644
--- a/internal/bootstrap/bootstrap.go
+++ b/internal/bootstrap/bootstrap.go
@@ -38,6 +38,7 @@ type upgrader interface {
 	HasParent() bool
 	Ready() error
 	Upgrade() error
+	Stop()
 }
 
 // New performs tableflip initialization
@@ -172,6 +173,7 @@ func (b *Bootstrap) Wait(gracefulTimeout time.Duration) error {
 		err = fmt.Errorf("graceful upgrade: %v", waitError)
 	case s := <-immediateShutdown:
 		err = fmt.Errorf("received signal %q", s)
+		b.upgrader.Stop()
 	case err = <-b.errChan:
 	}
 
diff --git a/internal/bootstrap/bootstrap_test.go b/internal/bootstrap/bootstrap_test.go
index 8b870395c..1fb6c3485 100644
--- a/internal/bootstrap/bootstrap_test.go
+++ b/internal/bootstrap/bootstrap_test.go
@@ -26,6 +26,8 @@ func (m *mockUpgrader) Exit() <-chan struct{} {
 	return m.exit
 }
 
+func (m *mockUpgrader) Stop() {}
+
 func (m *mockUpgrader) HasParent() bool {
 	return m.hasParent
 }
@@ -104,7 +106,10 @@ func waitWithTimeout(t *testing.T, waitCh <-chan error, timeout time.Duration) e
 }
 
 func TestImmediateTerminationOnSocketError(t *testing.T) {
-	b, server := makeBootstrap(t)
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	b, server := makeBootstrap(t, ctx)
 
 	waitCh := make(chan error)
 	go func() { waitCh <- b.Wait(2 * time.Second) }()
@@ -119,7 +124,10 @@ func TestImmediateTerminationOnSocketError(t *testing.T) {
 func TestImmediateTerminationOnSignal(t *testing.T) {
 	for _, sig := range []syscall.Signal{syscall.SIGTERM, syscall.SIGINT} {
 		t.Run(sig.String(), func(t *testing.T) {
-			b, server := makeBootstrap(t)
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+
+			b, server := makeBootstrap(t, ctx)
 
 			done := server.slowRequest(3 * time.Minute)
 
@@ -146,7 +154,10 @@ func TestImmediateTerminationOnSignal(t *testing.T) {
 }
 
 func TestGracefulTerminationStuck(t *testing.T) {
-	b, server := makeBootstrap(t)
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	b, server := makeBootstrap(t, ctx)
 
 	err := testGracefulUpdate(t, server, b, 3*time.Second, 2*time.Second, nil)
 	require.Contains(t, err.Error(), "grace period expired")
@@ -158,7 +169,9 @@ func TestGracefulTerminationWithSignals(t *testing.T) {
 
 	for _, sig := range []syscall.Signal{syscall.SIGTERM, syscall.SIGINT} {
 		t.Run(sig.String(), func(t *testing.T) {
-			b, server := makeBootstrap(t)
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+			b, server := makeBootstrap(t, ctx)
 
 			err := testGracefulUpdate(t, server, b, 1*time.Second, 2*time.Second, func() {
 				require.NoError(t, self.Signal(sig))
@@ -169,7 +182,9 @@ func TestGracefulTerminationWithSignals(t *testing.T) {
 }
 
 func TestGracefulTerminationServerErrors(t *testing.T) {
-	b, server := makeBootstrap(t)
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+	b, server := makeBootstrap(t, ctx)
 
 	done := make(chan error, 1)
 	// This is a simulation of receiving a listener error during waitGracePeriod
@@ -192,7 +207,9 @@ func TestGracefulTerminationServerErrors(t *testing.T) {
 }
 
 func TestGracefulTermination(t *testing.T) {
-	b, server := makeBootstrap(t)
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+	b, server := makeBootstrap(t, ctx)
 
 	// Using server.Close we bypass the graceful shutdown faking a completed shutdown
 	b.StopAction = func() { server.server.Close() }
@@ -202,21 +219,21 @@ func TestGracefulTermination(t *testing.T) {
 }
 
 func TestPortReuse(t *testing.T) {
-	var addr string
-
 	b, err := New()
 	require.NoError(t, err)
 
 	l, err := b.listen("tcp", "localhost:")
 	require.NoError(t, err, "failed to bind")
 
-	addr = l.Addr().String()
+	addr := l.Addr().String()
 	_, port, err := net.SplitHostPort(addr)
 	require.NoError(t, err)
 
 	l, err = b.listen("tcp", "localhost:"+port)
 	require.NoError(t, err, "failed to bind")
 	require.NoError(t, l.Close())
+
+	b.upgrader.Stop()
 }
 
 func testGracefulUpdate(t *testing.T, server *testServer, b *Bootstrap, waitTimeout, gracefulWait time.Duration, duringGracePeriodCallback func()) error {
@@ -251,7 +268,7 @@ func testGracefulUpdate(t *testing.T, server *testServer, b *Bootstrap, waitTime
 	return waitErr
 }
 
-func makeBootstrap(t *testing.T) (*Bootstrap, *testServer) {
+func makeBootstrap(t *testing.T, ctx context.Context) (*Bootstrap, *testServer) {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
 		w.WriteHeader(200)
@@ -260,12 +277,16 @@ func makeBootstrap(t *testing.T) (*Bootstrap, *testServer) {
 		sec, err := strconv.Atoi(r.URL.Query().Get("seconds"))
 		require.NoError(t, err)
 
-		time.Sleep(time.Duration(sec) * time.Second)
+		select {
+		case <-ctx.Done():
+		case <-time.After(time.Duration(sec) * time.Second):
+		}
 
 		w.WriteHeader(200)
 	})
 
 	s := http.Server{Handler: mux}
+	t.Cleanup(func() { testhelper.MustClose(t, &s) })
 	u := &mockUpgrader{exit: make(chan struct{})}
 
 	b, err := _new(u, net.Listen, false)