From 1c55c526590c5563dd4e0ebecab6ca9b0fbd5ad7 Mon Sep 17 00:00:00 2001 From: jmills Date: Thu, 4 Aug 2022 23:14:56 +0000 Subject: [PATCH] Implements ShutdownCode option and ShutdownSignal os.Signal wrapper --- shutdown.go | 51 +++++++++++++++++++++++++++++++++-- shutdown_code_example_test.go | 48 +++++++++++++++++++++++++++++++++ shutdown_test.go | 50 ++++++++++++++++++++++++++++++++++ 3 files changed, 147 insertions(+), 2 deletions(-) create mode 100644 shutdown_code_example_test.go diff --git a/shutdown.go b/shutdown.go index d5b8488c0c..6cab57baff 100644 --- a/shutdown.go +++ b/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 { diff --git a/shutdown_code_example_test.go b/shutdown_code_example_test.go new file mode 100644 index 0000000000..37105464f5 --- /dev/null +++ b/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) +} diff --git a/shutdown_test.go b/shutdown_test.go index b6af93f131..d544c6546c 100644 --- a/shutdown_test.go +++ b/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) {