diff --git a/diode/diode_test.go b/diode/diode_test.go index 29b3f867..a8d25076 100644 --- a/diode/diode_test.go +++ b/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) diff --git a/log.go b/log.go index c6baa3ab..9fec7cc3 100644 --- a/log.go +++ b/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 diff --git a/writer.go b/writer.go index 50d7653d..138a612e 100644 --- a/writer.go +++ b/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.