Adjusting to @abhinav review

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

logger_test.go

@@ -217,7 +217,7 @@ func TestLoggerAlwaysFatals(t *testing.T) {
 		assert.True(t, stub.Exited, "Expected calls to logger.Fatal to terminate process.")
 		assert.Equal(t, 1, stub.Code, "Expected calls to logger.Fatal to terminate process with predictable retcode.")
 
-		logger = logger.WithOptions(OnFatal(zapcore.WriteThenPosixExitCode))
+		logger = logger.WithOptions(OnFatal(zapcore.WriteThenPosixExit))
 		err := errors.New("bar")
 		stub = exit.WithStub(func() { logger.Fatal("foo", Error(err)) })
 		assert.True(t, stub.Exited, "Expected calls to logger.Fatal to terminate process.")

options.go

@@ -134,8 +134,13 @@ func IncreaseLevel(lvl zapcore.LevelEnabler) Option {
 
 // OnFatal sets the action to take on fatal logs.
 func OnFatal(action zapcore.CheckWriteAction) Option {
+	return OnFatalHook(action)
+}
+
+// OnFatalHook allows to set a custom action to take on fatal logs.
+func OnFatalHook(hook zapcore.CheckWriteHook) Option {
 	return optionFunc(func(log *Logger) {
-		log.onFatal = action
+		log.onFatal = hook
 	})
 }
 

zapcore/entry.go

@@ -28,10 +28,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 (
@@ -153,33 +152,42 @@ type Entry struct {
 	Stack      string
 }
 
+type CheckWriteHook interface {
+	OnWrite(*CheckedEntry, []Field)
+}
+
 // CheckWriteAction indicates what action to take after a log entry is
 // processed. Actions are ordered in increasing severity.
-type CheckWriteAction func(*CheckedEntry, []Field)
+type CheckWriteAction uint8
 
-var (
+const (
 	// WriteThenNoop indicates that nothing special needs to be done. It's the
 	// default behavior.
-	WriteThenNoop CheckWriteAction
+	WriteThenNoop CheckWriteAction = iota
 	// WriteThenGoexit runs runtime.Goexit after Write.
-	WriteThenGoexit CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
-		runtime.Goexit()
-	}
+	WriteThenGoexit
 	// WriteThenPanic causes a panic after Write.
-	WriteThenPanic CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+	WriteThenPanic
+	// WriteThenFatal causes an os.Exit(1) after Write.
+	WriteThenFatal
+	// WriteThenPosixExit causes an os.Exit(code) after Write. The exit code is
+	// calculated deterministically from the Fatal message, or from error Field
+	// if given.
+	WriteThenPosixExit
+)
+
+func (a CheckWriteAction) OnWrite(ce *CheckedEntry, fields []Field) {
+	switch a {
+	case WriteThenGoexit:
+		runtime.Goexit()
+	case WriteThenPanic:
 		panic(ce.Message)
-	}
-	// WriteThenFatal causes a fatal os.Exit(1) after Write.
-	WriteThenFatal CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+	case WriteThenFatal:
 		exit.With(1)
-	}
-	// WriteThenPosixExitCode causes an os.Exit(code) after Write. The code is
-	// calculated deterministically from the message, or from attached error
-	// Field.
-	WriteThenPosixExitCode CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+	case WriteThenPosixExit:
 		exit.With(retcode(ce, fields))
 	}
-)
+}
 
 // CheckedEntry is an Entry together with a collection of Cores that have
 // already agreed to log it.
@@ -191,7 +199,7 @@ type CheckedEntry struct {
 	Entry
 	ErrorOutput WriteSyncer
 	dirty       bool // best-effort detection of pool misuse
-	should      CheckWriteAction
+	should      CheckWriteHook
 	cores       []Core
 }
 
@@ -243,7 +251,7 @@ func (ce *CheckedEntry) Write(fields ...Field) {
 
 	// Terminal operation
 	if ce.should != nil {
-		ce.should(ce, fields)
+		ce.should.OnWrite(ce, fields)
 	}
 }
 
@@ -259,10 +267,10 @@ func (ce *CheckedEntry) AddCore(ent Entry, core Core) *CheckedEntry {
 	return ce
 }
 
-// Should sets this CheckedEntry's CheckWriteAction, which controls whether a
+// Should sets this CheckedEntry's CheckWriteHook, 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.
-func (ce *CheckedEntry) Should(ent Entry, should CheckWriteAction) *CheckedEntry {
+func (ce *CheckedEntry) Should(ent Entry, should CheckWriteHook) *CheckedEntry {
 	if ce == nil {
 		ce = getCheckedEntry()
 		ce.Entry = ent

zapcore/entry_test.go

@@ -131,23 +131,29 @@ func TestCheckedEntryWrite(t *testing.T) {
 		assert.Equal(t, 1, stub.Code, "Expected to exit when WriteThenFatal is set.")
 	})
 
-	t.Run("WriteThenPosixExitCode", func(t *testing.T) {
+	t.Run("WriteThenPosixExit", func(t *testing.T) {
 		var ce *CheckedEntry
-		ce = ce.Should(Entry{Message: "foo"}, WriteThenPosixExitCode)
+		ce = ce.Should(Entry{Message: "foo"}, WriteThenPosixExit)
 		stub := exit.WithStub(func() {
 			ce.Write()
 		})
-		assert.True(t, stub.Exited, "Expected to exit when WriteThenPosixExitCode is set.")
-		assert.Equal(t, 38, stub.Code, "Expected to exit with specific code when WriteThenPosixExitCode is set.")
+		assert.True(t, stub.Exited, "Expected to exit when WriteThenPosixExit is set.")
+		assert.Equal(t, 38, stub.Code, "Expected to exit with specific code when WriteThenPosixExit is set.")
 	})
 
 	t.Run("Custom", func(t *testing.T) {
-		var called bool
 		var ce *CheckedEntry
-		ce = ce.Should(Entry{}, func(entry *CheckedEntry, fields []Field) {
-			called = true
-		})
+		hook := &customHook{}
+		ce = ce.Should(Entry{}, hook)
 		ce.Write()
-		assert.True(t, called, "Expected to call custom action after 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
+}