diff --git a/zapio/writer.go b/zapio/writer.go new file mode 100644 index 000000000..db509121c --- /dev/null +++ b/zapio/writer.go @@ -0,0 +1,88 @@ +package zapio + +import ( + "bytes" + "io" + + "go.uber.org/zap" + "go.uber.org/zap/zapcore" +) + +// Writer is an io.Writer that writes to the provided Zap logger, splitting log +// messages on line boundaries. +// +// Writer must be closed when finished to flush buffered data to the logger. +type Writer struct { + Log *zap.Logger // log to write to + Level zapcore.Level // log level to write at + + buff bytes.Buffer +} + +var ( + _ zapcore.WriteSyncer = (*Writer)(nil) + _ io.Closer = (*Writer)(nil) +) + +// Write writes the provided bytes to the underlying logger at the configured +// log level and returns the length of the bytes. +func (w *Writer) Write(bs []byte) (n int, err error) { + // Skip all checks if the level isn't enabled. + if !w.Log.Core().Enabled(w.Level) { + return len(bs), nil + } + + n = len(bs) + for len(bs) > 0 { + bs = w.writeLine(bs) + } + + return n, nil +} + +// writeLine writes a single line from the input, returning the remaining, +// unconsumed bytes. +func (w *Writer) writeLine(line []byte) (remaining []byte) { + idx := bytes.IndexByte(line, '\n') + if idx < 0 { + // If there are no newlines, buffer the entire string. + w.buff.Write(line) + return nil + } + + // Split on the newline, buffer and flush the left. + line, remaining = line[:idx], line[idx+1:] + w.buff.Write(line) + + // Log empty messages in the middle of the stream so that we don't lose + // information when the user writes "foo\n\nbar". + w.flush(true /* allowEmpty */) + + return remaining +} + +// Close closes the writer, flushing any buffered data in the process. +func (w *Writer) Close() error { + return w.Sync() +} + +// Sync flushes the buffered data from the writer, even if it doesn't end with +// a newline. +func (w *Writer) Sync() error { + // Don't allow empty messages on explicit Sync calls or on Close + // because we don't want an extraneous empty message at the end of the + // stream -- it's common for files to end with a newline. + w.flush(false /* allowEmpty */) + return nil +} + +// flush flushes the buffered data to the logger, allowing empty messages only +// if the bool is set. +func (w *Writer) flush(allowEmpty bool) { + if allowEmpty || w.buff.Len() > 0 { + if ce := w.Log.Check(w.Level, w.buff.String()); ce != nil { + ce.Write() + } + } + w.buff.Reset() +} diff --git a/zapio/writer_test.go b/zapio/writer_test.go new file mode 100644 index 000000000..dac236397 --- /dev/null +++ b/zapio/writer_test.go @@ -0,0 +1,138 @@ +package zapio + +import ( + "io" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/zap" + "go.uber.org/zap/zapcore" + "go.uber.org/zap/zaptest/observer" +) + +func TestWriter(t *testing.T) { + t.Parallel() + + tests := []struct { + desc string + level zapcore.Level // defaults to info + writes []string + want []zapcore.Entry + }{ + { + desc: "simple", + writes: []string{ + "foo\n", + "bar\n", + "baz\n", + }, + want: []zapcore.Entry{ + {Level: zap.InfoLevel, Message: "foo"}, + {Level: zap.InfoLevel, Message: "bar"}, + {Level: zap.InfoLevel, Message: "baz"}, + }, + }, + { + desc: "level too low", + level: zap.DebugLevel, + writes: []string{ + "foo\n", + "bar\n", + }, + want: []zapcore.Entry{}, + }, + { + desc: "multiple newlines in a message", + level: zap.WarnLevel, + writes: []string{ + "foo\nbar\n", + "baz\n", + "qux\nquux\n", + }, + want: []zapcore.Entry{ + {Level: zap.WarnLevel, Message: "foo"}, + {Level: zap.WarnLevel, Message: "bar"}, + {Level: zap.WarnLevel, Message: "baz"}, + {Level: zap.WarnLevel, Message: "qux"}, + {Level: zap.WarnLevel, Message: "quux"}, + }, + }, + { + desc: "message split across multiple writes", + level: zap.ErrorLevel, + writes: []string{ + "foo", + "bar\nbaz", + "qux", + }, + want: []zapcore.Entry{ + {Level: zap.ErrorLevel, Message: "foobar"}, + {Level: zap.ErrorLevel, Message: "bazqux"}, + }, + }, + { + desc: "blank lines in the middle", + writes: []string{ + "foo\n\nbar\nbaz", + }, + want: []zapcore.Entry{ + {Level: zap.InfoLevel, Message: "foo"}, + {Level: zap.InfoLevel, Message: ""}, + {Level: zap.InfoLevel, Message: "bar"}, + {Level: zap.InfoLevel, Message: "baz"}, + }, + }, + { + desc: "blank line at the end", + writes: []string{ + "foo\nbar\nbaz\n", + }, + want: []zapcore.Entry{ + {Level: zap.InfoLevel, Message: "foo"}, + {Level: zap.InfoLevel, Message: "bar"}, + {Level: zap.InfoLevel, Message: "baz"}, + }, + }, + { + desc: "multiple blank line at the end", + writes: []string{ + "foo\nbar\nbaz\n\n", + }, + want: []zapcore.Entry{ + {Level: zap.InfoLevel, Message: "foo"}, + {Level: zap.InfoLevel, Message: "bar"}, + {Level: zap.InfoLevel, Message: "baz"}, + {Level: zap.InfoLevel, Message: ""}, + }, + }, + } + + for _, tt := range tests { + tt := tt // for t.Parallel + t.Run(tt.desc, func(t *testing.T) { + t.Parallel() + + core, observed := observer.New(zap.InfoLevel) + + w := Writer{ + Log: zap.New(core), + Level: tt.level, + } + + for _, s := range tt.writes { + _, err := io.WriteString(&w, s) + require.NoError(t, err, "Writer.Write failed.") + } + + assert.NoError(t, w.Close(), "Writer.Close failed.") + + // Turn []observer.LoggedEntry => []zapcore.Entry + got := make([]zapcore.Entry, observed.Len()) + for i, ent := range observed.AllUntimed() { + got[i] = ent.Entry + } + assert.Equal(t, tt.want, got, "Logged entries do not match.") + }) + } +}