Make Log.Fatal() call Close on the writer before os.Exit(1) (#634)

Author: gmbuell

diode/diode_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"os/exec"
 	"testing"
 	"time"
 
@@ -38,6 +39,43 @@ func TestClose(t *testing.T) {
 	w.Close()
 }
 
+func TestFatal(t *testing.T) {
+	if os.Getenv("TEST_FATAL") == "1" {
+		w := diode.NewWriter(os.Stderr, 1000, 0, func(missed int) {
+			fmt.Printf("Dropped %d messages\n", missed)
+		})
+		defer w.Close()
+		log := zerolog.New(w)
+		log.Fatal().Msg("test")
+		return
+	}
+
+	cmd := exec.Command(os.Args[0], "-test.run=TestFatal")
+	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)
+	}
+	slurp, 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")
+	}
+
+	want := "{\"level\":\"fatal\",\"message\":\"test\"}\n"
+	got := string(slurp)
+	if got != want {
+		t.Errorf("Diode Fatal Test failed. got:%s, want:%s!", got, want)
+	}
+}
+
 func Benchmark(b *testing.B) {
 	log.SetOutput(io.Discard)
 	defer log.SetOutput(os.Stderr)

log.go

@@ -387,7 +387,14 @@ func (l *Logger) Err(err error) *Event {
 //
 // 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) { os.Exit(1) })
+	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.
+			closer.Close()
+		}
+		os.Exit(1)
+	})
 }
 
 // Panic starts a new message with panic level. The panic() function

writer.go

@@ -27,6 +27,13 @@ func (lw LevelWriterAdapter) WriteLevel(l Level, p []byte) (n int, err error) {
 	return lw.Write(p)
 }
 
+func (lw LevelWriterAdapter) Close() error {
+	if closer, ok := lw.Writer.(io.Closer); ok {
+		return closer.Close()
+	}
+	return nil
+}
+
 type syncWriter struct {
 	mu sync.Mutex
 	lw LevelWriter
@@ -57,6 +64,15 @@ func (s *syncWriter) WriteLevel(l Level, p []byte) (n int, err error) {
 	return s.lw.WriteLevel(l, p)
 }
 
+func (s *syncWriter) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if closer, ok := s.lw.(io.Closer); ok {
+		return closer.Close()
+	}
+	return nil
+}
+
 type multiLevelWriter struct {
 	writers []LevelWriter
 }
@@ -89,6 +105,17 @@ func (t multiLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) {
 	return n, err
 }
 
+func (t multiLevelWriter) Close() error {
+	for _, w := range t.writers {
+		if closer, ok := w.(io.Closer); ok {
+			if err := closer.Close(); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 // MultiLevelWriter creates a writer that duplicates its writes to all the
 // provided writers, similar to the Unix tee(1) command. If some writers
 // implement LevelWriter, their WriteLevel method will be used instead of Write.