diff --git a/logger.go b/logger.go index 4726456d7..ea484aed1 100644 --- a/logger.go +++ b/logger.go @@ -49,6 +49,7 @@ type Logger struct { addStack zapcore.LevelEnabler callerSkip int + onFatal zapcore.CheckWriteAction // default is WriteThenFatal } // New constructs a new Logger from the provided zapcore.Core and Options. If @@ -280,7 +281,13 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry { case zapcore.PanicLevel: ce = ce.Should(ent, zapcore.WriteThenPanic) case zapcore.FatalLevel: - ce = ce.Should(ent, zapcore.WriteThenFatal) + onFatal := log.onFatal + // Noop is the default value for CheckWriteAction, and it leads to + // continued execution after a Fatal which is unexpected. + if onFatal == zapcore.WriteThenNoop { + onFatal = zapcore.WriteThenFatal + } + ce = ce.Should(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 6992a9a19..0fb0600e7 100644 --- a/logger_test.go +++ b/logger_test.go @@ -535,6 +535,51 @@ func TestLoggerConcurrent(t *testing.T) { }) } +func TestLoggerCustomOnFatal(t *testing.T) { + tests := []struct { + msg string + onFatal zapcore.CheckWriteAction + recoverValue interface{} + }{ + { + msg: "panic", + onFatal: zapcore.WriteThenPanic, + recoverValue: "fatal", + }, + { + msg: "goexit", + onFatal: zapcore.WriteThenGoexit, + recoverValue: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.msg, func(t *testing.T) { + withLogger(t, InfoLevel, opts(OnFatal(tt.onFatal)), func(logger *Logger, logs *observer.ObservedLogs) { + + var finished bool + recovered := make(chan interface{}) + go func() { + defer func() { + recovered <- recover() + }() + + logger.Fatal("fatal") + finished = true + }() + + assert.Equal(t, tt.recoverValue, <-recovered, "unexpected value from recover()") + assert.False(t, finished, "expect goroutine to not finish after Fatal") + + assert.Equal(t, []observer.LoggedEntry{{ + Entry: zapcore.Entry{Level: FatalLevel, Message: "fatal"}, + Context: []Field{}, + }}, logs.AllUntimed(), "unexpected logs") + }) + }) + } +} + func infoLog(logger *Logger, msg string, fields ...Field) { logger.Info(msg, fields...) } diff --git a/options.go b/options.go index 2bf2cb222..0135c2092 100644 --- a/options.go +++ b/options.go @@ -131,3 +131,10 @@ func IncreaseLevel(lvl zapcore.LevelEnabler) Option { } }) } + +// OnFatal sets the action to take on fatal logs. +func OnFatal(action zapcore.CheckWriteAction) Option { + return optionFunc(func(log *Logger) { + log.onFatal = action + }) +} diff --git a/zapcore/entry.go b/zapcore/entry.go index 717dc0c4f..4aa8b4f90 100644 --- a/zapcore/entry.go +++ b/zapcore/entry.go @@ -22,6 +22,7 @@ package zapcore import ( "fmt" + "runtime" "strings" "sync" "time" @@ -159,6 +160,8 @@ const ( // WriteThenNoop indicates that nothing special needs to be done. It's the // default behavior. WriteThenNoop CheckWriteAction = iota + // WriteThenGoexit runs runtime.Goexit after Write. + WriteThenGoexit // WriteThenPanic causes a panic after Write. WriteThenPanic // WriteThenFatal causes a fatal os.Exit after Write. @@ -231,6 +234,8 @@ func (ce *CheckedEntry) Write(fields ...Field) { panic(msg) case WriteThenFatal: exit.Exit() + case WriteThenGoexit: + runtime.Goexit() } } diff --git a/zapcore/entry_test.go b/zapcore/entry_test.go index 569c4e1e0..4c2d67eae 100644 --- a/zapcore/entry_test.go +++ b/zapcore/entry_test.go @@ -29,6 +29,22 @@ import ( "github.com/stretchr/testify/assert" ) +func assertGoexit(t *testing.T, f func()) { + var finished bool + recovered := make(chan interface{}) + go func() { + defer func() { + recovered <- recover() + }() + + f() + finished = true + }() + + assert.Nil(t, <-recovered, "Goexit should cause recover to return nil") + assert.False(t, finished, "Goroutine should not finish after Goexit") +} + func TestPutNilEntry(t *testing.T) { // Pooling nil entries defeats the purpose. var wg sync.WaitGroup @@ -88,20 +104,29 @@ func TestEntryCaller(t *testing.T) { } func TestCheckedEntryWrite(t *testing.T) { - // Nil checked entries are safe. - var ce *CheckedEntry - assert.NotPanics(t, func() { ce.Write() }, "Unexpected panic writing nil CheckedEntry.") - - // WriteThenPanic - ce = ce.Should(Entry{}, WriteThenPanic) - assert.Panics(t, func() { ce.Write() }, "Expected to panic when WriteThenPanic is set.") - ce.reset() - - // WriteThenFatal - ce = ce.Should(Entry{}, WriteThenFatal) - stub := exit.WithStub(func() { - ce.Write() + t.Run("nil is safe", func(t *testing.T) { + var ce *CheckedEntry + assert.NotPanics(t, func() { ce.Write() }, "Unexpected panic writing nil CheckedEntry.") + }) + + t.Run("WriteThenPanic", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenPanic) + assert.Panics(t, func() { ce.Write() }, "Expected to panic when WriteThenPanic is set.") + }) + + t.Run("WriteThenGoexit", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenGoexit) + assertGoexit(t, func() { ce.Write() }) + }) + + t.Run("WriteThenFatal", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenFatal) + stub := exit.WithStub(func() { + ce.Write() + }) + assert.True(t, stub.Exited, "Expected to exit when WriteThenFatal is set.") }) - assert.True(t, stub.Exited, "Expected to exit when WriteThenFatal is set.") - ce.reset() }