Deterministic calc of POSIX return code for Fatal

Author: cardil

internal/exit/exit.go

@@ -24,24 +24,32 @@ package exit
 
 import "os"
 
-var real = func() { os.Exit(1) }
+var realFn = func(code int) { os.Exit(code) }
 
 // 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 ExitWith() instead.
 func Exit() {
-	real()
+	ExitWith(1)
+}
+
+// ExitWith terminates the process by calling os.Exit(code). If the package
+// is stubbed, it instead records a call in the testing spy.
+func ExitWith(code int) {
+	realFn(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: realFn}
+	realFn = 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
+	realFn = se.prev
 }
 
-func (se *StubbedExit) exit() {
+func (se *StubbedExit) exit(code int) {
 	se.Exited = true
+	se.Code = code
 }

internal/exit/exit_test.go

@@ -18,25 +18,34 @@
 // 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
+		f func()
+		want
 	}{
-		{Exit, true},
-		{func() {}, false},
+		{func() {
+			exit.ExitWith(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.")
 	}
 }

logger.go

@@ -290,7 +290,7 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry {
 		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 {
+		if onFatal == nil {
 			onFatal = zapcore.WriteThenFatal
 		}
 		ce = ce.Should(ent, onFatal)

logger_test.go

@@ -213,15 +213,23 @@ func TestLoggerAlwaysFatals(t *testing.T) {
 	// Users can disable writing out fatal-level logs, but calls to logger.Fatal()
 	// should still terminate the process.
 	withLogger(t, FatalLevel+1, nil, func(logger *Logger, logs *observer.ObservedLogs) {
-		stub := exit.WithStub(func() { logger.Fatal("") })
+		stub := exit.WithStub(func() { logger.Fatal("foo") })
 		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.WriteThenExitCode))
+		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.")
+		assert.Equal(t, 129, stub.Code, "Expected calls to logger.Fatal to terminate process with predictable retcode.")
 
 		stub = exit.WithStub(func() {
 			if ce := logger.Check(FatalLevel, ""); ce != nil {
 				ce.Write()
 			}
 		})
 		assert.True(t, stub.Exited, "Expected calls to logger.Check(FatalLevel, ...) to terminate process.")
+		assert.Equal(t, 1, stub.Code, "Expected calls to logger.Check(FatalLevel, ...) to terminate process with predictable retcode.")
 
 		assert.Equal(t, 0, logs.Len(), "Shouldn't write out logs when fatal-level logging is disabled.")
 	})

zapcore/entry.go

@@ -22,6 +22,7 @@ package zapcore
 
 import (
 	"fmt"
+	"hash/crc32"
 	"runtime"
 	"strings"
 	"sync"
@@ -154,18 +155,30 @@ type Entry struct {
 
 // CheckWriteAction indicates what action to take after a log entry is
 // processed. Actions are ordered in increasing severity.
-type CheckWriteAction uint8
+type CheckWriteAction func(*CheckedEntry, []Field)
 
-const (
+var (
 	// WriteThenNoop indicates that nothing special needs to be done. It's the
 	// default behavior.
-	WriteThenNoop CheckWriteAction = iota
+	WriteThenNoop CheckWriteAction
 	// WriteThenGoexit runs runtime.Goexit after Write.
-	WriteThenGoexit
+	WriteThenGoexit CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+		runtime.Goexit()
+	}
 	// WriteThenPanic causes a panic after Write.
-	WriteThenPanic
-	// WriteThenFatal causes a fatal os.Exit after Write.
-	WriteThenFatal
+	WriteThenPanic CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+		panic(ce.Message)
+	}
+	// WriteThenFatal causes a fatal os.Exit(1) after Write.
+	WriteThenFatal CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+		exit.ExitWith(1)
+	}
+	// WriteThenExitCode causes an os.Exit(code) after Write. The code is
+	// calculated deterministically from the message, or from attached error
+	// Field.
+	WriteThenExitCode CheckWriteAction = func(ce *CheckedEntry, fields []Field) {
+		exit.ExitWith(retcode(ce, fields))
+	}
 )
 
 // CheckedEntry is an Entry together with a collection of Cores that have
@@ -186,7 +199,9 @@ func (ce *CheckedEntry) reset() {
 	ce.Entry = Entry{}
 	ce.ErrorOutput = nil
 	ce.dirty = false
-	ce.should = WriteThenNoop
+	if ce.should != nil {
+		ce.should = WriteThenNoop
+	}
 	for i := range ce.cores {
 		// don't keep references to cores
 		ce.cores[i] = nil
@@ -224,16 +239,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()
+	// Terminal operation
+	if ce.should != nil {
+		ce.should(ce, fields)
 	}
 }
 
@@ -260,3 +270,14 @@ func (ce *CheckedEntry) Should(ent Entry, should CheckWriteAction) *CheckedEntry
 	ce.should = should
 	return ce
 }
+
+func retcode(ce *CheckedEntry, fields []Field) int {
+	msg := ce.Message
+	for _, field := range fields {
+		if field.Type == ErrorType {
+			msg = field.Interface.(error).Error()
+			break
+		}
+	}
+	return int(crc32.ChecksumIEEE([]byte(msg)))%254 + 1
+}

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.should, "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.")
 		}
@@ -124,9 +124,22 @@ func TestCheckedEntryWrite(t *testing.T) {
 	t.Run("WriteThenFatal", func(t *testing.T) {
 		var ce *CheckedEntry
 		ce = ce.Should(Entry{}, WriteThenFatal)
+		ce.Message = t.Name()
 		stub := exit.WithStub(func() {
 			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("WriteThenExitCode", func(t *testing.T) {
+		var ce *CheckedEntry
+		ce = ce.Should(Entry{}, WriteThenExitCode)
+		ce.Message = t.Name()
+		stub := exit.WithStub(func() {
+			ce.Write()
+		})
+		assert.True(t, stub.Exited, "Expected to exit when WriteThenExitCode is set.")
+		assert.Equal(t, 243, stub.Code, "Expected to exit with specific code when WriteThenExitCode is set.")
 	})
 }