test: Shutdown from fx.Invoke with App.Start

Adds a test for calling Shutdown from an fx.Invoke.
This is partially tested in uber-go#1075 already,
but as reported in uber-go#1074, it doesn't work if Start is used.
This adds a test for that case as well.

internal/e2e/shutdowner_wait_exitcode/main.go

@@ -0,0 +1,54 @@
+// Copyright (c) 2023 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"time"
+
+	"go.uber.org/fx"
+)
+
+func main() {
+	app := fx.New(
+		fx.Invoke(func(shutdowner fx.Shutdowner) error {
+			shutdowner.Shutdown(fx.ExitCode(20))
+			return nil
+		}),
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	if err := app.Start(ctx); err != nil {
+		log.Fatal(err)
+	}
+
+	sig := <-app.Wait()
+
+	if err := app.Stop(ctx); err != nil {
+		log.Fatal(err)
+	}
+
+	os.Exit(sig.ExitCode)
+}

internal/e2e/shutdowner_wait_exitcode/main_test.go

@@ -0,0 +1,72 @@
+// Copyright (c) 2023 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package main
+
+import (
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx/internal/testutil"
+)
+
+// Hijacks the test binary so that the test can run main() as a subprocess
+// instead of trying to compile the program and run it directly.
+func TestMain(m *testing.M) {
+	// If the test binary is named "app", then we're running as a subprocess.
+	// Otherwise, run the tests.
+	switch filepath.Base(os.Args[0]) {
+	case "app":
+		main()
+		os.Exit(0)
+	default:
+		os.Exit(m.Run())
+	}
+}
+
+// Verifies that an Fx program running with Run
+// exits with the exit code passed to Shutdowner.
+//
+// Regression test for https://github.com/uber-go/fx/issues/1074.
+func TestShutdownExitCode(t *testing.T) {
+	exe, err := os.Executable()
+	require.NoError(t, err)
+
+	out := testutil.WriteSyncer{T: t}
+
+	// Run the test binary with the name 'app' so that it runs main().
+	cmd := exec.Command(exe)
+	cmd.Args[0] = "app"
+	cmd.Stdout = &out
+	cmd.Stderr = &out
+
+	// The program should exit with code 20.
+	err = cmd.Run()
+	require.Error(t, err)
+
+	var exitErr *exec.ExitError
+	require.ErrorAs(t, err, &exitErr)
+
+	assert.Equal(t, 20, exitErr.ExitCode())
+}

shutdown_test.go

@@ -128,6 +128,29 @@ func TestShutdown(t *testing.T) {
 
 		assert.NoError(t, s.Shutdown(fx.ExitCode(2), fx.ShutdownTimeout(time.Second)))
 	})
+
+	t.Run("from invoke", func(t *testing.T) {
+		t.Parallel()
+
+		app := fxtest.New(
+			t,
+			fx.Invoke(func(s fx.Shutdowner) {
+				s.Shutdown()
+			}),
+		)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		require.NoError(t, app.Start(ctx), "error starting app")
+		defer app.Stop(ctx)
+		select {
+		case <-ctx.Done():
+			assert.Fail(t, "app did not shutdown in time")
+		case <-app.Wait():
+			// success
+		}
+	})
 }
 
 func TestDataRace(t *testing.T) {