Repository: uber-go/fx
Pull request #952 (2022-09-30 14:01:39)

Add support for various Hook-convertible functions
Adds new `Hook`-based helper functions:

- `StartHook`
- `StopHook`
- `StartStopHook`

and enables those function to accept any function of the form `func([context.Context] [error]` (see the `HookFunc` constraint), for example:

- `func()`
- `func() error`
- `func(context.Context)`
- `func(context.Context) error`

So, rather than users needing to do things like:

```go
func start() error {
  // ...
}

func stop(context.Context) {
  // ...
}

lifecycle.Append(fx.Hook{
  OnStart: func(context.Context) error {
    return start()
  },
  OnStop: func(ctx context.Context) error {
    stop(ctx)
    return nil
  },
})
```

which is an increasingly common pattern, users can now do:

```go
lifecycle.Append(fx.StartStopHook(start, stop))
```

or

```go
lifecycle.Append(fx.StartHook(start))
lifecycle.Append(fx.StopHook(stop))
```

and the provided functions will be converted to `func(context.Context) error` under the hood.


Patch:
diff --git a/internal/lifecycle/lifecycle.go b/internal/lifecycle/lifecycle.go
index 7c6d102..17fe24a 100644
--- a/internal/lifecycle/lifecycle.go
+++ b/internal/lifecycle/lifecycle.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Uber Technologies, Inc.
+// Copyright (c) 2022 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
@@ -25,6 +25,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -35,11 +36,89 @@ import (
 	"go.uber.org/multierr"
 )
 
+// Reflection types for each of the supported hook function signatures. These
+// are used in cases in which the Callable constraint matches a user-defined
+// function type that cannot be converted to an underlying function type with
+// a conventional conversion or type switch.
+var (
+	_reflFunc             = reflect.TypeOf(Func(nil))
+	_reflErrorFunc        = reflect.TypeOf(ErrorFunc(nil))
+	_reflContextFunc      = reflect.TypeOf(ContextFunc(nil))
+	_reflContextErrorFunc = reflect.TypeOf(ContextErrorFunc(nil))
+)
+
+// Discrete function signatures that are allowed as part of a [Callable].
+type (
+	// A Func can be converted to a ContextErrorFunc.
+	Func = func()
+	// An ErrorFunc can be converted to a ContextErrorFunc.
+	ErrorFunc = func() error
+	// A ContextFunc can be converted to a ContextErrorFunc.
+	ContextFunc = func(context.Context)
+	// A ContextErrorFunc is used as a [Hook.OnStart] or [Hook.OnStop]
+	// function.
+	ContextErrorFunc = func(context.Context) error
+)
+
+// A Callable is a constraint that matches functions that are, or can be
+// converted to, functions suitable for a Hook.
+//
+// Callable must be identical to [fx.HookFunc].
+type Callable interface {
+	~Func | ~ErrorFunc | ~ContextFunc | ~ContextErrorFunc
+}
+
+// Wrap wraps x into a ContextErrorFunc suitable for a Hook.
+func Wrap[T Callable](x T) (ContextErrorFunc, string) {
+	if x == nil {
+		return nil, ""
+	}
+
+	switch fn := any(x).(type) {
+	case Func:
+		return func(context.Context) error {
+			fn()
+			return nil
+		}, fxreflect.FuncName(x)
+	case ErrorFunc:
+		return func(context.Context) error {
+			return fn()
+		}, fxreflect.FuncName(x)
+	case ContextFunc:
+		return func(ctx context.Context) error {
+			fn(ctx)
+			return nil
+		}, fxreflect.FuncName(x)
+	case ContextErrorFunc:
+		return fn, fxreflect.FuncName(x)
+	}
+
+	// Since (1) we're already using reflect in Fx, (2) we're not particularly
+	// concerned with performance, and (3) unsafe would require discrete build
+	// targets for appengine (etc), just use reflect to convert user-defined
+	// function types to their underlying function types and then call Wrap
+	// again with the converted value.
+	reflVal := reflect.ValueOf(x)
+	switch {
+	case reflVal.CanConvert(_reflFunc):
+		return Wrap(reflVal.Convert(_reflFunc).Interface().(Func))
+	case reflVal.CanConvert(_reflErrorFunc):
+		return Wrap(reflVal.Convert(_reflErrorFunc).Interface().(ErrorFunc))
+	case reflVal.CanConvert(_reflContextFunc):
+		return Wrap(reflVal.Convert(_reflContextFunc).Interface().(ContextFunc))
+	default:
+		// Is already convertible to ContextErrorFunc.
+		return Wrap(reflVal.Convert(_reflContextErrorFunc).Interface().(ContextErrorFunc))
+	}
+}
+
 // A Hook is a pair of start and stop callbacks, either of which can be nil,
 // plus a string identifying the supplier of the hook.
 type Hook struct {
-	OnStart func(context.Context) error
-	OnStop  func(context.Context) error
+	OnStart     func(context.Context) error
+	OnStop      func(context.Context) error
+	OnStartName string
+	OnStopName  string
 
 	callerFrame fxreflect.Frame
 }
@@ -112,7 +191,11 @@ func (l *Lifecycle) Start(ctx context.Context) error {
 }
 
 func (l *Lifecycle) runStartHook(ctx context.Context, hook Hook) (runtime time.Duration, err error) {
-	funcName := fxreflect.FuncName(hook.OnStart)
+	funcName := hook.OnStartName
+	if len(funcName) == 0 {
+		funcName = fxreflect.FuncName(hook.OnStart)
+	}
+
 	l.logger.LogEvent(&fxevent.OnStartExecuting{
 		CallerName:   hook.callerFrame.Function,
 		FunctionName: funcName,
@@ -176,7 +259,10 @@ func (l *Lifecycle) Stop(ctx context.Context) error {
 }
 
 func (l *Lifecycle) runStopHook(ctx context.Context, hook Hook) (runtime time.Duration, err error) {
-	funcName := fxreflect.FuncName(hook.OnStop)
+	funcName := hook.OnStopName
+	if len(funcName) == 0 {
+		funcName = fxreflect.FuncName(hook.OnStop)
+	}
 
 	l.logger.LogEvent(&fxevent.OnStopExecuting{
 		CallerName:   hook.callerFrame.Function,
diff --git a/lifecycle.go b/lifecycle.go
index d48f7b6..72492a2 100644
--- a/lifecycle.go
+++ b/lifecycle.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Uber Technologies, Inc.
+// Copyright (c) 2022 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
@@ -26,6 +26,11 @@ import (
 	"go.uber.org/fx/internal/lifecycle"
 )
 
+// A HookFunc is a function that can be used as a [Hook].
+type HookFunc interface {
+	~func() | ~func() error | ~func(context.Context) | ~func(context.Context) error
+}
+
 // Lifecycle allows constructors to register callbacks that are executed on
 // application start and stop. See the documentation for App for details on Fx
 // applications' initialization, startup, and shutdown logic.
@@ -40,14 +45,104 @@ type Lifecycle interface {
 type Hook struct {
 	OnStart func(context.Context) error
 	OnStop  func(context.Context) error
+
+	onStartName string
+	onStopName  string
 }
 
-type lifecycleWrapper struct{ *lifecycle.Lifecycle }
+// StartHook returns a new Hook with start as its [Hook.OnStart] function,
+// wrapping its signature as needed. For example, given the following function:
+//
+//	func myhook() {
+//	  fmt.Println("hook called")
+//	}
+//
+// then calling:
+//
+//	lifecycle.Append(StartHook(myfunc))
+//
+// is functionally equivalent to calling:
+//
+//	lifecycle.Append(fx.Hook{
+//	  OnStart: func(context.Context) error {
+//	    myfunc()
+//	    return nil
+//	  },
+//	})
+//
+// The same is true for all functions that satisfy the HookFunc constraint.
+// Note that any context.Context parameter or error return will be propagated
+// as expected. If propagation is not intended, users should instead provide a
+// closure that discards the undesired value(s), or construct a Hook directly.
+func StartHook[T HookFunc](start T) Hook {
+	onstart, startname := lifecycle.Wrap(start)
+
+	return Hook{
+		OnStart:     onstart,
+		onStartName: startname,
+	}
+}
+
+// StopHook returns a new Hook with stop as its [Hook.OnStop] function,
+// wrapping its signature as needed. For example, given the following function:
+//
+//	func myhook() {
+//	  fmt.Println("hook called")
+//	}
+//
+// then calling:
+//
+//	lifecycle.Append(StopHook(myfunc))
+//
+// is functionally equivalent to calling:
+//
+//	lifecycle.Append(fx.Hook{
+//	  OnStop: func(context.Context) error {
+//	    myfunc()
+//	    return nil
+//	  },
+//	})
+//
+// The same is true for all functions that satisfy the HookFunc constraint.
+// Note that any context.Context parameter or error return will be propagated
+// as expected. If propagation is not intended, users should instead provide a
+// closure that discards the undesired value(s), or construct a Hook directly.
+func StopHook[T HookFunc](stop T) Hook {
+	onstop, stopname := lifecycle.Wrap(stop)
+
+	return Hook{
+		OnStop:     onstop,
+		onStopName: stopname,
+	}
+}
+
+// StartStopHook returns a new Hook with start as its [Hook.OnStart] function
+// and stop as its [Hook.OnStop] function, independently wrapping the signature
+// of each as needed.
+func StartStopHook[T1 HookFunc, T2 HookFunc](start T1, stop T2) Hook {
+	var (
+		onstart, startname = lifecycle.Wrap(start)
+		onstop, stopname   = lifecycle.Wrap(stop)
+	)
+
+	return Hook{
+		OnStart:     onstart,
+		OnStop:      onstop,
+		onStartName: startname,
+		onStopName:  stopname,
+	}
+}
+
+type lifecycleWrapper struct {
+	*lifecycle.Lifecycle
+}
 
 func (l *lifecycleWrapper) Append(h Hook) {
 	l.Lifecycle.Append(lifecycle.Hook{
-		OnStart: h.OnStart,
-		OnStop:  h.OnStop,
+		OnStart:     h.OnStart,
+		OnStop:      h.OnStop,
+		OnStartName: h.onStartName,
+		OnStopName:  h.onStopName,
 	})
 }
 


Test patch:
diff --git a/app_test.go b/app_test.go
index 0011272..d99a58a 100644
--- a/app_test.go
+++ b/app_test.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021 Uber Technologies, Inc.
+// Copyright (c) 2022 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
@@ -205,7 +205,8 @@ func TestNewApp(t *testing.T) {
 				func(b struct {
 					In
 					Foos []int `group:"foo"`
-				}) {
+				},
+				) {
 					assert.Len(t, b.Foos, 1)
 				},
 			),
@@ -381,7 +382,6 @@ func TestNewApp(t *testing.T) {
 			[]string{"Provided", "Provided", "Provided", "Provided", "Decorated", "Decorated", "LoggerInitialized", "Started"},
 			spy.EventTypes())
 	})
-
 }
 
 func TestWithLoggerErrorUseDefault(t *testing.T) {
@@ -793,13 +793,13 @@ func TestAppRunTimeout(t *testing.T) {
 	t.Parallel()
 
 	// Fails with an error immediately.
-	failure := func(context.Context) error {
+	failure := func() error {
 		return errors.New("great sadness")
 	}
 
 	// Builds a hook that takes much longer than the application timeout.
-	takeVeryLong := func(clock *clock.Mock) func(ctx context.Context) error {
-		return func(ctx context.Context) error {
+	takeVeryLong := func(clock *clock.Mock) func() error {
+		return func() error {
 			// We'll exceed the start and stop timeouts,
 			// and then some.
 			for i := 0; i < 3; i++ {
@@ -825,7 +825,7 @@ func TestAppRunTimeout(t *testing.T) {
 			desc: "OnStart timeout",
 			buildHooks: func(clock *clock.Mock) []Hook {
 				return []Hook{
-					{OnStart: takeVeryLong(clock)},
+					StartHook(takeVeryLong(clock)),
 				}
 			},
 			wantEventType: &fxevent.Started{},
@@ -837,8 +837,8 @@ func TestAppRunTimeout(t *testing.T) {
 				return []Hook{
 					// The hooks are separate because
 					// OnStop will not be run if that hook failed.
-					{OnStop: takeVeryLong(clock)},
-					{OnStart: failure},
+					StartHook(takeVeryLong(clock)),
+					StopHook(failure),
 				}
 			},
 			wantEventType: &fxevent.Started{},
@@ -848,7 +848,7 @@ func TestAppRunTimeout(t *testing.T) {
 			desc: "OnStop timeout",
 			buildHooks: func(clock *clock.Mock) []Hook {
 				return []Hook{
-					{OnStop: takeVeryLong(clock)},
+					StopHook(takeVeryLong(clock)),
 				}
 			},
 			wantEventType: &fxevent.Stopped{},
@@ -1428,6 +1428,275 @@ func TestValidateApp(t *testing.T) {
 	})
 }
 
+func TestHookConstructors(t *testing.T) {
+	t.Run("all", func(t *testing.T) {
+		var (
+			calls = map[string]int{
+				"start func":         0,
+				"start func err":     0,
+				"start ctx func":     0,
+				"start ctx func err": 0,
+				"stop func":          0,
+				"stop func err":      0,
+				"stop ctx func":      0,
+				"stop ctx func err":  0,
+			}
+			nilFunc             func()
+			nilErrorFunc        func() error
+			nilContextFunc      func(context.Context)
+			nilContextErrorFunc func(context.Context) error
+		)
+
+		fxtest.New(
+			t,
+			Invoke(func(lc Lifecycle) {
+				// Nil functions
+				lc.Append(StartStopHook(nilFunc, nilErrorFunc))
+				lc.Append(StartStopHook(nilContextFunc, nilContextErrorFunc))
+
+				// Start hooks
+				lc.Append(StartHook(func() {
+					calls["start func"]++
+				}))
+				lc.Append(StartHook(func() error {
+					calls["start func err"]++
+					return nil
+				}))
+				lc.Append(StartHook(func(context.Context) {
+					calls["start ctx func"]++
+				}))
+				lc.Append(StartHook(func(context.Context) error {
+					calls["start ctx func err"]++
+					return nil
+				}))
+
+				// Stop hooks
+				lc.Append(StopHook(func() {
+					calls["stop func"]++
+				}))
+				lc.Append(StopHook(func() error {
+					calls["stop func err"]++
+					return nil
+				}))
+				lc.Append(StopHook(func(context.Context) {
+					calls["stop ctx func"]++
+				}))
+				lc.Append(StopHook(func(context.Context) error {
+					calls["stop ctx func err"]++
+					return nil
+				}))
+
+				// StartStop hooks
+				lc.Append(StartStopHook(
+					func() {
+						calls["start func"]++
+					},
+					func() error {
+						calls["stop func err"]++
+						return nil
+					},
+				))
+				lc.Append(StartStopHook(
+					func() error {
+						calls["start func err"]++
+						return nil
+					},
+					func(context.Context) {
+						calls["stop ctx func"]++
+					},
+				))
+				lc.Append(StartStopHook(
+					func(context.Context) {
+						calls["start ctx func"]++
+					},
+					func(context.Context) error {
+						calls["stop ctx func err"]++
+						return nil
+					},
+				))
+				lc.Append(StartStopHook(
+					func(context.Context) error {
+						calls["start ctx func err"]++
+						return nil
+					},
+					func() {
+						calls["stop func"]++
+					},
+				))
+			}),
+		).RequireStart().RequireStop()
+
+		for name, count := range calls {
+			require.Equalf(t, 2, count, "bad call count: %s (%d)", name, count)
+		}
+	})
+
+	t.Run("start errors", func(t *testing.T) {
+		wantErr := errors.New("oh no")
+
+		// Verify that wrapped `func() error` funcs produce the expected error.
+		err := New(Invoke(func(lc Lifecycle) {
+			lc.Append(StartHook(func() error {
+				return wantErr
+			}))
+		})).Start(context.Background())
+		require.ErrorContains(t, err, wantErr.Error())
+
+		// Verify that wrapped `func(context.Context) error` funcs produce the
+		// expected error.
+		err = New(Invoke(func(lc Lifecycle) {
+			lc.Append(StartHook(func(ctx context.Context) error {
+				return wantErr
+			}))
+		})).Start(context.Background())
+		require.ErrorContains(t, err, wantErr.Error())
+	})
+
+	t.Run("start deadline", func(t *testing.T) {
+		wantCtx, cancel := context.WithTimeout(context.Background(), 123*time.Second)
+		defer cancel()
+
+		// Verify that `func(context.Context)` funcs receive the expected
+		// deadline.
+		app := New(Invoke(func(lc Lifecycle) {
+			lc.Append(StartHook(func(ctx context.Context) {
+				var (
+					want, wantOK = wantCtx.Deadline()
+					give, giveOK = ctx.Deadline()
+				)
+
+				require.Equal(t, wantOK, giveOK)
+				require.True(t, want.Equal(give))
+			}))
+		}))
+		require.NoError(t, app.Start(wantCtx))
+		require.NoError(t, app.Stop(wantCtx))
+
+		// Verify that `func(context.Context) error` funcs receive the expected
+		// deadline.
+		app = New(Invoke(func(lc Lifecycle) {
+			lc.Append(StartHook(func(ctx context.Context) error {
+				var (
+					want, wantOK = wantCtx.Deadline()
+					give, giveOK = ctx.Deadline()
+				)
+
+				require.Equal(t, wantOK, giveOK)
+				require.True(t, want.Equal(give))
+
+				return nil
+			}))
+		}))
+		require.NoError(t, app.Start(wantCtx))
+		require.NoError(t, app.Stop(wantCtx))
+	})
+
+	t.Run("stop errors", func(t *testing.T) {
+		var (
+			ctx     = context.Background()
+			wantErr = errors.New("oh no")
+		)
+
+		// Verify that wrapped `func() error` funcs produce the expected error.
+		app := New(Invoke(func(lc Lifecycle) {
+			lc.Append(StopHook(func() error {
+				return wantErr
+			}))
+		}))
+		require.NoError(t, app.Start(ctx))
+		require.ErrorContains(t, app.Stop(ctx), wantErr.Error())
+
+		// Verify that wrapped `func(context.Context) error` funcs produce the
+		// expected error.
+		app = New(Invoke(func(lc Lifecycle) {
+			lc.Append(StopHook(func(ctx context.Context) error {
+				return wantErr
+			}))
+		}))
+		require.NoError(t, app.Start(ctx))
+		require.ErrorContains(t, app.Stop(ctx), wantErr.Error())
+	})
+
+	t.Run("stop deadline", func(t *testing.T) {
+		wantCtx, cancel := context.WithTimeout(
+			context.Background(),
+			123*time.Second,
+		)
+		defer cancel()
+
+		// Verify that `func(context.Context)` funcs receive the expected
+		// deadline.
+		app := New(Invoke(func(lc Lifecycle) {
+			lc.Append(StopHook(func(ctx context.Context) {
+				var (
+					want, wantOK = wantCtx.Deadline()
+					give, giveOK = ctx.Deadline()
+				)
+
+				require.Equal(t, wantOK, giveOK)
+				require.True(t, want.Equal(give))
+			}))
+		}))
+		require.NoError(t, app.Start(wantCtx))
+		require.NoError(t, app.Stop(wantCtx))
+
+		// Verify that `func(context.Context) error` funcs receive the expected
+		// deadline.
+		app = New(Invoke(func(lc Lifecycle) {
+			lc.Append(StopHook(func(ctx context.Context) error {
+				var (
+					want, wantOK = wantCtx.Deadline()
+					give, giveOK = ctx.Deadline()
+				)
+
+				require.Equal(t, wantOK, giveOK)
+				require.True(t, want.Equal(give))
+
+				return nil
+			}))
+		}))
+		require.NoError(t, app.Start(wantCtx))
+		require.NoError(t, app.Stop(wantCtx))
+	})
+
+	t.Run("typed", func(t *testing.T) {
+		type (
+			funcType       func()
+			funcErrType    func() error
+			ctxFuncType    func(context.Context)
+			ctxFuncErrType func(context.Context) error
+		)
+
+		var (
+			calls  int
+			myFunc = funcType(func() {
+				calls++
+			})
+			myFuncErr = funcErrType(func() error {
+				calls++
+				return nil
+			})
+			myCtxFunc = ctxFuncType(func(context.Context) {
+				calls++
+			})
+			myCtxFuncErr = ctxFuncErrType(func(context.Context) error {
+				calls++
+				return nil
+			})
+		)
+
+		// Ensure that user function types that are assignable to supported
+		// base function types are converted as expected.
+		require.NotPanics(t, func() {
+			fxtest.New(t, Invoke(func(lc Lifecycle) {
+				lc.Append(StartStopHook(myFunc, myFuncErr))
+				lc.Append(StartStopHook(myCtxFunc, myCtxFuncErr))
+			})).RequireStart().RequireStop()
+		})
+		require.Equal(t, 4, calls)
+	})
+}
+
 func TestDone(t *testing.T) {
 	t.Parallel()
 
@@ -1661,7 +1930,6 @@ func TestErrorHook(t *testing.T) {
 		assert.Contains(t, graphStr, `"fx_test.B" [color=red];`)
 		assert.Contains(t, graphStr, `"fx_test.A" [color=orange];`)
 	})
-
 }
 
 func TestOptionString(t *testing.T) {