diff --git a/internal/exit/exit.go b/internal/exit/exit.go index dfc5b05fe..187fe5bfe 100644 --- a/internal/exit/exit.go +++ b/internal/exit/exit.go @@ -24,24 +24,32 @@ package exit import "os" -var real = func() { os.Exit(1) } +var _exit = os.Exit // Exit normally terminates the process by calling os.Exit(1). If the package // is stubbed, it instead records a call in the testing spy. +// Deprecated: use With() instead. func Exit() { - real() + With(1) +} + +// With terminates the process by calling os.Exit(code). If the package is +// stubbed, it instead records a call in the testing spy. +func With(code int) { + _exit(code) } // A StubbedExit is a testing fake for os.Exit. type StubbedExit struct { Exited bool - prev func() + Code int + prev func(code int) } // Stub substitutes a fake for the call to os.Exit(1). func Stub() *StubbedExit { - s := &StubbedExit{prev: real} - real = s.exit + s := &StubbedExit{prev: _exit} + _exit = s.exit return s } @@ -56,9 +64,10 @@ func WithStub(f func()) *StubbedExit { // Unstub restores the previous exit function. func (se *StubbedExit) Unstub() { - real = se.prev + _exit = se.prev } -func (se *StubbedExit) exit() { +func (se *StubbedExit) exit(code int) { se.Exited = true + se.Code = code } diff --git a/internal/exit/exit_test.go b/internal/exit/exit_test.go index 300cdc309..0f86a8844 100644 --- a/internal/exit/exit_test.go +++ b/internal/exit/exit_test.go @@ -18,25 +18,32 @@ // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN // THE SOFTWARE. -package exit +package exit_test import ( "testing" "github.com/stretchr/testify/assert" + "go.uber.org/zap/internal/exit" ) func TestStub(t *testing.T) { + type want struct { + exit bool + code int + } tests := []struct { f func() - want bool + want want }{ - {Exit, true}, - {func() {}, false}, + {func() { exit.With(42) }, want{exit: true, code: 42}}, + {exit.Exit, want{exit: true, code: 1}}, + {func() {}, want{}}, } for _, tt := range tests { - s := WithStub(tt.f) - assert.Equal(t, tt.want, s.Exited, "Stub captured unexpected exit value.") + s := exit.WithStub(tt.f) + assert.Equal(t, tt.want.exit, s.Exited, "Stub captured unexpected exit value.") + assert.Equal(t, tt.want.code, s.Code, "Stub captured unexpected exit value.") } } diff --git a/logger.go b/logger.go index 087c74222..e04de7cb3 100644 --- a/logger.go +++ b/logger.go @@ -42,7 +42,7 @@ type Logger struct { development bool addCaller bool - onFatal zapcore.CheckWriteAction // default is WriteThenFatal + onFatal zapcore.CheckWriteHook // default is WriteThenFatal name string errorOutput zapcore.WriteSyncer @@ -288,12 +288,12 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry { ce = ce.Should(ent, zapcore.WriteThenPanic) case zapcore.FatalLevel: onFatal := log.onFatal - // Noop is the default value for CheckWriteAction, and it leads to + // nil is the default value for CheckWriteAction, and it leads to // continued execution after a Fatal which is unexpected. - if onFatal == zapcore.WriteThenNoop { + if onFatal == nil { onFatal = zapcore.WriteThenFatal } - ce = ce.Should(ent, onFatal) + ce = ce.After(ent, onFatal) case zapcore.DPanicLevel: if log.development { ce = ce.Should(ent, zapcore.WriteThenPanic) diff --git a/logger_test.go b/logger_test.go index 6fbf12266..50f8c6599 100644 --- a/logger_test.go +++ b/logger_test.go @@ -579,6 +579,23 @@ func TestLoggerCustomOnFatal(t *testing.T) { } } +type customWriteHook struct { + called bool +} + +func (h *customWriteHook) OnWrite(_ *zapcore.CheckedEntry, _ []Field) { + h.called = true +} + +func TestLoggerWithFatalHook(t *testing.T) { + var h customWriteHook + withLogger(t, InfoLevel, opts(WithFatalHook(&h)), func(logger *Logger, logs *observer.ObservedLogs) { + logger.Fatal("great sadness") + assert.True(t, h.called) + assert.Equal(t, 1, logs.FilterLevelExact(FatalLevel).Len()) + }) +} + func TestNopLogger(t *testing.T) { logger := NewNop() diff --git a/options.go b/options.go index e9e66161f..d6c3449ca 100644 --- a/options.go +++ b/options.go @@ -133,9 +133,15 @@ func IncreaseLevel(lvl zapcore.LevelEnabler) Option { } // OnFatal sets the action to take on fatal logs. +// Deprecated: Use WithFatalHook instead. func OnFatal(action zapcore.CheckWriteAction) Option { + return WithFatalHook(action) +} + +// WithFatalHook sets a CheckWriteHook to run on fatal logs. +func WithFatalHook(hook zapcore.CheckWriteHook) Option { return optionFunc(func(log *Logger) { - log.onFatal = action + log.onFatal = hook }) } diff --git a/zapcore/entry.go b/zapcore/entry.go index 0885505b7..b035915f1 100644 --- a/zapcore/entry.go +++ b/zapcore/entry.go @@ -27,10 +27,9 @@ import ( "sync" "time" + "go.uber.org/multierr" "go.uber.org/zap/internal/bufferpool" "go.uber.org/zap/internal/exit" - - "go.uber.org/multierr" ) var ( @@ -152,6 +151,13 @@ type Entry struct { Stack string } +// CheckWriteHook allows to customize the action to take after a Fatal log entry +// is processed. +type CheckWriteHook interface { + // OnWrite gets invoked when an entry is written + OnWrite(*CheckedEntry, []Field) +} + // CheckWriteAction indicates what action to take after a log entry is // processed. Actions are ordered in increasing severity. type CheckWriteAction uint8 @@ -164,10 +170,25 @@ const ( WriteThenGoexit // WriteThenPanic causes a panic after Write. WriteThenPanic - // WriteThenFatal causes a fatal os.Exit after Write. + // WriteThenFatal causes an os.Exit(1) after Write. WriteThenFatal ) +// OnWrite implements the OnWrite method to keep CheckWriteAction compatible +// with the new CheckWriteHook interface which deprecates CheckWriteAction. +func (a CheckWriteAction) OnWrite(ce *CheckedEntry, _ []Field) { + switch a { + case WriteThenGoexit: + runtime.Goexit() + case WriteThenPanic: + panic(ce.Message) + case WriteThenFatal: + exit.Exit() + } +} + +var _ CheckWriteHook = CheckWriteAction(0) + // CheckedEntry is an Entry together with a collection of Cores that have // already agreed to log it. // @@ -178,7 +199,7 @@ type CheckedEntry struct { Entry ErrorOutput WriteSyncer dirty bool // best-effort detection of pool misuse - should CheckWriteAction + after CheckWriteHook cores []Core } @@ -186,7 +207,7 @@ func (ce *CheckedEntry) reset() { ce.Entry = Entry{} ce.ErrorOutput = nil ce.dirty = false - ce.should = WriteThenNoop + ce.after = nil for i := range ce.cores { // don't keep references to cores ce.cores[i] = nil @@ -224,17 +245,11 @@ func (ce *CheckedEntry) Write(fields ...Field) { ce.ErrorOutput.Sync() } - should, msg := ce.should, ce.Message - putCheckedEntry(ce) - - switch should { - case WriteThenPanic: - panic(msg) - case WriteThenFatal: - exit.Exit() - case WriteThenGoexit: - runtime.Goexit() + hook := ce.after + if hook != nil { + hook.OnWrite(ce, fields) } + putCheckedEntry(ce) } // AddCore adds a Core that has agreed to log this CheckedEntry. It's intended to be @@ -252,11 +267,19 @@ func (ce *CheckedEntry) AddCore(ent Entry, core Core) *CheckedEntry { // Should sets this CheckedEntry's CheckWriteAction, which controls whether a // Core will panic or fatal after writing this log entry. Like AddCore, it's // safe to call on nil CheckedEntry references. +// Deprecated: Use After(ent Entry, after CheckWriteHook) instead. func (ce *CheckedEntry) Should(ent Entry, should CheckWriteAction) *CheckedEntry { + return ce.After(ent, should) +} + +// After sets this CheckEntry's CheckWriteHook, which will be called after this +// log entry has been written. It's safe to call this on nil CheckedEntry +// references. +func (ce *CheckedEntry) After(ent Entry, hook CheckWriteHook) *CheckedEntry { if ce == nil { ce = getCheckedEntry() ce.Entry = ent } - ce.should = should + ce.after = hook return ce } diff --git a/zapcore/entry_test.go b/zapcore/entry_test.go index 4c2d67eae..d04cad9e7 100644 --- a/zapcore/entry_test.go +++ b/zapcore/entry_test.go @@ -64,7 +64,7 @@ func TestPutNilEntry(t *testing.T) { assert.NotNil(t, ce, "Expected only non-nil CheckedEntries in pool.") assert.False(t, ce.dirty, "Unexpected dirty bit set.") assert.Nil(t, ce.ErrorOutput, "Non-nil ErrorOutput.") - assert.Equal(t, WriteThenNoop, ce.should, "Unexpected terminal behavior.") + assert.Nil(t, ce.after, "Unexpected terminal behavior.") assert.Equal(t, 0, len(ce.cores), "Expected empty slice of cores.") assert.True(t, cap(ce.cores) > 0, "Expected pooled CheckedEntries to pre-allocate slice of Cores.") } @@ -128,5 +128,22 @@ func TestCheckedEntryWrite(t *testing.T) { ce.Write() }) assert.True(t, stub.Exited, "Expected to exit when WriteThenFatal is set.") + assert.Equal(t, 1, stub.Code, "Expected to exit when WriteThenFatal is set.") }) + + t.Run("After", func(t *testing.T) { + var ce *CheckedEntry + hook := &customHook{} + ce = ce.After(Entry{}, hook) + ce.Write() + assert.True(t, hook.called, "Expected to call custom action after Write.") + }) +} + +type customHook struct { + called bool +} + +func (c *customHook) OnWrite(_ *CheckedEntry, _ []Field) { + c.called = true }