diff --git a/zapio/writer.go b/zapio/writer.go new file mode 100644 index 000000000..29699c86d --- /dev/null +++ b/zapio/writer.go @@ -0,0 +1,120 @@ +// Copyright (c) 2021 Uber Technologies, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +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:] + + // Fast path: if we don't have a partial message from a previous write + // in the buffer, skip the buffer and log directly. + if w.buff.Len() == 0 { + w.log(line) + return + } + + 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 { + w.log(w.buff.Bytes()) + } + w.buff.Reset() +} + +func (w *Writer) log(b []byte) { + if ce := w.Log.Check(w.Level, string(b)); ce != nil { + ce.Write() + } +} diff --git a/zapio/writer_test.go b/zapio/writer_test.go new file mode 100644 index 000000000..9bdf3488d --- /dev/null +++ b/zapio/writer_test.go @@ -0,0 +1,248 @@ +// Copyright (c) 2021 Uber Technologies, Inc. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +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.") + }) + } +} + +func TestWrite_Sync(t *testing.T) { + t.Parallel() + + core, observed := observer.New(zap.InfoLevel) + + w := Writer{ + Log: zap.New(core), + Level: zap.InfoLevel, + } + + io.WriteString(&w, "foo") + io.WriteString(&w, "bar") + + t.Run("no sync", func(t *testing.T) { + assert.Zero(t, observed.Len(), "Expected no logs yet") + }) + + t.Run("sync", func(t *testing.T) { + defer observed.TakeAll() + + require.NoError(t, w.Sync(), "Sync must not fail") + + assert.Equal(t, []observer.LoggedEntry{ + {Entry: zapcore.Entry{Message: "foobar"}, Context: []zapcore.Field{}}, + }, observed.AllUntimed(), "Log messages did not match") + }) + + t.Run("sync on empty", func(t *testing.T) { + require.NoError(t, w.Sync(), "Sync must not fail") + assert.Zero(t, observed.Len(), "Expected no logs yet") + }) +} + +func BenchmarkWriter(b *testing.B) { + tests := []struct { + name string + writes [][]byte + }{ + { + name: "single", + writes: [][]byte{ + []byte("foobar\n"), + []byte("bazqux\n"), + }, + }, + { + name: "splits", + writes: [][]byte{ + []byte("foo"), + []byte("bar\nbaz"), + []byte("qux\n"), + }, + }, + } + + writer := Writer{ + Log: zap.New(new(partiallyNopCore)), + Level: zapcore.DebugLevel, + } + + for _, tt := range tests { + b.Run(tt.name, func(b *testing.B) { + b.ResetTimer() + + for i := 0; i < b.N; i++ { + for _, bs := range tt.writes { + writer.Write(bs) + } + } + }) + } +} + +// partiallyNopCore behaves exactly like NopCore except it always returns true +// for whether the provided level is enabled, and accepts all Check requests. +// +// This lets us measure the overhead of the writer without measuring the cost +// of logging. +type partiallyNopCore struct{} + +func (*partiallyNopCore) Enabled(zapcore.Level) bool { return true } + +func (c *partiallyNopCore) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry { + return ce.AddCore(ent, c) +} + +func (c *partiallyNopCore) With([]zapcore.Field) zapcore.Core { return c } +func (*partiallyNopCore) Write(zapcore.Entry, []zapcore.Field) error { return nil } +func (*partiallyNopCore) Sync() error { return nil }