From b1fc1c0931b8ce76e0805d3974a47af6fea4c2bd Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <mail@abhinavg.net>
Date: Sat, 29 Apr 2023 10:49:18 -0700
Subject: [PATCH] app.Run: Respect Shutdower exit code

We added support for changing the exit code for a Shutdowner with the
ExitCode option in #989, but this was somehow not respected by App.Run.

This changes App.Run to use the same underlying machinery (`Wait()`)
to decide on the exit code to use.

Resolves #1074
---
 app.go               | 9 +++++----
 app_internal_test.go | 5 ++---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/app.go b/app.go
index d2b8caf17..bc3ed09d7 100644
--- a/app.go
+++ b/app.go
@@ -574,12 +574,12 @@ func (app *App) Run() {
 	// Historically, we do not os.Exit(0) even though most applications
 	// cede control to Fx with they call app.Run. To avoid a breaking
 	// change, never os.Exit for success.
-	if code := app.run(app.Done()); code != 0 {
+	if code := app.run(app.Wait()); code != 0 {
 		app.exit(code)
 	}
 }
 
-func (app *App) run(done <-chan os.Signal) (exitCode int) {
+func (app *App) run(done <-chan ShutdownSignal) (exitCode int) {
 	startCtx, cancel := app.clock.WithTimeout(context.Background(), app.StartTimeout())
 	defer cancel()
 
@@ -588,7 +588,8 @@ func (app *App) run(done <-chan os.Signal) (exitCode int) {
 	}
 
 	sig := <-done
-	app.log().LogEvent(&fxevent.Stopping{Signal: sig})
+	app.log().LogEvent(&fxevent.Stopping{Signal: sig.Signal})
+	exitCode = sig.ExitCode
 
 	stopCtx, cancel := app.clock.WithTimeout(context.Background(), app.StopTimeout())
 	defer cancel()
@@ -597,7 +598,7 @@ func (app *App) run(done <-chan os.Signal) (exitCode int) {
 		return 1
 	}
 
-	return 0
+	return exitCode
 }
 
 // Err returns any error encountered during New's initialization. See the
diff --git a/app_internal_test.go b/app_internal_test.go
index 6e54f8638..a377dd711 100644
--- a/app_internal_test.go
+++ b/app_internal_test.go
@@ -23,7 +23,6 @@ package fx
 import (
 	"errors"
 	"fmt"
-	"os"
 	"sync"
 	"testing"
 
@@ -42,7 +41,7 @@ func TestAppRun(t *testing.T) {
 	app := New(
 		WithLogger(func() fxevent.Logger { return spy }),
 	)
-	done := make(chan os.Signal)
+	done := make(chan ShutdownSignal)
 
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -51,7 +50,7 @@ func TestAppRun(t *testing.T) {
 		app.run(done)
 	}()
 
-	done <- _sigINT
+	done <- ShutdownSignal{Signal: _sigINT}
 	wg.Wait()
 
 	assert.Equal(t, []string{