From 40b848842e13554fe27270def4266ebe411352d5 Mon Sep 17 00:00:00 2001 From: luoboton Date: Sat, 11 Mar 2023 10:09:05 +0800 Subject: [PATCH] Add fx.Decorate (#833) This adds `fx.Decorate`, which lets you specify decorators to an fx app. A decorator can take in one or more dependencies that have already been `Provide`d to the app, and produce one or more values that will be used as replacements in the object graph. For example, suppose there is a simple app like this: ```go fx.New( fx.Provide(func() *Logger { return &Logger{Name: "logger"} }), fx.Invoke(func(l *Logger) { fmt.Println(l.Name) }), ) ``` Running this app will print "logger" on the console. Now let us suppose a decorator was provided: ```go fx.New( fx.Provide(...), // Provide same function as above fx.Decorate(func(l *Logger) *Logger { return &Logger{Name: "decorated " + l.Name} }), fx.Invoke(...), // Invoke same function as above ) ``` The decorator here will take in the provided Logger and replace it with another logger whose `Name` is `decorated logger`. The `Invoke`d function is then executed with this replacement value, so running this app will print "decorated logger" on the console. In terms of implementation, a decorator is represented by the target decorator function and the call stack it was provided from, similar to a provider. `module` contains a list of decorators that were specified within its scope. The dig dependency had to be updated to the latest master branch of Dig to ensure the fix for https://github.com/uber-go/dig/issues/316 is in. Following this PR, there are two additional pieces I will be adding: 1. An eventing system for fx.Decorate. 2. fx.Replace, which takes in a value instead of a function to replace a value in the object graph. This is similar to what fx.Supply is to fx.Provide. This PR along with the two PRs above should make the long-awaited feature of graph modifications in fx finally possible. --- Refs #653, #649, #825, uber-go/dig#230, GO-1203, GO-736 --- app.go | 7 + app_test.go | 5 + decorate.go | 137 ++++++++++++++++++ decorate_test.go | 356 +++++++++++++++++++++++++++++++++++++++++++++++ go.mod | 2 +- go.sum | 4 +- module.go | 30 +++- 7 files changed, 531 insertions(+), 10 deletions(-) create mode 100644 decorate.go create mode 100644 decorate_test.go diff --git a/app.go b/app.go index 0942d98..0598891 100644 --- a/app.go +++ b/app.go @@ -518,6 +518,13 @@ func New(opts ...Option) *App { } } + // Run decorators before executing any Invokes. + if err := app.root.decorate(); err != nil { + app.err = err + + return app + } + // This error might have come from the provide loop above. We've // already flushed to the custom logger, so we can return. if app.err != nil { diff --git a/app_test.go b/app_test.go index 128c535..ad8d095 100644 --- a/app_test.go +++ b/app_test.go @@ -1693,6 +1693,11 @@ func TestOptionString(t *testing.T) { give: Supply(Annotated{Target: bytes.NewReader(nil)}), want: "fx.Supply(*bytes.Reader)", }, + { + desc: "Decorate", + give: Decorate(bytes.NewBufferString), + want: "fx.Decorate(bytes.NewBufferString())", + }, } for _, tt := range tests { diff --git a/decorate.go b/decorate.go new file mode 100644 index 0000000..1e863d3 --- /dev/null +++ b/decorate.go @@ -0,0 +1,137 @@ +// 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 +// 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 + +import ( + "fmt" + "strings" + + "go.uber.org/dig" + "go.uber.org/fx/internal/fxreflect" +) + +// Decorate specifies one or more decorator functions to an Fx application. +// Decorator functions let users augment objects in the graph. They can take in +// zero or more dependencies that must be provided to the application with fx.Provide, +// and produce one or more values that can be used by other invoked values. +// +// An example decorator is the following function which accepts a value, augments that value, +// and returns the replacement value. +// +// fx.Decorate(func(log *zap.Logger) *zap.Logger { +// return log.Named("myapp") +// }) +// +// The following decorator accepts multiple dependencies from the graph, augments and returns +// one of them. +// +// fx.Decorate(func(log *zap.Logger, cfg *Config) *zap.Logger { +// return log.Named(cfg.Name) +// }) +// +// Similar to fx.Provide, functions passed to fx.Decorate may optionally return an error +// as their last result. If a decorator returns a non-nil error, it will halt application startup. +// +// All modifications in the object graph due to a decorator are scoped to the fx.Module it was +// specified from. Decorations specified in the top-level fx.New call apply across the application. +// +// Decorators can be annotated using fx.Annotate, but not with fx.Annotated. Refer to documentation +// on fx.Annotate() to learn how to use it for annotating functions. +// +// Decorators support fx.In and fx.Out structs, similar to fx.Provide and fx.Invoke. +// +// Decorators support value groups as well. For example, the following code shows a decorator +// which takes in a value group using fx.In struct, and returns another value group. +// +// type HandlerParam struct { +// fx.In +// +// Handlers []Handler `group:"server" +// } +// +// type HandlerResult struct { +// fx.Out +// +// Handlers []Handler `group:"server" +// } +// +// fx.New( +// // ... +// fx.Decorate(func(p HandlerParam) HandlerResult { +// // ... +// }), +// ) +func Decorate(decorators ...interface{}) Option { + return decorateOption{ + Targets: decorators, + Stack: fxreflect.CallerStack(1, 0), + } +} + +type decorateOption struct { + Targets []interface{} + Stack fxreflect.Stack +} + +func (o decorateOption) apply(mod *module) { + for _, target := range o.Targets { + mod.decorators = append(mod.decorators, decorator{ + Target: target, + Stack: o.Stack, + }) + } +} + +func (o decorateOption) String() string { + items := make([]string, len(o.Targets)) + for i, f := range o.Targets { + items[i] = fxreflect.FuncName(f) + } + return fmt.Sprintf("fx.Decorate(%s)", strings.Join(items, ", ")) +} + +// decorator is a single decorator used in Fx. +type decorator struct { + // Decorator provided to Fx. + Target interface{} + + // Stack trace of where this provide was made. + Stack fxreflect.Stack +} + +func runDecorator(c container, d decorator, opts ...dig.DecorateOption) (err error) { + decorator := d.Target + defer func() { + if err != nil { + err = fmt.Errorf("fx.Decorate(%v) from:\n%+vFailed: %v", decorator, d.Stack, err) + } + }() + + switch decorator := decorator.(type) { + case annotated: + if dcor, derr := decorator.Build(); derr == nil { + err = c.Decorate(dcor, opts...) + } + default: + err = c.Decorate(decorator, opts...) + } + return +} diff --git a/decorate_test.go b/decorate_test.go new file mode 100644 index 0000000..cf3f494 --- /dev/null +++ b/decorate_test.go @@ -0,0 +1,356 @@ +// 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 +// 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 ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/fx" + "go.uber.org/fx/fxtest" +) + +func TestDecorateSuccess(t *testing.T) { + type Logger struct { + Name string + } + + t.Run("objects provided by other modules are decorated", func(t *testing.T) { + redis := fx.Module("redis", + fx.Provide(func() *Logger { + return &Logger{Name: "redis"} + }), + ) + + testRedis := fx.Module("testRedis", + redis, + fx.Decorate(func() *Logger { + return &Logger{Name: "testRedis"} + }), + fx.Invoke(func(l *Logger) { + assert.Equal(t, "testRedis", l.Name) + }), + ) + + app := fxtest.New(t, + testRedis, + fx.Invoke(func(l *Logger) { + assert.Equal(t, "redis", l.Name) + }), + ) + defer app.RequireStart().RequireStop() + }) + + t.Run("objects in child modules are decorated.", func(t *testing.T) { + redis := fx.Module("redis", + fx.Decorate(func() *Logger { + return &Logger{Name: "redis"} + }), + fx.Invoke(func(l *Logger) { + assert.Equal(t, "redis", l.Name) + }), + ) + app := fxtest.New(t, + redis, + fx.Provide(func() *Logger { + assert.Fail(t, "should not run this") + return &Logger{Name: "root"} + }), + ) + defer app.RequireStart().RequireStop() + }) + + t.Run("root decoration applies to all modules", func(t *testing.T) { + redis := fx.Module("redis", + fx.Invoke(func(l *Logger) { + assert.Equal(t, "decorated logger", l.Name) + }), + ) + logger := fx.Module("logger", + fx.Provide(func() *Logger { + return &Logger{Name: "logger"} + }), + ) + app := fxtest.New(t, + redis, + logger, + fx.Decorate(func(l *Logger) *Logger { + return &Logger{Name: "decorated " + l.Name} + }), + ) + defer app.RequireStart().RequireStop() + }) + + t.Run("use Decorate with Annotate", func(t *testing.T) { + type Coffee struct { + Name string + Price int + } + + cafe := fx.Module("cafe", + fx.Provide(fx.Annotate(func() *Coffee { + return &Coffee{Name: "Americano", Price: 3} + }, fx.ResultTags(`group:"coffee"`))), + fx.Provide(fx.Annotate(func() *Coffee { + return &Coffee{Name: "Cappucino", Price: 4} + }, fx.ResultTags(`group:"coffee"`))), + fx.Provide(fx.Annotate(func() *Coffee { + return &Coffee{Name: "Cold Brew", Price: 4} + }, fx.ResultTags(`group:"coffee"`))), + ) + + takeout := fx.Module("takeout", + cafe, + fx.Decorate(fx.Annotate(func(coffee []*Coffee) []*Coffee { + var newC []*Coffee + for _, c := range coffee { + newC = append(newC, &Coffee{ + Name: c.Name, + Price: c.Price + 1, + }) + } + return newC + }, fx.ParamTags(`group:"coffee"`), fx.ResultTags(`group:"coffee"`))), + fx.Invoke(fx.Annotate(func(coffee []*Coffee) { + assert.Equal(t, 3, len(coffee)) + totalPrice := 0 + for _, c := range coffee { + totalPrice += c.Price + } + assert.Equal(t, 4+5+5, totalPrice) + }, fx.ParamTags(`group:"coffee"`))), + ) + + app := fxtest.New(t, + takeout, + ) + defer app.RequireStart().RequireStop() + }) + + t.Run("use Decorate with parameter/result struct", func(t *testing.T) { + type Logger struct { + Name string + } + type A struct { + fx.In + + Log *Logger + Version int `name:"versionNum"` + } + type B struct { + fx.Out + + Log *Logger + Version int `name:"versionNum"` + } + app := fxtest.New(t, + fx.Provide( + fx.Annotate(func() int { return 1 }, + fx.ResultTags(`name:"versionNum"`)), + func() *Logger { + return &Logger{Name: "logger"} + }, + ), + fx.Decorate(func(a A) B { + return B{ + Log: &Logger{Name: a.Log.Name + " decorated"}, + Version: a.Version + 1, + } + }), + fx.Invoke(fx.Annotate(func(l *Logger, ver int) { + assert.Equal(t, "logger decorated", l.Name) + assert.Equal(t, 2, ver) + }, fx.ParamTags(``, `name:"versionNum"`))), + ) + defer app.RequireStart().RequireStop() + }) + + t.Run("decorator with optional parameter", func(t *testing.T) { + type Config struct { + Name string + } + type Logger struct { + Name string + } + type DecoratorParam struct { + fx.In + + Cfg *Config `optional:"true"` + Log *Logger + } + + app := fxtest.New(t, + fx.Provide(func() *Logger { return &Logger{Name: "log"} }), + fx.Decorate(func(p DecoratorParam) *Logger { + if p.Cfg != nil { + return &Logger{Name: p.Cfg.Name} + } + return &Logger{Name: p.Log.Name} + }), + fx.Invoke(func(l *Logger) { + assert.Equal(t, l.Name, "log") + }), + ) + defer app.RequireStart().RequireStop() + }) +} + +func TestDecorateFailure(t *testing.T) { + t.Run("decorator returns an error", func(t *testing.T) { + type Logger struct { + Name string + } + + app := NewForTest(t, + fx.Provide(func() *Logger { + return &Logger{Name: "root"} + }), + fx.Decorate(func(l *Logger) (*Logger, error) { + return &Logger{Name: l.Name + "decorated"}, errors.New("minor sadness") + }), + fx.Invoke(func(l *Logger) { + assert.Fail(t, "this should not be executed") + }), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "minor sadness") + }) + + t.Run("decorator in a nested module returns an error", func(t *testing.T) { + type Logger struct { + Name string + } + + app := NewForTest(t, + fx.Provide(func() *Logger { + return &Logger{Name: "root"} + }), + fx.Module("child", + fx.Decorate(func(l *Logger) *Logger { + return &Logger{Name: l.Name + "decorated"} + }), + fx.Decorate(func(l *Logger) *Logger { + return &Logger{Name: l.Name + "decorated"} + }), + fx.Invoke(func(l *Logger) { + assert.Fail(t, "this should not be executed") + }), + ), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "*fx_test.Logger already decorated") + }) + + t.Run("decorating a type more than once in the same Module errors", func(t *testing.T) { + type Logger struct { + Name string + } + + app := NewForTest(t, + fx.Provide(func() *Logger { + return &Logger{Name: "root"} + }), + fx.Decorate(func(l *Logger) *Logger { + return &Logger{Name: "dec1 " + l.Name} + }), + fx.Decorate(func(l *Logger) *Logger { + return &Logger{Name: "dec2 " + l.Name} + }), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "*fx_test.Logger already decorated") + }) + + t.Run("annotated decorator returns an error", func(t *testing.T) { + type Logger struct { + Name string + } + + tag := `name:"decoratedLogger"` + app := NewForTest(t, + fx.Provide(fx.Annotate(func() *Logger { + return &Logger{Name: "root"} + }, fx.ResultTags(tag))), + fx.Decorate(fx.Annotate(func(l *Logger) (*Logger, error) { + return &Logger{Name: "dec1 " + l.Name}, errors.New("major sadness") + }, fx.ParamTags(tag), fx.ResultTags(tag))), + fx.Invoke(fx.Annotate(func(l *Logger) { + assert.Fail(t, "this should never run") + }, fx.ParamTags(tag))), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "major sadness") + }) + + t.Run("all decorator dependencies must be provided", func(t *testing.T) { + type Logger struct { + Name string + } + type Config struct { + Name string + } + + app := NewForTest(t, + fx.Provide(func() *Logger { + return &Logger{Name: "logger"} + }), + fx.Decorate(func(l *Logger, c *Config) *Logger { + return &Logger{Name: l.Name + c.Name} + }), + fx.Invoke(func(l *Logger) { + assert.Fail(t, "this should never run") + }), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "missing dependencies") + }) + + t.Run("decorate cannot provide a non-existent type", func(t *testing.T) { + type Logger struct { + Name string + } + + app := NewForTest(t, + fx.Decorate(func() *Logger { + return &Logger{Name: "decorator"} + }), + fx.Invoke(func(l *Logger) { + assert.Fail(t, "this should never run") + }), + ) + + err := app.Err() + require.Error(t, err) + assert.Contains(t, err.Error(), "missing dependencies") + }) +} diff --git a/go.mod b/go.mod index eac0e7f..d53179e 100644 --- a/go.mod +++ b/go.mod @@ -12,4 +12,4 @@ require ( golang.org/x/sys v0.0.0-20210903071746-97244b99971b ) -replace go.uber.org/dig => github.com/uber-go/dig v1.13.1-0.20220106194054-29dd17211ed4 +replace go.uber.org/dig => github.com/uber-go/dig v1.13.1-0.20220208182428-8193c7fedade diff --git a/go.sum b/go.sum index 613bd45..067f9b9 100644 --- a/go.sum +++ b/go.sum @@ -22,8 +22,8 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/uber-go/dig v1.13.1-0.20220106194054-29dd17211ed4 h1:lTskNvD4R8nSk2vH8AwwU7N1ZpFE7AzBv2CGTfCAxlw= -github.com/uber-go/dig v1.13.1-0.20220106194054-29dd17211ed4/go.mod h1:jHAn/z1Ld1luVVyGKOAIFYz/uBFqKjjEEdIqVAqfQ2o= +github.com/uber-go/dig v1.13.1-0.20220208182428-8193c7fedade h1:q2FxCCuMWy3p904qjF7NyaFj4EP/G5T7eACVKLvMgrU= +github.com/uber-go/dig v1.13.1-0.20220208182428-8193c7fedade/go.mod h1:jHAn/z1Ld1luVVyGKOAIFYz/uBFqKjjEEdIqVAqfQ2o= github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k= go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk= go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= diff --git a/module.go b/module.go index 47d99ec..b07bb5a 100644 --- a/module.go +++ b/module.go @@ -36,6 +36,7 @@ import ( type container interface { Invoke(interface{}, ...dig.InvokeOption) error Provide(interface{}, ...dig.ProvideOption) error + Decorate(interface{}, ...dig.DecorateOption) error } // Module is a named group of zero or more fx.Options. @@ -76,13 +77,14 @@ func (o moduleOption) apply(mod *module) { } type module struct { - parent *module - name string - scope *dig.Scope - provides []provide - invokes []invoke - modules []*module - app *App + parent *module + name string + scope *dig.Scope + provides []provide + invokes []invoke + decorators []decorator + modules []*module + app *App } // builds the Scopes using the App's Container. Note that this happens @@ -174,3 +176,17 @@ func (m *module) executeInvoke(i invoke) (err error) { }) return err } + +func (m *module) decorate() (err error) { + for _, decorator := range m.decorators { + if err := runDecorator(m.scope, decorator); err != nil { + return err + } + } + for _, m := range m.modules { + if err := m.decorate(); err != nil { + return err + } + } + return nil +}