ensure withTimeout doesn't have if runtime.Goexit() is called from wi…

…thin the callback. (#890)

app.go

@@ -765,9 +765,26 @@ type withTimeoutParams struct {
 	lifecycle *lifecycleWrapper
 }
 
+// errHookCallbackExited is returned when a hook callback does not finish executing
+var errHookCallbackExited = fmt.Errorf("goroutine exited without returning")
+
 func withTimeout(ctx context.Context, param *withTimeoutParams) error {
 	c := make(chan error, 1)
-	go func() { c <- param.callback(ctx) }()
+	go func() {
+		// If runtime.Goexit() is called from within the callback
+		// then nothing is written to the chan.
+		// However the defer will still be called, so we can write to the chan,
+		// to avoid hanging until the timeout is reached.
+		callbackExited := false
+		defer func() {
+			if !callbackExited {
+				c <- errHookCallbackExited
+			}
+		}()
+
+		c <- param.callback(ctx)
+		callbackExited = true
+	}()
 
 	var err error
 
@@ -782,7 +799,7 @@ func withTimeout(ctx context.Context, param *withTimeoutParams) error {
 			err = ctx.Err()
 		}
 	}
-	if err != context.DeadlineExceeded {
+	if err != context.DeadlineExceeded && err != errHookCallbackExited {
 		return err
 	}
 	// On timeout, report running hook's caller and recorded

app_test.go

@@ -30,6 +30,7 @@ import (
 	"log"
 	"os"
 	"reflect"
+	"runtime"
 	"strings"
 	"testing"
 	"time"
@@ -381,6 +382,7 @@ func TestNewApp(t *testing.T) {
 			[]string{"Provided", "Provided", "Provided", "Provided", "Decorated", "Decorated", "LoggerInitialized", "Started"},
 			spy.EventTypes())
 	})
+
 }
 
 func TestWithLoggerErrorUseDefault(t *testing.T) {
@@ -1262,6 +1264,24 @@ func TestAppStart(t *testing.T) {
 		assert.Contains(t, err.Error(), "go.uber.org/fx_test.TestAppStart")
 		assert.Contains(t, err.Error(), "/app_test.go")
 	})
+
+	t.Run("HookGoroutineExitsErrorMsg", func(t *testing.T) {
+		t.Parallel()
+
+		addHook := func(lc Lifecycle) {
+			lc.Append(Hook{
+				OnStart: func(ctx context.Context) error {
+					runtime.Goexit()
+					return nil
+				},
+			})
+		}
+		app := fxtest.New(t,
+			Invoke(addHook),
+		)
+		err := app.Start(context.Background()).Error()
+		assert.Contains(t, err, "OnStart hook added by go.uber.org/fx_test.TestAppStart.func10.1 failed: goroutine exited without returning")
+	})
 }
 
 func TestAppStop(t *testing.T) {