/
writer.go
120 lines (101 loc) · 3.51 KB
/
writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
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()
}
}