Repository: uber-go/fx (go)

Issue:
Allow specifying fx.Module-specific custom loggers with WithLogger
Currently, passing `fx.WithLogger` to an `fx.Module` throws an error like this:

```
[Fx] ERROR              Failed to start: fx.WithLogger Option should be passed to top-level App, not to fx.Module
```

This isn't ideal because a library author that exports their app as an `fx.Module` cannot specify a custom logger whose scope applies to just their `Module`.

This tracks extending the current implementation of `WithLogger` option so that it can be passed directly to an `fx.Module`, such that when it is, it only affects the scope of that `Module`. 

Internal Ref: GO-1488

Pull request description:
Allow fx.Module-specific custom loggers
Currently passing fx.WithLogger option to fx.Module does not work because fx currently does not support specifying module-level custom loggers.

Now fx.WithLogger can be passed to fx.Modules for specifying custom loggers for fx.Modules. Also added some test cases to test fx.Module-specific loggers.

Fixes #894

Patch:
diff --git a/app.go b/app.go
index 8d0132c..dd29c3b 100644
--- a/app.go
+++ b/app.go
@@ -193,9 +193,10 @@ type withLoggerOption struct {
 
 func (l withLoggerOption) apply(m *module) {
 	if m.parent != nil {
-		// loggers shouldn't differ based on Module.
-		m.app.err = fmt.Errorf("fx.WithLogger Option should be passed to top-level App, " +
-			"not to fx.Module")
+		m.logConstructor = &provide{
+			Target: l.constructor,
+			Stack:  l.Stack,
+		}
 	} else {
 		m.app.logConstructor = &provide{
 			Target: l.constructor,
@@ -521,6 +522,10 @@ func New(opts ...Option) *App {
 		}
 	}
 
+	for _, m := range app.modules {
+		m.constructAllCustomLoggers()
+	}
+
 	// 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/module.go b/module.go
index 4a769d4..9839429 100644
--- a/module.go
+++ b/module.go
@@ -26,6 +26,7 @@ import (
 	"go.uber.org/dig"
 	"go.uber.org/fx/fxevent"
 	"go.uber.org/fx/internal/fxreflect"
+	"go.uber.org/multierr"
 )
 
 // A container represents a set of constructors to provide
@@ -79,14 +80,16 @@ func (o moduleOption) apply(mod *module) {
 }
 
 type module struct {
-	parent     *module
-	name       string
-	scope      scope
-	provides   []provide
-	invokes    []invoke
-	decorators []decorator
-	modules    []*module
-	app        *App
+	parent         *module
+	name           string
+	scope          scope
+	provides       []provide
+	invokes        []invoke
+	decorators     []decorator
+	modules        []*module
+	app            *App
+	log            fxevent.Logger
+	logConstructor *provide
 }
 
 // scope is a private wrapper interface for dig.Container and dig.Scope.
@@ -107,9 +110,20 @@ type scope interface {
 func (m *module) build(app *App, root *dig.Container) {
 	if m.parent == nil {
 		m.scope = root
+		m.log = app.log
 	} else {
 		parentScope := m.parent.scope
 		m.scope = parentScope.Scope(m.name)
+		// use parent module's logger by default
+		m.log = m.parent.log
+	}
+
+	if m.logConstructor != nil {
+		// Since user supplied a custom logger, use a buffered logger
+		// to hold all messages until user supplied logger is
+		// instantiated. Then we flush those messages after fully
+		// constructing the custom logger.
+		m.log = new(logBuffer)
 	}
 
 	for _, mod := range m.modules {
@@ -158,7 +172,49 @@ func (m *module) provide(p provide) {
 			Err:             m.app.err,
 		}
 	}
-	m.app.log.LogEvent(ev)
+	m.log.LogEvent(ev)
+}
+
+// Constructs custom loggers for all modules in the tree
+func (m *module) constructAllCustomLoggers() {
+	if m.logConstructor != nil {
+		if buffer, ok := m.log.(*logBuffer); ok {
+			// default to parent's logger if custom logger constructor fails
+			if err := m.constructCustomLogger(buffer); err != nil {
+				m.app.err = multierr.Append(m.app.err, err)
+				m.log = m.parent.log
+				buffer.Connect(m.log)
+			}
+		}
+	}
+
+	for _, mod := range m.modules {
+		mod.constructAllCustomLoggers()
+	}
+}
+
+// Mirroring the behavior of app.constructCustomLogger
+func (m *module) constructCustomLogger(buffer *logBuffer) (err error) {
+	p := m.logConstructor
+	fname := fxreflect.FuncName(p.Target)
+	defer func() {
+		m.log.LogEvent(&fxevent.LoggerInitialized{
+			Err:             err,
+			ConstructorName: fname,
+		})
+	}()
+
+	// TODO: Use dig.FillProvideInfo to inspect the provided constructor
+	// and fail the application if its signature didn't match.
+	if err := m.scope.Provide(p.Target); err != nil {
+		return fmt.Errorf("fx.WithLogger(%v) from:\n%+v\nin Module: %q\nFailed: %w",
+			fname, p.Stack, m.name, err)
+	}
+
+	return m.scope.Invoke(func(log fxevent.Logger) {
+		m.log = log
+		buffer.Connect(log)
+	})
 }
 
 func (m *module) executeInvokes() error {
@@ -179,12 +235,12 @@ func (m *module) executeInvokes() error {
 
 func (m *module) executeInvoke(i invoke) (err error) {
 	fnName := fxreflect.FuncName(i.Target)
-	m.app.log.LogEvent(&fxevent.Invoking{
+	m.log.LogEvent(&fxevent.Invoking{
 		FunctionName: fnName,
 		ModuleName:   m.name,
 	})
 	err = runInvoke(m.scope, i)
-	m.app.log.LogEvent(&fxevent.Invoked{
+	m.log.LogEvent(&fxevent.Invoked{
 		FunctionName: fnName,
 		ModuleName:   m.name,
 		Err:          err,
@@ -203,14 +259,14 @@ func (m *module) decorate() (err error) {
 		}
 
 		if decorator.IsReplace {
-			m.app.log.LogEvent(&fxevent.Replaced{
+			m.log.LogEvent(&fxevent.Replaced{
 				ModuleName:      m.name,
 				OutputTypeNames: outputNames,
 				Err:             err,
 			})
 		} else {
 
-			m.app.log.LogEvent(&fxevent.Decorated{
+			m.log.LogEvent(&fxevent.Decorated{
 				DecoratorName:   fxreflect.FuncName(decorator.Target),
 				ModuleName:      m.name,
 				OutputTypeNames: outputNames,


Test patch:
diff --git a/module_test.go b/module_test.go
index 13be8be..469dc8f 100644
--- a/module_test.go
+++ b/module_test.go
@@ -21,6 +21,7 @@
 package fx_test
 
 import (
+	"bytes"
 	"errors"
 	"testing"
 	"time"
@@ -31,6 +32,7 @@ import (
 	"go.uber.org/fx/fxevent"
 	"go.uber.org/fx/fxtest"
 	"go.uber.org/fx/internal/fxlog"
+	"go.uber.org/zap"
 )
 
 func TestModuleSuccess(t *testing.T) {
@@ -40,6 +42,10 @@ func TestModuleSuccess(t *testing.T) {
 		Name string
 	}
 
+	type Foo struct {
+		Name string
+	}
+
 	t.Run("provide a dependency from a submodule", func(t *testing.T) {
 		t.Parallel()
 
@@ -241,6 +247,191 @@ func TestModuleSuccess(t *testing.T) {
 		)
 		require.NoError(t, app.Err())
 	})
+
+	t.Run("custom logger for module", func(t *testing.T) {
+		t.Parallel()
+
+		tests := []struct {
+			desc           string
+			giveWithLogger fx.Option
+			wantEvents     []string
+		}{
+			{
+				desc:           "custom logger for module",
+				giveWithLogger: fx.NopLogger,
+				wantEvents: []string{
+					"Supplied", "Provided", "Provided", "Provided",
+					"LoggerInitialized", "Invoking", "Invoked",
+				},
+			},
+			{
+				desc:           "Not using a custom logger for module defaults to app logger",
+				giveWithLogger: fx.Options(),
+				wantEvents: []string{
+					"Supplied", "Provided", "Provided", "Provided", "Provided",
+					"LoggerInitialized", "Invoking", "Invoked", "Invoking", "Invoked",
+				},
+			},
+		}
+
+		for _, tt := range tests {
+			tt := tt
+			t.Run(tt.desc, func(t *testing.T) {
+				t.Parallel()
+				var spy fxlog.Spy
+
+				redis := fx.Module("redis",
+					fx.Provide(func() *Foo {
+						return &Foo{Name: "redis"}
+					}),
+					fx.Invoke(func(r *Foo) {
+						assert.Equal(t, "redis", r.Name)
+					}),
+					tt.giveWithLogger,
+				)
+
+				app := fxtest.New(t,
+					redis,
+					fx.Supply(&spy),
+					fx.WithLogger(func(spy *fxlog.Spy) fxevent.Logger {
+						return spy
+					}),
+					fx.Invoke(func(r *Foo) {
+						assert.Equal(t, "redis", r.Name)
+					}),
+				)
+
+				// events from module with a custom logger not logged in app logger
+				assert.Equal(t, tt.wantEvents, spy.EventTypes())
+
+				spy.Reset()
+				app.RequireStart().RequireStop()
+
+				require.NoError(t, app.Err())
+
+				assert.Equal(t, []string{"Started", "Stopped"}, spy.EventTypes())
+			})
+		}
+	})
+
+	type NamedSpy struct {
+		fxlog.Spy
+		name string
+	}
+
+	t.Run("decorator on module logger does not affect app logger", func(t *testing.T) {
+		t.Parallel()
+
+		appSpy := NamedSpy{name: "app"}
+		moduleSpy := NamedSpy{name: "redis"}
+
+		redis := fx.Module("redis",
+			fx.Provide(func() *Foo {
+				return &Foo{Name: "redis"}
+			}),
+			fx.Supply(&appSpy),
+			fx.Replace(&moduleSpy),
+			fx.WithLogger(func(spy *NamedSpy) fxevent.Logger {
+				assert.Equal(t, "redis", spy.name)
+				return spy
+			}),
+			fx.Invoke(func(r *Foo) {
+				assert.Equal(t, "redis", r.Name)
+			}),
+		)
+
+		app := fxtest.New(t,
+			redis,
+			fx.WithLogger(func(spy *NamedSpy) fxevent.Logger {
+				assert.Equal(t, "app", spy.name)
+				return spy
+			}),
+			fx.Invoke(func(r *Foo) {
+				assert.Equal(t, "redis", r.Name)
+			}),
+		)
+
+		assert.Equal(t, []string{
+			"Provided", "Supplied", "Replaced",
+			"LoggerInitialized", "Invoking", "Invoked",
+		}, moduleSpy.EventTypes())
+
+		assert.Equal(t, []string{
+			"Provided", "Provided", "Provided",
+			"LoggerInitialized", "Invoking", "Invoked",
+		}, appSpy.EventTypes())
+
+		appSpy.Reset()
+		moduleSpy.Reset()
+
+		app.RequireStart().RequireStop()
+
+		require.NoError(t, app.Err())
+
+		assert.Equal(t, []string{"Started", "Stopped"}, appSpy.EventTypes())
+		assert.Empty(t, moduleSpy.EventTypes())
+	})
+
+	t.Run("module uses parent module's logger to log events", func(t *testing.T) {
+		t.Parallel()
+
+		appSpy := NamedSpy{name: "app"}
+		childSpy := NamedSpy{name: "child"}
+
+		grandchild := fx.Module("grandchild",
+			fx.Provide(func() *Foo {
+				return &Foo{Name: "grandchild"}
+			}),
+			fx.Invoke(func(r *Foo) {
+				assert.Equal(t, "grandchild", r.Name)
+			}),
+		)
+
+		child := fx.Module("child",
+			grandchild,
+			fx.Supply(&appSpy),
+			fx.Replace(&childSpy),
+			fx.WithLogger(func(spy *NamedSpy) fxevent.Logger {
+				assert.Equal(t, "child", spy.name)
+				return spy
+			}),
+			fx.Invoke(func(r *Foo) {
+				assert.Equal(t, "grandchild", r.Name)
+			}),
+		)
+
+		app := fxtest.New(t,
+			child,
+			fx.WithLogger(func(spy *NamedSpy) fxevent.Logger {
+				assert.Equal(t, "app", spy.name)
+				return spy
+			}),
+			fx.Invoke(func(r *Foo) {
+				assert.Equal(t, "grandchild", r.Name)
+			}),
+		)
+
+		assert.Equal(t, []string{
+			"Supplied", "Provided", "Replaced", "LoggerInitialized",
+			//Invoke logged twice, once from child and another from grandchild
+			"Invoking", "Invoked", "Invoking", "Invoked",
+		}, childSpy.EventTypes(), "events from grandchild also logged in child logger")
+
+		assert.Equal(t, []string{
+			"Provided", "Provided", "Provided",
+			"LoggerInitialized", "Invoking", "Invoked",
+		}, appSpy.EventTypes(), "events from modules do not appear in app logger")
+
+		appSpy.Reset()
+		childSpy.Reset()
+
+		app.RequireStart().RequireStop()
+
+		require.NoError(t, app.Err())
+
+		assert.Equal(t, []string{"Started", "Stopped"}, appSpy.EventTypes())
+		assert.Empty(t, childSpy.EventTypes())
+	})
 }
 
 func TestModuleFailures(t *testing.T) {
@@ -400,10 +591,6 @@ func TestModuleFailures(t *testing.T) {
 				desc: "StopTimeout Option",
 				opt:  fx.StopTimeout(time.Second),
 			},
-			{
-				desc: "WithLogger Option",
-				opt:  fx.WithLogger(func() fxevent.Logger { return new(fxlog.Spy) }),
-			},
 		}
 
 		for _, tt := range tests {
@@ -420,4 +607,115 @@ func TestModuleFailures(t *testing.T) {
 			})
 		}
 	})
+
+	t.Run("invalid WithLogger in Module", func(t *testing.T) {
+		t.Parallel()
+
+		var spy fxlog.Spy
+
+		spyAsLogger := fx.Options(
+			fx.Supply(&spy),
+			fx.WithLogger(func(spy *fxlog.Spy) fxevent.Logger {
+				return spy
+			}),
+		)
+
+		defaultModuleOpts := fx.Options(
+			fx.Provide(func() string {
+				return "redis"
+			}),
+			fx.Invoke(func(s string) {
+				assert.Fail(t, "this should never run")
+			}),
+		)
+
+		tests := []struct {
+			desc            string
+			giveModuleOpts  fx.Option
+			giveAppOpts     fx.Option
+			wantErrContains []string
+			wantEvents      []string
+		}{
+			{
+				desc:           "error in Provide shows logs in module",
+				giveModuleOpts: fx.Options(spyAsLogger, fx.Provide(&bytes.Buffer{})), // not passing in a constructor
+				giveAppOpts:    fx.Options(),
+				wantErrContains: []string{
+					"must provide constructor function, got  (type *bytes.Buffer)",
+				},
+				wantEvents: []string{
+					"Supplied", "Provided", "LoggerInitialized",
+				},
+			},
+			{
+				desc: "logger in module failed to build",
+				giveModuleOpts: fx.WithLogger(func() (fxevent.Logger, error) {
+					return nil, errors.New("error building logger")
+				}),
+				giveAppOpts:     spyAsLogger,
+				wantErrContains: []string{"error building logger"},
+				wantEvents: []string{
+					"Supplied", "Provided", "Provided", "Provided",
+					"LoggerInitialized", "Provided", "LoggerInitialized",
+				},
+			},
+			{
+				desc: "logger dependency in module failed to build",
+				giveModuleOpts: fx.Options(
+					fx.Provide(func() (*zap.Logger, error) {
+						return nil, errors.New("error building logger dependency")
+					}),
+					fx.WithLogger(func(log *zap.Logger) fxevent.Logger {
+						t.Errorf("WithLogger must not be called")
+						panic("must not be called")
+					}),
+				),
+				giveAppOpts:     spyAsLogger,
+				wantErrContains: []string{"error building logger dependency"},
+				wantEvents: []string{
+					"Supplied", "Provided", "Provided", "Provided",
+					"LoggerInitialized", "Provided", "Provided", "LoggerInitialized",
+				},
+			},
+			{
+				desc:           "Invalid input for WithLogger",
+				giveModuleOpts: fx.WithLogger(&fxlog.Spy{}), // not passing in a constructor for WithLogger
+				giveAppOpts:    spyAsLogger,
+				wantErrContains: []string{
+					"fx.WithLogger", "from:", "Failed",
+				},
+				wantEvents: []string{
+					"Supplied", "Provided", "Provided", "Provided",
+					"LoggerInitialized", "Provided", "LoggerInitialized",
+				},
+			},
+		}
+		for _, tt := range tests {
+			spy.Reset()
+			t.Run(tt.desc, func(t *testing.T) {
+				redis := fx.Module("redis",
+					tt.giveModuleOpts,
+					defaultModuleOpts,
+				)
+
+				app := fx.New(
+					tt.giveAppOpts,
+					redis,
+					fx.Invoke(func(s string) {
+						assert.Fail(t, "this should never run")
+					}),
+				)
+				err := app.Err()
+				require.Error(t, err)
+				for _, contains := range tt.wantErrContains {
+					assert.Contains(t, err.Error(), contains)
+				}
+
+				assert.Equal(t,
+					tt.wantEvents,
+					spy.EventTypes(),
+				)
+			})
+		}
+	})
 }