Repository: uber-go/fx (go)

Issue:
fx.Decorate called on a dependency is not called when fx.WithLogger is involved with that dependency
First off, thank you for building, maintaining an publishing this amazing library. We use to great effect at our company.

We've decided to give fx.Decorate a shot to conditionally decorate a slice of functional options whenever some OTHER option is set.  Somethling like this: 
```Go
var Module = fx.Module("sink",
	fx.Decorate(func(os ...Option) []Option {
                fmt.Println("Decorated Called")
		// TODO: append some option to "os" in some cases
		return os
	}),
       // New is a constructor, see below. fx.Annotate is required to make fx provide the optional 
       // options, hard to debug but OK.
	fx.Provide(fx.Annotate(New)),
)

// a logging sink
type Sink struct{ opts Options }

// init the sink
func New(os ...Option) (s *Sink) {
	s = &Sink{}
	for _, o := range os {
		o(&s.opts)
	}
	return
}
```

In our case the module above provides a sink that our zap.Logger will write too. This is setup in another module dedicated to logging. Something like this:

```Go
// the dedicated logging module
var Module = fx.Module("log",
	fx.Provide(
		zap.New,
		zapcore.NewCore,
		zap.NewProductionEncoderConfig,
		zapcore.NewJSONEncoder,
		fx.Annotate(NewEnabler),
		fx.Annotate(WriteSyncers),
	),
)

func NewEnabler() zapcore.LevelEnabler { return zap.DebugLevel}

// Note: the sink need to be provided by the other module
func WriteSyncers(rs *sink2.Sink) zapcore.WriteSyncer {
	return zapcore.AddSync(rs)
}
```

This all seems to work well. We can bring the above together in a top level module like this:
```Go
var app = fx.Module("app",
	sink2.Module,
	log2.Module,
	
	// give fx something to provide for us
	fx.Invoke(func(_ *zap.Logger) {}),
)
```

We also have a option that includes the zap.Logger as an fx logger:
```Go
// FxLogger is a convenient option that configures fx to use the zap logger.
func FxLogger() fx.Option {
	return fx.WithLogger(func(l *zap.Logger) fxevent.Logger {
		return &fxevent.ZapLogger{Logger: l}
	})
}
```

So if we construct the final fx.App like this (without the fx.WithLogger): 
```Go
app := fx.New(app, fx.Invoke(func(_ *zap.Logger) {}))
```
The  decorator will be called correctly.

But if we now create the app WITH the fx logger: 
```Go
app := fx.New(FxLogger(), app, fx.Invoke(func(_ *zap.Logger) {}))
```
It does NOT work. I.e. The decorator doesn't get called.

It took me a while to track this down so I've setup a dedicated repo that reproduces the problem: https://github.com/advanderveer/fxzapreproduce

It should be possible to clone it and simply run `go test -v`. There are two test cases: one called the decorate, the other doesn't.





Pull request:
WithLogger: Support decoration
`fx.WithLogger` does not currently respect decorations made with `fx.Decorate`
or `fx.Replace`. The reasons for this are two fold:

1. App runs decorations *after* constructing the custom logger. This means that
   dig doesn't know about the decorations at this time.
2. App provides and invokes the custom logger on the `dig.Container`, while the
   decorations begin running at the root `dig.Scope` held by Fx. So the
   container that we run the custom logger off doesn't know about the decorator
   because it's one layer down.

This PR moves decoration to run before the custom logger is constructed,
and changes the custom logger instantiation code to operate on the root scope.

Resolves #860
Refs GO-1341

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4cc866f..733ab84 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Fixed
 - Upgrade Dig dependency to v1.14.1 to address a couple of issues with decorations. Refer to
   Dig v1.14.1 release notes for more details.
+- `fx.WithLogger` no longer ignores decorations and replacements of types that
+  it depends on.
 
 ## [1.17.1] - 2021-03-23
 ### Added
diff --git a/app.go b/app.go
index 0598891..99ac9ec 100644
--- a/app.go
+++ b/app.go
@@ -404,7 +404,7 @@ func (app *App) constructCustomLogger(buffer *logBuffer) (err error) {
 		})
 	}()
 
-	if err := app.container.Provide(p.Target); err != nil {
+	if err := app.root.scope.Provide(p.Target); err != nil {
 		return fmt.Errorf("fx.WithLogger(%v) from:\n%+vFailed: %v",
 			fname, p.Stack, err)
 	}
@@ -412,7 +412,7 @@ func (app *App) constructCustomLogger(buffer *logBuffer) (err error) {
 	// TODO: Use dig.FillProvideInfo to inspect the provided constructor
 	// and fail the application if its signature didn't match.
 
-	return app.container.Invoke(func(log fxevent.Logger) {
+	return app.root.scope.Invoke(func(log fxevent.Logger) {
 		app.log = log
 		buffer.Connect(log)
 	})
@@ -501,6 +501,10 @@ func New(opts ...Option) *App {
 	app.root.provide(provide{Target: app.shutdowner, Stack: frames})
 	app.root.provide(provide{Target: app.dotGraph, Stack: frames})
 
+	// Run decorators before executing any Invokes -- including the one
+	// inside constructCustomLogger.
+	app.err = multierr.Append(app.err, app.root.decorate())
+
 	// If you are thinking about returning here after provides: do not (just yet)!
 	// If a custom logger was being used, we're still buffering messages.
 	// We'll want to flush them to the logger.
@@ -518,13 +522,6 @@ 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 {


Test patch:
diff --git a/app_test.go b/app_test.go
index e38a735..c617894 100644
--- a/app_test.go
+++ b/app_test.go
@@ -357,7 +357,7 @@ func TestNewApp(t *testing.T) {
 		defer app.RequireStart().RequireStop()
 
 		require.Equal(t,
-			[]string{"Provided", "Provided", "Provided", "Provided", "LoggerInitialized", "Decorated", "Invoking", "Invoked", "Started"},
+			[]string{"Provided", "Provided", "Provided", "Provided", "Decorated", "LoggerInitialized", "Invoking", "Invoked", "Started"},
 			spy.EventTypes())
 	})
 
@@ -378,7 +378,7 @@ func TestNewApp(t *testing.T) {
 		defer app.RequireStart().RequireStop()
 
 		require.Equal(t,
-			[]string{"Provided", "Provided", "Provided", "Provided", "LoggerInitialized", "Decorated", "Decorated", "Started"},
+			[]string{"Provided", "Provided", "Provided", "Provided", "Decorated", "Decorated", "LoggerInitialized", "Started"},
 			spy.EventTypes())
 	})
 }
diff --git a/log_test.go b/log_test.go
index 94ed4de..c0d3008 100644
--- a/log_test.go
+++ b/log_test.go
@@ -26,6 +26,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/fx/fxevent"
 	"go.uber.org/fx/internal/fxlog"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest"
+	"go.uber.org/zap/zaptest/observer"
 )
 
 func TestLogBufferConnect(t *testing.T) {
@@ -57,3 +60,19 @@ func TestLogBufferLog(t *testing.T) {
 	lb.Connect(spy)
 	assert.Equal(t, fxlog.Events{event}, spy.Events())
 }
+
+func TestWithLoggerDecorate(t *testing.T) {
+	t.Parallel()
+
+	core, logs := observer.New(zap.DebugLevel)
+	New(
+		Supply(zaptest.NewLogger(t)), // provide a logger
+		Replace(zap.New(core)),       // and replace it
+		WithLogger(func(log *zap.Logger) fxevent.Logger {
+			return &fxevent.ZapLogger{Logger: log}
+		}),
+	)
+
+	assert.NotZero(t, logs.Len(),
+		"should post to replacement logger")
+}