Implements ShutdownCode option and ShutdownSignal os.Signal wrapper

Author: jasonmills

shutdown.go

@@ -39,8 +39,48 @@ type ShutdownOption interface {
 	apply(*shutdowner)
 }
 
+type shutdownCode int
+
+func (c shutdownCode) apply(s *shutdowner) {
+	s.exitCode = int(c)
+}
+
+// ShutdownCode implements a shutdown option that allows a user specify the
+// os.Exit code that an application should exit with.
+func ShutdownCode(code int) ShutdownOption {
+	return shutdownCode(code)
+}
+
 type shutdowner struct {
-	app *App
+	exitCode int
+	app      *App
+}
+
+// ShutdownerSignal defines an os.Signal interface with an extension that allows
+// the querying of a exit code defined for the signal. ExitCode defaults to 0.
+type ShutdownerSignal interface {
+	os.Signal
+	ExitCode() int
+}
+
+type shutdownSignal struct {
+	signal   os.Signal
+	exitCode int
+}
+
+func (s shutdownSignal) String() string { return s.signal.String() }
+func (s shutdownSignal) Signal()        { s.signal.Signal() }
+func (s shutdownSignal) ExitCode() int  { return s.exitCode }
+
+// ShutdownSignal will return a ShutdownerSignal for a given os.Signal. If
+// the signal type was not originally a ShutdownerSignal, a new ShutdownerSignal
+// will be created from the os.Signal, and exit code will be set to default.
+func ShutdownSignal(signal os.Signal) ShutdownerSignal {
+	if s, ok := signal.(ShutdownerSignal); ok {
+		return s
+	}
+
+	return &shutdownSignal{signal: signal}
 }
 
 // Shutdown broadcasts a signal to all of the application's Done channels
@@ -49,7 +89,14 @@ 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 {
-	return s.app.broadcastSignal(_sigTERM)
+	for _, opt := range opts {
+		opt.apply(s)
+	}
+
+	return s.app.broadcastSignal(&shutdownSignal{
+		exitCode: s.exitCode,
+		signal:   _sigTERM,
+	})
 }
 
 func (app *App) shutdowner() Shutdowner {

shutdown_code_example_test.go

@@ -0,0 +1,48 @@
+// Copyright (c) 2019 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 fx_test
+
+import (
+	"fmt"
+
+	"go.uber.org/fx"
+)
+
+func ExampleShutdownCode() {
+	app := fx.New(
+		fx.Invoke(func(shutdowner fx.Shutdowner) {
+			// Call the shutdowner Shutdown method with a shutdown code
+			// option
+			shutdowner.Shutdown(fx.ShutdownCode(1))
+		}),
+	)
+
+	app.Run()
+
+	// Extract the shutdown signal from the os.Signal returned by app.Done
+	signal := fx.ShutdownSignal(<-app.Done())
+
+	// Retrieve the exit code
+	fmt.Printf("os.Exit(%v)\n", signal.ExitCode())
+
+	// Output:
+	// os.Exit(1)
+}

shutdown_test.go

@@ -22,6 +22,7 @@ package fx_test
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"testing"
 
@@ -87,6 +88,55 @@ 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 exit code(s)", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("default", func(t *testing.T) {
+			t.Parallel()
+			var s fx.Shutdowner
+			app := fxtest.New(
+				t,
+				fx.Populate(&s),
+			)
+
+			done := app.Done()
+			defer app.RequireStart().RequireStop()
+
+			assert.NoError(t, s.Shutdown(), "error in app shutdown")
+			signal := <-done
+
+			assert.NotNil(t, signal, "done channel did not receive signal")
+			shutdownSignal := fx.ShutdownSignal(signal)
+			assert.NotNil(t, shutdownSignal, "done channel did not send shutdown signal")
+			assert.NotPanics(t, shutdownSignal.Signal)
+			assert.Zero(t, shutdownSignal.ExitCode())
+			assert.NotNil(t, fx.ShutdownSignal(nil))
+		})
+
+		for expected := 0; expected < 3; expected++ {
+			expected := expected
+			t.Run(fmt.Sprintf("with exit code %v", expected), func(t *testing.T) {
+				t.Parallel()
+				var s fx.Shutdowner
+				app := fxtest.New(
+					t,
+					fx.Populate(&s),
+				)
+
+				done := app.Done()
+				defer app.RequireStart().RequireStop()
+
+				assert.NoError(t, s.Shutdown(fx.ShutdownCode(expected)), "error in app shutdown")
+				signal := <-done
+
+				assert.NotNil(t, signal, "done channel did not receive signal")
+				shutdownSignal := fx.ShutdownSignal(signal)
+				assert.NotNil(t, shutdownSignal, "done channel did not send shutdown signal")
+				assert.Equal(t, expected, shutdownSignal.ExitCode())
+			})
+		}
+	})
 }
 
 func TestDataRace(t *testing.T) {