Add options to customize Fatal behaviour for better testability (uber…

…-go#861)

Fixes uber-go#846.

There's currently no easy way to test Fatal from outside of zap, as it
triggers an os.Exit(1). Add options to customize the behaviour (allowing
for panic, or Goexit), which can be used with `recover()` in tests.

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)

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

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

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

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