Add WithPanicHook logger option for panic log tests (#1416)

Add a `WithPanicHook` logger option that allows callers to specify custom behavior besides panicking on Panic/DPanic logs. This is similar to what we already have with the WithFatal hook implemented in #861. This will make it possible to unit test Panic log cases like the one we had with our periodic runner which was impossible because of unrecoverable panics in another go routine. Added unit tests and they pass. ``` $ make test ? go.uber.org/zap/internal [no test files] ? go.uber.org/zap/internal/bufferpool [no test files] ok go.uber.org/zap (cached) ? go.uber.org/zap/internal/readme [no test files] ok go.uber.org/zap/buffer (cached) ok go.uber.org/zap/internal/color (cached) ok go.uber.org/zap/internal/exit (cached) ok go.uber.org/zap/internal/pool (cached) ok go.uber.org/zap/internal/stacktrace (cached) ok go.uber.org/zap/internal/ztest (cached) ok go.uber.org/zap/zapcore (cached) ok go.uber.org/zap/zapgrpc (cached) ok go.uber.org/zap/zapio (cached) ok go.uber.org/zap/zaptest (cached) ok go.uber.org/zap/zaptest/observer (cached) ok go.uber.org/zap/exp/zapfield (cached) ok go.uber.org/zap/exp/zapslog (cached) ok go.uber.org/zap/benchmarks (cached) [no tests to run] ok go.uber.org/zap/zapgrpc/internal/test (cached) ``` Closes #1415
uber-go · Feb 20, 2024 · e5a56ee · e5a56ee
1 parent 0e2aa4e
commit e5a56ee
Show file tree

Hide file tree

Showing 3 changed files with 160 additions and 18 deletions.
diff --git a/logger.go b/logger.go
@@ -43,6 +43,7 @@ type Logger struct {
 
 	development bool
 	addCaller   bool
+	onPanic     zapcore.CheckWriteHook // default is WriteThenPanic
 	onFatal     zapcore.CheckWriteHook // default is WriteThenFatal
 
 	name        string
@@ -345,27 +346,12 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry {
 	// Set up any required terminal behavior.
 	switch ent.Level {
 	case zapcore.PanicLevel:
-		ce = ce.After(ent, zapcore.WriteThenPanic)
+		ce = ce.After(ent, terminalHookOverride(zapcore.WriteThenPanic, log.onPanic))
 	case zapcore.FatalLevel:
-		onFatal := log.onFatal
-		// nil or WriteThenNoop will lead to continued execution after
-		// a Fatal log entry, which is unexpected. For example,
-		//
-		//   f, err := os.Open(..)
-		//   if err != nil {
-		//     log.Fatal("cannot open", zap.Error(err))
-		//   }
-		//   fmt.Println(f.Name())
-		//
-		// The f.Name() will panic if we continue execution after the
-		// log.Fatal.
-		if onFatal == nil || onFatal == zapcore.WriteThenNoop {
-			onFatal = zapcore.WriteThenFatal
-		}
-		ce = ce.After(ent, onFatal)
+		ce = ce.After(ent, terminalHookOverride(zapcore.WriteThenFatal, log.onFatal))
 	case zapcore.DPanicLevel:
 		if log.development {
-			ce = ce.After(ent, zapcore.WriteThenPanic)
+			ce = ce.After(ent, terminalHookOverride(zapcore.WriteThenPanic, log.onPanic))
 		}
 	}
 
@@ -430,3 +416,20 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry {
 
 	return ce
 }
+
+func terminalHookOverride(defaultHook, override zapcore.CheckWriteHook) zapcore.CheckWriteHook {
+	// A nil or WriteThenNoop hook will lead to continued execution after
+	// a Panic or Fatal log entry, which is unexpected. For example,
+	//
+	//   f, err := os.Open(..)
+	//   if err != nil {
+	//     log.Fatal("cannot open", zap.Error(err))
+	//   }
+	//   fmt.Println(f.Name())
+	//
+	// The f.Name() will panic if we continue execution after the log.Fatal.
+	if override == nil || override == zapcore.WriteThenNoop {
+		return defaultHook
+	}
+	return override
+}
diff --git a/logger_test.go b/logger_test.go
@@ -836,6 +836,130 @@ func TestLoggerFatalOnNoop(t *testing.T) {
 	assert.Equal(t, 1, exitStub.Code, "must exit with status 1 for WriteThenNoop")
 }
 
+func TestLoggerCustomOnPanic(t *testing.T) {
+	tests := []struct {
+		msg          string
+		level        zapcore.Level
+		opts         []Option
+		finished     bool
+		want         []observer.LoggedEntry
+		recoverValue any
+	}{
+		{
+			msg:      "panic with nil hook",
+			level:    PanicLevel,
+			opts:     opts(WithPanicHook(nil)),
+			finished: false,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: PanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: "foobar",
+		},
+		{
+			msg:      "panic with noop hook",
+			level:    PanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenNoop)),
+			finished: false,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: PanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: "foobar",
+		},
+		{
+			msg:      "no panic with goexit hook",
+			level:    PanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenGoexit)),
+			finished: false,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: PanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: nil,
+		},
+		{
+			msg:      "dpanic no panic in development mode with goexit hook",
+			level:    DPanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenGoexit), Development()),
+			finished: false,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: DPanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: nil,
+		},
+		{
+			msg:      "dpanic panic in development mode with noop hook",
+			level:    DPanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenNoop), Development()),
+			finished: false,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: DPanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: "foobar",
+		},
+		{
+			msg:      "dpanic no exit in production mode with goexit hook",
+			level:    DPanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenPanic)),
+			finished: true,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: DPanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: nil,
+		},
+		{
+			msg:      "dpanic no panic in production mode with panic hook",
+			level:    DPanicLevel,
+			opts:     opts(WithPanicHook(zapcore.WriteThenPanic)),
+			finished: true,
+			want: []observer.LoggedEntry{
+				{
+					Entry:   zapcore.Entry{Level: DPanicLevel, Message: "foobar"},
+					Context: []Field{},
+				},
+			},
+			recoverValue: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.msg, func(t *testing.T) {
+			withLogger(t, InfoLevel, tt.opts, func(logger *Logger, logs *observer.ObservedLogs) {
+				var finished bool
+				recovered := make(chan any)
+				go func() {
+					defer func() {
+						recovered <- recover()
+					}()
+
+					logger.Log(tt.level, "foobar")
+					finished = true
+				}()
+
+				assert.Equal(t, tt.recoverValue, <-recovered, "unexpected value from recover()")
+				assert.Equal(t, tt.finished, finished, "expect goroutine finished state doesn't match")
+				assert.Equal(t, tt.want, logs.AllUntimed(), "unexpected logs")
+			})
+		})
+	}
+}
+
 func TestLoggerCustomOnFatal(t *testing.T) {
 	tests := []struct {
 		msg          string

diff --git a/options.go b/options.go
@@ -132,6 +132,21 @@ func IncreaseLevel(lvl zapcore.LevelEnabler) Option {
 	})
 }
 
+// WithPanicHook sets a CheckWriteHook to run on Panic/DPanic logs.
+// Zap will call this hook after writing a log statement with a Panic/DPanic level.
+//
+// For example, the following builds a logger that will exit the current
+// goroutine after writing a Panic/DPanic log message, but it will not start a panic.
+//
+//	zap.New(core, zap.WithPanicHook(zapcore.WriteThenGoexit))
+//
+// This is useful for testing Panic/DPanic log output.
+func WithPanicHook(hook zapcore.CheckWriteHook) Option {
+	return optionFunc(func(log *Logger) {
+		log.onPanic = hook
+	})
+}
+
 // OnFatal sets the action to take on fatal logs.
 //
 // Deprecated: Use [WithFatalHook] instead.