WithLogger: Support decoration (#872)

`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

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

app.go

@@ -404,15 +404,15 @@ 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)
 	}
 
 	// 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 {

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())
 	})
 }

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")
+}