From 02a2658c222c3250ebaad9a9b6386d210a059e4e Mon Sep 17 00:00:00 2001
From: Misha Abramovich <abram13504@yandex.ru>
Date: Wed, 7 Dec 2022 21:41:56 +0100
Subject: [PATCH] Add fx.ShutdownError option

---
 shutdown.go      | 49 +++++++++++++++++++++++++++++++++++++++++-------
 shutdown_test.go | 23 +++++++++++++++++++++++
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/shutdown.go b/shutdown.go
index aa81e68d3..564ab7cb3 100644
--- a/shutdown.go
+++ b/shutdown.go
@@ -23,6 +23,8 @@ package fx
 import (
 	"context"
 	"time"
+
+	"go.uber.org/multierr"
 )
 
 // Shutdowner provides a method that can manually trigger the shutdown of the
@@ -34,17 +36,11 @@ 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 (code exitCodeOption) apply(s *shutdowner) {
-	s.exitCode = int(code)
-}
-
 var _ ShutdownOption = exitCodeOption(0)
 
 // ExitCode is a [ShutdownOption] that may be passed to the Shutdown method of the
@@ -71,6 +67,45 @@ func ShutdownTimeout(timeout time.Duration) ShutdownOption {
 	return shutdownTimeoutOption(timeout)
 }
 
+type shutdownErrorOption []error
+
+func (errs shutdownErrorOption) apply(s *shutdowner) {
+	s.app.err = multierr.Append(s.app.err, multierr.Combine(errs...))
+}
+
+// ShutdownError registers any number of errors with the application shutdown.
+// If more than one error is given, the errors are combined into a
+// single error. Similar to invocations, errors are applied in order.
+//
+// You can use these errors, for example, to decide what to do after the app shutdown.
+//
+//	customErr := errors.New("something went wrong")
+//	app := fx.New(
+//		...
+//		fx.Provide(func(s fx.Shutdowner, a A) B {
+//			s.Shutdown(fx.ShutdownError(customErr))
+//	    }),
+//	    ...
+//	)
+//	err := app.Start(context.Background())
+//	if err != nil {
+//		panic(err)
+//	}
+//	defer app.Stop(context.Background())
+//
+//	if err := app.Err(); errors.Is(err, customErr) {
+//	   // custom logic here
+//	}
+func ShutdownError(errs ...error) ShutdownOption {
+	return shutdownErrorOption(errs)
+}
+
+type exitCodeOption int
+
+func (code exitCodeOption) apply(s *shutdowner) {
+	s.exitCode = int(code)
+}
+
 type shutdowner struct {
 	app             *App
 	exitCode        int
diff --git a/shutdown_test.go b/shutdown_test.go
index 1f322a2cb..165d07486 100644
--- a/shutdown_test.go
+++ b/shutdown_test.go
@@ -22,6 +22,7 @@ package fx_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"testing"
@@ -128,6 +129,28 @@ func TestShutdown(t *testing.T) {
 
 		assert.NoError(t, s.Shutdown(fx.ExitCode(2), fx.ShutdownTimeout(time.Second)))
 	})
+
+	t.Run("with shutdown error", func(t *testing.T) {
+		t.Parallel()
+
+		var s fx.Shutdowner
+		app := fxtest.New(
+			t,
+			fx.Populate(&s),
+		)
+
+		done := app.Done()
+		wait := app.Wait()
+		defer app.RequireStart().RequireStop()
+
+		var expectedError = errors.New("shutdown error")
+
+		assert.NoError(t, s.Shutdown(fx.ShutdownError(expectedError)), "error in app shutdown")
+		assert.NotNil(t, <-done, "done channel did not receive signal")
+		assert.NotNil(t, <-wait, "wait channel did not receive signal")
+		assert.ErrorIs(t, app.Err(), expectedError,
+			"unexpected error, expected: %q, got: %q", expectedError, app.Err())
+	})
 }
 
 func TestDataRace(t *testing.T) {