Support app shutdown with exit codes

app.go

@@ -382,6 +382,7 @@ type App struct {
  donesMu sync.Mutex // guards dones and shutdownSig
  dones []chan os.Signal
  shutdownSig os.Signal
+ exitCode int
 
  osExit func(code int) // os.Exit override; used for testing only
 }
@@ -665,7 +666,11 @@ 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 {
+ code := app.run(app.Done())
+ if app.exitCode != 0 {
+ code = app.exitCode
+ }
+ if code != 0 {
  app.exit(code)
  }
 }
@@ -804,6 +809,25 @@ func (app *App) Done() <-chan os.Signal {
  return c
 }
 
+// Shutdown captures a signal and exit code.
+type Shutdown struct {
+ Signal os.Signal
+ ExitCode int
+}
+
+// DoneWithCode returns a channel of the Shutdown struct to block on after starting
+// the application. This behaves exactly as Done() does with the addition of
+// capturing an application exit code, if one exists.
+func (app *App) DoneWithCode() <-chan Shutdown {
+ c := make(chan Shutdown, 1)
+ // farm out a goroutine to wait for
+ // app.Done() channel to be populated
+ go func() {
+ c <- Shutdown{Signal: <-app.Done(), ExitCode: app.exitCode}
+ }()
+ return c
+}
+
 // StartTimeout returns the configured startup timeout. Apps default to using
 // DefaultTimeout, but users can configure this behavior using the
 // StartTimeout option.

shutdown.go

@@ -34,11 +34,27 @@ type Shutdowner interface {
 }
 
 // ShutdownOption provides a way to configure properties of the shutdown
-// process. Currently, no options have been implemented.
+// process.
 type ShutdownOption interface {
  apply(*shutdowner)
 }
 
+type exitCodeOption int
+
+func (e exitCodeOption) apply(s *shutdowner) {
+ s.app.exitCode = int(e)
+}
+
+func (e exitCodeOption) String() string {
+ return fmt.Sprintf("s.Shutdown(%v)", int(e))
+}
+
+// WithExitCode allows the user to configure the exitCode upon application
+// shutdown.
+func WithExitCode(exitCode int) ShutdownOption {
+ return exitCodeOption(exitCode)
+}
+
 type shutdowner struct {
  app *App
 }
@@ -49,6 +65,10 @@ type shutdowner struct {
 // In practice this means Shutdowner.Shutdown should not be called from an
 // fx.Invoke, but from a fx.Lifecycle.OnStart hook.
 func (s *shutdowner) Shutdown(opts ...ShutdownOption) error {
+ // Apply the options to shutdowner
+ for _, opt := range opts {
+ opt.apply(s)
+ }
  return s.app.broadcastSignal(_sigTERM)
 }
 

shutdown_test.go

@@ -87,6 +87,59 @@ func TestShutdown(t *testing.T) {
  assert.NotNil(t, <-done1, "done channel 1 did not receive signal")
  assert.NotNil(t, <-done2, "done channel 2 did not receive signal")
  })
+
+ t.Run("shutdown app with non-zero, user-specified exit code with app.Run()", func(t *testing.T) {
+ t.Parallel()
+
+ var (
+ exited bool
+ exitCode int
+ )
+ exit := func(code int) {
+ exited = true
+ exitCode = code
+ }
+
+ var s fx.Shutdowner
+ expectedExitCode := 2
+ shutdown := func(s fx.Shutdowner) error {
+ return s.Shutdown(fx.WithExitCode(expectedExitCode))
+ }
+
+ app := fxtest.New(
+ t,
+ fx.Populate(&s),
+ fx.WithExit(exit),
+ fx.Invoke(shutdown),
+ )
+
+ app.Run()
+
+ assert.Equal(t, expectedExitCode, exitCode)
+ assert.True(t, exited)
+ })
+
+ t.Run("shutdown app with non-zero, user-specified exit code without app.Run()", func(t *testing.T) {
+ t.Parallel()
+
+ var s fx.Shutdowner
+ expectedExitCode := 2
+ app := fxtest.New(
+ t,
+ fx.Populate(&s),
+ )
+
+ require.NoError(t, app.Start(context.Background()), "error starting app")
+ assert.NoError(t, s.Shutdown(fx.WithExitCode(expectedExitCode)), "error in app shutdown")
+ done1, done2 := app.DoneWithCode(), app.DoneWithCode()
+ done := app.Done()
+ defer app.Stop(context.Background())
+ result1, result2 := <-done1, <-done2
+
+ assert.NotNil(t, <-done, "done channel did not receive signal")
+ assert.Equal(t, expectedExitCode, result1.ExitCode, "mismatched exit code")
+ assert.Equal(t, expectedExitCode, result2.ExitCode, "mismatched exit code")
+ })
 }
 
 func TestDataRace(t *testing.T) {