From 54caddb9ed8291ce97de16bd381ab96b675d5c40 Mon Sep 17 00:00:00 2001 From: Mitar Date: Sun, 15 Oct 2023 23:45:29 +0200 Subject: [PATCH] Add TriggerLevelWriter. See: https://github.com/rs/zerolog/issues/583 --- writer.go | 143 +++++++++++++++++++++++++++++++++++++++++++++++++ writer_test.go | 55 +++++++++++++++++++ 2 files changed, 198 insertions(+) diff --git a/writer.go b/writer.go index 9b9ef88e..1b2f7ec5 100644 --- a/writer.go +++ b/writer.go @@ -2,6 +2,7 @@ package zerolog import ( "bytes" + "errors" "io" "path" "runtime" @@ -180,3 +181,145 @@ func (w *FilteredLevelWriter) WriteLevel(level Level, p []byte) (int, error) { } return len(p), nil } + +var triggerWriterPool = &sync.Pool{ + New: func() interface{} { + return new(bytes.Buffer) + }, +} + +// NewTriggerLevelWriter returns a new TriggerLevelWriter. +// +// It obtains a buffer from the pool and you must +// call Close to return the buffer to the pool. +func NewTriggerLevelWriter(w io.Writer, conditionalLevel, triggerLevel Level) *TriggerLevelWriter { + return &TriggerLevelWriter{ + Writer: w, + ConditionalLevel: conditionalLevel, + TriggerLevel: triggerLevel, + buf: triggerWriterPool.Get().(*bytes.Buffer), + } +} + +// TriggerLevelWriter buffers log lines at the ConditionalLevel or below +// until a trigger level (or higher) line is emitted. Log lines with level +// higher than ConditionalLevel are always written out to the destination +// writer. If trigger never happens, buffered log lines are never written out. +// +// It can be used to configure "log level per request". You should create a +// new TriggerLevelWriter (using NewTriggerLevelWriter) per request. +type TriggerLevelWriter struct { + // Destination writer. If LevelWriter is provided (usually), its WriteLevel is used + // instead of Write. + io.Writer + + // ConditionalLevel is the level (and below) at which lines are buffered until + // a trigger level (or higher) line is emitted. Usually this is set to DebugLevel. + ConditionalLevel Level + + // TriggerLevel is the lowest level that triggers the sending of the conditional + // level lines. Usually this is set to ErrorLevel. + TriggerLevel Level + + buf *bytes.Buffer + triggered bool + mu sync.Mutex +} + +func (w *TriggerLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) { + w.mu.Lock() + defer w.mu.Unlock() + + if w.buf == nil { + return 0, errors.New("invalid writer") + } + + // At first trigger level or above log line, we flush the buffer and change the + // trigger state to triggered. + if !w.triggered && l >= w.TriggerLevel { + err := w.trigger() + if err != nil { + return 0, err + } + } + + // Unless triggered, we buffer everything at and below ConditionalLevel. + if !w.triggered && l <= w.ConditionalLevel { + // We prefix each log line with a byte with the level. + // Hopefully we will never have a level value which equals a newline + // (which could interfere with reconstruction of log lines in the trigger method). + w.buf.WriteByte(byte(l)) + w.buf.Write(p) + return len(p), nil + } + + // Anything above ConditionalLevel is always passed through. + // Once triggered, everything is passed through. + if lw, ok := w.Writer.(LevelWriter); ok { + return lw.WriteLevel(l, p) + } + return w.Write(p) +} + +// trigger expects lock to be held. +func (w *TriggerLevelWriter) trigger() error { + if w.triggered { + return nil + } + w.triggered = true + defer w.buf.Reset() + + p := w.buf.Bytes() + for len(p) > 0 { + // We do not use bufio.Scanner here because we already have full buffer + // in the memory and we do not want extra copying from the buffer to + // scanner's token slice, nor we want to hit scanner's token size limit, + // and we also want to preserve newlines. + i := bytes.IndexByte(p, '\n') + line := p[0 : i+1] + p = p[i+1:] + // We prefixed each log line with a byte with the level. + level := Level(line[0]) + line = line[1:] + var err error + if lw, ok := w.Writer.(LevelWriter); ok { + _, err = lw.WriteLevel(level, line) + } else { + _, err = w.Write(line) + } + if err != nil { + return err + } + } + + return nil +} + +// Trigger forces flushing the buffer and change the trigger state to +// triggered, if the writer has not already been triggered before. +func (w *TriggerLevelWriter) Trigger() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.buf == nil { + return errors.New("invalid writer") + } + + return w.trigger() +} + +// Close closes the writer and returns the buffer to the pool. +func (w *TriggerLevelWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.buf == nil { + return nil + } + + w.buf.Reset() + triggerWriterPool.Put(w.buf) + w.buf = nil + + return nil +} diff --git a/writer_test.go b/writer_test.go index 60595ba6..0ff20833 100644 --- a/writer_test.go +++ b/writer_test.go @@ -195,3 +195,58 @@ func TestFilteredLevelWriter(t *testing.T) { t.Errorf("Expected %q, got %q.", want, p) } } + +type testWrite struct { + Level + Line []byte +} + +func TestTriggerLevelWriter(t *testing.T) { + tests := []struct { + write []testWrite + want []byte + all []byte + }{{ + []testWrite{ + {DebugLevel, []byte("no\n")}, + {InfoLevel, []byte("yes\n")}, + }, + []byte("yes\n"), + []byte("yes\nno\n"), + }, { + []testWrite{ + {DebugLevel, []byte("yes1\n")}, + {InfoLevel, []byte("yes2\n")}, + {ErrorLevel, []byte("yes3\n")}, + {DebugLevel, []byte("yes4\n")}, + }, + []byte("yes2\nyes1\nyes3\nyes4\n"), + []byte("yes2\nyes1\nyes3\nyes4\n"), + }} + + for k, tt := range tests { + t.Run(fmt.Sprintf("case=%d", k), func(t *testing.T) { + buf := bytes.Buffer{} + writer := NewTriggerLevelWriter(LevelWriterAdapter{&buf}, DebugLevel, ErrorLevel) + t.Cleanup(func() { writer.Close() }) + for _, w := range tt.write { + _, err := writer.WriteLevel(w.Level, w.Line) + if err != nil { + t.Error(err) + } + } + p := buf.Bytes() + if want := tt.want; !bytes.Equal([]byte(want), p) { + t.Errorf("Expected %q, got %q.", want, p) + } + err := writer.Trigger() + if err != nil { + t.Error(err) + } + p = buf.Bytes() + if want := tt.all; !bytes.Equal([]byte(want), p) { + t.Errorf("Expected %q, got %q.", want, p) + } + }) + } +}