Repository: uber-go/fx
Pull request #953 (2022-09-30 14:07:51)

Add support for context- and error-less annotated hook funcs
Similar to #952, this adds support for any `HookFunc`-compatible funcs as part of `fx.OnStart`/`fx.OnStop` annotations.

~Unrelated change: also changes bind addresses for tests to be localhost.~ Was done in #960 

codecov[bot]: # [Codecov](https://codecov.io/gh/uber-go/fx/pull/953?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) Report
> Merging [#953](https://codecov.io/gh/uber-go/fx/pull/953?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (65ba282) into [master](https://codecov.io/gh/uber-go/fx/commit/ee6d9f3dfde6f1d436a845a6bbcdf5b35ec8d423?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (ee6d9f3) will **increase** coverage by `0.11%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #953      +/-   ##
==========================================
+ Coverage   98.67%   98.78%   +0.11%     
==========================================
  Files          38       32       -6     
  Lines        1581     1478     -103     
==========================================
- Hits         1560     1460     -100     
+ Misses         15       13       -2     
+ Partials        6        5       -1     
```


| [Impacted Files](https://codecov.io/gh/uber-go/fx/pull/953?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) | Coverage Δ | |
|---|---|---|
| [annotated.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-YW5ub3RhdGVkLmdv) | `99.69% <100.00%> (-0.01%)` | :arrow_down: |
| [docs/internal/httptest/http.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC9odHRwdGVzdC9odHRwLmdv) | | |
| [docs/internal/test/fake.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC90ZXN0L2Zha2UuZ28=) | | |
| [docs/internal/iotest/read.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC9pb3Rlc3QvcmVhZC5nbw==) | | |
| [docs/internal/exectest/cmd.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC9leGVjdGVzdC9jbWQuZ28=) | | |
| [docs/internal/exectest/output.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC9leGVjdGVzdC9vdXRwdXQuZ28=) | | |
| [docs/internal/apptest/run.go](https://codecov.io/gh/uber-go/fx/pull/953/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-ZG9jcy9pbnRlcm5hbC9hcHB0ZXN0L3J1bi5nbw==) | | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)


Patch:
diff --git a/annotated.go b/annotated.go
index 12bbbcc..7937ce9 100644
--- a/annotated.go
+++ b/annotated.go
@@ -250,23 +250,14 @@ func (la *lifecycleHookAnnotation) apply(ann *annotated) error {
 		)
 	}
 
-	if ft.NumIn() < 1 || ft.In(0) != _typeOfContext {
-		return fmt.Errorf(
-			"first argument of hook must be context.Context, got %v (%T)",
-			la.Target,
-			la.Target,
-		)
-	}
-
-	hasOut := ft.NumOut() == 1
-	returnsErr := hasOut && ft.Out(0) == _typeOfError
-
-	if !hasOut || !returnsErr {
-		return fmt.Errorf(
-			"hooks must return only an error type, got %v (%T)",
-			la.Target,
-			la.Target,
-		)
+	if n := ft.NumOut(); n > 0 {
+		if n > 1 || ft.Out(0) != _typeOfError {
+			return fmt.Errorf(
+				"optional hook return may only be an error, got %v (%T)",
+				la.Target,
+				la.Target,
+			)
+		}
 	}
 
 	if ft.IsVariadic() {
@@ -349,10 +340,18 @@ func (la *lifecycleHookAnnotation) parameters(results ...reflect.Type) (
 	}
 
 	ft := reflect.TypeOf(la.Target)
-	resolverIdx := make([]argSource, 1)
+	resolverIdx := make([]argSource, 0)
 
-	for i := 1; i < ft.NumIn(); i++ {
+	for i := 0; i < ft.NumIn(); i++ {
 		t := ft.In(i)
+
+		// If the first parameter is a context, add a sentinel value and skip
+		// it; it will be injected as part of the actual function call.
+		if i == 0 && t == _typeOfContext {
+			resolverIdx = append(resolverIdx, argSource{})
+			continue
+		}
+
 		result, isProvidedByResults := resultMap[t]
 
 		if isProvidedByResults {
@@ -383,7 +382,7 @@ func (la *lifecycleHookAnnotation) parameters(results ...reflect.Type) (
 	) (lc Lifecycle, remapped []reflect.Value) {
 		remapped = make([]reflect.Value, ft.NumIn())
 
-		if len(args) != 0 {
+		if len(args) > 0 {
 			var (
 				results reflect.Value
 				p       = args[0]
@@ -394,7 +393,13 @@ func (la *lifecycleHookAnnotation) parameters(results ...reflect.Type) (
 			}
 
 			lc, _ = p.FieldByName("Lifecycle").Interface().(Lifecycle)
-			for i := 1; i < ft.NumIn(); i++ {
+			for i := 0; i < ft.NumIn(); i++ {
+				// If the first field is a context, skip it; as it will be
+				// injected separately as part of the actual function call.
+				if i == 0 && ft.In(i) == _typeOfContext {
+					continue
+				}
+
 				resolver := resolverIdx[i]
 				source := p
 				if resolver.result {
@@ -428,13 +433,23 @@ func (la *lifecycleHookAnnotation) Build(results ...reflect.Type) reflect.Value
 		}
 	}
 
-	origFn := reflect.ValueOf(la.Target)
-	newFnType := reflect.FuncOf(params, nil, false)
+	var (
+		origFn    = reflect.ValueOf(la.Target)
+		origFnT   = origFn.Type()
+		newFnType = reflect.FuncOf(params, nil, false)
+	)
+
 	newFn := reflect.MakeFunc(newFnType, func(args []reflect.Value) []reflect.Value {
-		var lc Lifecycle
-		lc, args = paramMap(args)
+		lc, args := paramMap(args)
+
 		hookFn := func(ctx context.Context) (err error) {
-			args[0] = reflect.ValueOf(ctx)
+			// If the hook function has multiple parameters, and the first
+			// parameter is a context, inject the provided context.
+			if origFnT.NumIn() > 0 {
+				if origFnT.In(0) == _typeOfContext {
+					args[0] = reflect.ValueOf(ctx)
+				}
+			}
 
 			results := origFn.Call(args)
 			if len(results) > 0 && results[0].Type() == _typeOfError {
diff --git a/go.mod b/go.mod
index cab2fc5..ffa4059 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/benbjohnson/clock v1.3.0
 	github.com/stretchr/testify v1.7.1
+	go.uber.org/atomic v1.6.0
 	go.uber.org/dig v1.15.0
 	go.uber.org/goleak v1.1.11
 	go.uber.org/multierr v1.5.0
@@ -15,6 +16,5 @@ require (
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	go.uber.org/atomic v1.6.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )


Test patch:
diff --git a/annotated_test.go b/annotated_test.go
index 4b4c6db..f7ebc8e 100644
--- a/annotated_test.go
+++ b/annotated_test.go
@@ -31,6 +31,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 	"go.uber.org/fx"
 	"go.uber.org/fx/fxevent"
 	"go.uber.org/fx/fxtest"
@@ -1239,7 +1240,7 @@ func TestHookAnnotations(t *testing.T) {
 				in.WriteString("decorated")
 				return in
 			},
-			fx.OnStart(func(_ context.Context, a A) error {
+			fx.OnStart(func(_ context.Context, _ A) error {
 				called = assert.Equal(t, "decorated", buf.String())
 				return nil
 			}),
@@ -1369,14 +1370,6 @@ func TestHookAnnotationFailures(t *testing.T) {
 				fx.OnStart(func(context.Context) error { return nil }),
 			),
 		},
-		{
-			name:        "with hook that doesn't return an error",
-			errContains: "must return only an error",
-			annotation: fx.Annotate(
-				func() A { return nil },
-				fx.OnStart(func(context.Context) {}),
-			),
-		},
 		{
 			name:        "with constructor that errors",
 			errContains: "hooks should not be installed",
@@ -1400,11 +1393,23 @@ func TestHookAnnotationFailures(t *testing.T) {
 			),
 		},
 		{
-			name:        "without context.Context as first parameter",
-			errContains: "must be context.Context",
+			name:        "invalid return: non-error return",
+			errContains: "optional hook return may only be an error",
+			annotation: fx.Annotate(
+				func() A { return nil },
+				fx.OnStart(func(context.Context) any {
+					return nil
+				}),
+			),
+		},
+		{
+			name:        "invalid return: too many returns",
+			errContains: "optional hook return may only be an error",
 			annotation: fx.Annotate(
 				func() A { return nil },
-				fx.OnStart(func() {}),
+				fx.OnStart(func(context.Context) (any, any) {
+					return nil, nil
+				}),
 			),
 		},
 		{
@@ -1450,3 +1455,68 @@ func TestHookAnnotationFailures(t *testing.T) {
 		})
 	}
 }
+
+func TestHookAnnotationFunctionFlexibility(t *testing.T) {
+	type A interface{}
+
+	table := []struct {
+		name       string
+		annotation interface{}
+	}{
+		{
+			name: "without error return",
+			annotation: fx.Annotate(
+				func() A { return nil },
+				fx.OnStart(func(_ context.Context, called *atomic.Bool) {
+					called.Store(true)
+				}),
+			),
+		},
+		{
+			name: "without context param",
+			annotation: fx.Annotate(
+				func() A { return nil },
+				fx.OnStart(func(called *atomic.Bool) error {
+					called.Store(true)
+					return nil
+				}),
+			),
+		},
+		{
+			name: "without context param or error return",
+			annotation: fx.Annotate(
+				func() A { return nil },
+				fx.OnStart(func(called *atomic.Bool) {
+					called.Store(true)
+				}),
+			),
+		},
+		{
+			name: "with context param and error return",
+			annotation: fx.Annotate(
+				func() A { return nil },
+				fx.OnStart(func(_ context.Context, called *atomic.Bool) error {
+					called.Store(true)
+					return nil
+				}),
+			),
+		},
+	}
+
+	for _, tt := range table {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			var (
+				called = atomic.NewBool(false)
+				opts   = fx.Options(
+					fx.Provide(tt.annotation),
+					fx.Supply(called),
+					fx.Invoke(func(A) {}),
+				)
+			)
+
+			fxtest.New(t, opts).RequireStart().RequireStop()
+			require.True(t, called.Load())
+		})
+	}
+}
diff --git a/error_example_test.go b/error_example_test.go
index d0fd4f2..c260bf4 100644
--- a/error_example_test.go
+++ b/error_example_test.go
@@ -39,7 +39,7 @@ func ExampleError() {
 			return fx.Error(errors.New("$PORT is not set"))
 		}
 		return fx.Provide(&http.Server{
-			Addr: fmt.Sprintf(":%s", port),
+			Addr: fmt.Sprintf("127.0.0.1:%s", port),
 		})
 	}
 
diff --git a/example_test.go b/example_test.go
index 8206cb7..38e458e 100644
--- a/example_test.go
+++ b/example_test.go
@@ -122,7 +122,7 @@ func NewMux(lc fx.Lifecycle, logger *log.Logger) *http.ServeMux {
 	// until all handlers are registered.
 	mux := http.NewServeMux()
 	server := &http.Server{
-		Addr:    ":8080",
+		Addr:    "127.0.0.1:8080",
 		Handler: mux,
 	}
 	// If NewMux is called, we know that another function is using the mux. In