Skip to content

Commit

Permalink
Make Log.Fatal() call Close on the writer before os.Exit(1) (#634)
Browse files Browse the repository at this point in the history
  • Loading branch information
gmbuell committed Jan 13, 2024
1 parent 417580d commit c1ab4ed
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 1 deletion.
38 changes: 38 additions & 0 deletions diode/diode_test.go
Expand Up @@ -6,6 +6,7 @@ import (
"io"
"log"
"os"
"os/exec"
"testing"
"time"

Expand Down Expand Up @@ -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)
Expand Down
9 changes: 8 additions & 1 deletion log.go
Expand Up @@ -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
Expand Down
27 changes: 27 additions & 0 deletions writer.go
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit c1ab4ed

Please sign in to comment.