Add ExitFunc for FatalLevel log messages

log.go

@@ -221,19 +221,25 @@ func (l Level) MarshalText() ([]byte, error) {
 	return []byte(LevelFieldMarshalFunc(l)), nil
 }
 
+// ExitFunc is a function type that takes an integer exit code and terminates the program.
+// By default, it is set to os.Exit, but it can be overridden to customize the behavior
+// of program termination (for example, for graceful shutdown or testing purposes).
+type ExitFunc func(int)
+
 // A Logger represents an active logging object that generates lines
 // of JSON output to an io.Writer. Each logging operation makes a single
 // call to the Writer's Write method. There is no guarantee on access
 // serialization to the Writer. If your Writer is not thread safe,
 // you may consider a sync wrapper.
 type Logger struct {
-	w       LevelWriter
-	level   Level
-	sampler Sampler
-	context []byte
-	hooks   []Hook
-	stack   bool
-	ctx     context.Context
+	w        LevelWriter
+	level    Level
+	sampler  Sampler
+	context  []byte
+	hooks    []Hook
+	stack    bool
+	ctx      context.Context
+	exitFunc ExitFunc
 }
 
 // New creates a root logger with given output writer. If the output writer implements
@@ -335,6 +341,12 @@ func (l Logger) Hook(hooks ...Hook) Logger {
 	return l
 }
 
+// ExitFunc returns a logger with the e ExitFunc.
+func (l Logger) ExitFunc(e ExitFunc) Logger {
+	l.exitFunc = e
+	return l
+}
+
 // Trace starts a new message with trace level.
 //
 // You must call Msg on the returned event in order to send the event.
@@ -382,18 +394,23 @@ func (l *Logger) Err(err error) *Event {
 	return l.Info()
 }
 
-// Fatal starts a new message with fatal level. The os.Exit(1) function
-// is called by the Msg method, which terminates the program immediately.
+// Fatal starts a new message with fatal level. The ExitFunc(1) function
+// (os.Exit by default) is called by the Msg method, which terminates the
+// program immediately.
 //
 // You must call Msg on the returned event in order to send the event.
 func (l *Logger) Fatal() *Event {
 	return l.newEvent(FatalLevel, func(msg string) {
 		if closer, ok := l.w.(io.Closer); ok {
 			// Close the writer to flush any buffered message. Otherwise the message
-			// will be lost as os.Exit() terminates the program immediately.
+			// will be lost as ExitFunc() (os.Exit by default) terminates the program immediately.
 			closer.Close()
 		}
-		os.Exit(1)
+		exitFunc := l.exitFunc
+		if exitFunc == nil {
+			exitFunc = os.Exit
+		}
+		exitFunc(1)
 	})
 }
 

log_test.go

@@ -5,7 +5,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
+	"os"
+	"os/exec"
 	"reflect"
 	"runtime"
 	"strconv"
@@ -78,6 +81,57 @@ func TestInfo(t *testing.T) {
 	})
 }
 
+func TestFatal(t *testing.T) {
+	t.Run("should exit", func(t *testing.T) {
+		if os.Getenv("TEST_FATAL") == "1" {
+			log := New(os.Stderr)
+			log.Fatal().Msg("")
+			return
+		}
+
+		cmd := exec.Command(os.Args[0], "-test.run=TestFatal/should_exit")
+		cmd.Env = append(os.Environ(), "TEST_FATAL=1")
+		stderr, err := cmd.StderrPipe()
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cmd.Start()
+		if err != nil {
+			t.Fatal(err)
+		}
+		out, err := io.ReadAll(stderr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cmd.Wait()
+		if err == nil {
+			t.Error("Expected log Fatal to exit with non-zero status")
+		}
+
+		if got, want := decodeIfBinaryToString(out), `{"level":"fatal"}`+"\n"; got != want {
+			t.Errorf("invalid log output:\n got:  %v\nwant: %v", got, want)
+		}
+	})
+
+	t.Run("should not exit", func(t *testing.T) {
+		out := &bytes.Buffer{}
+		log := New(out).ExitFunc(func(_ int) {})
+		log.Fatal().Msg("")
+		if got, want := decodeIfBinaryToString(out.Bytes()), `{"level":"fatal"}`+"\n"; got != want {
+			t.Errorf("invalid log output:\n got:  %v\nwant: %v", got, want)
+		}
+	})
+
+	t.Run("with level fatal should not exit", func(t *testing.T) {
+		out := &bytes.Buffer{}
+		log := New(out)
+		log.WithLevel(FatalLevel).Msg("")
+		if got, want := decodeIfBinaryToString(out.Bytes()), `{"level":"fatal"}`+"\n"; got != want {
+			t.Errorf("invalid log output:\n got:  %v\nwant: %v", got, want)
+		}
+	})
+}
+
 func TestEmptyLevelFieldName(t *testing.T) {
 	fieldName := LevelFieldName
 	LevelFieldName = ""