-
Notifications
You must be signed in to change notification settings - Fork 0
/
writer.go
113 lines (93 loc) · 2.53 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
package helper
import (
"bytes"
"io"
"strings"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
type Writer struct {
Log *zap.Logger
buff bytes.Buffer
}
var (
_ zapcore.WriteSyncer = (*Writer)(nil)
_ io.Closer = (*Writer)(nil)
)
func (w *Writer) Write(bs []byte) (n int, err error) {
msg := string(bs)
level := zapcore.InfoLevel
if strings.HasPrefix(msg, "WARNING:") {
level = zapcore.WarnLevel
} else if strings.HasPrefix(msg, "ERROR:") {
level = zapcore.ErrorLevel
}
// Skip all checks if the level isn't enabled.
if !w.Log.Core().Enabled(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.
//
// Always call Close once you're done with the Writer to ensure that it flushes
// all data.
func (w *Writer) Close() error {
return w.Sync()
}
// Sync flushes buffered data to the logger as a new log entry even if it
// doesn't contain 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) {
msg := string(b)
level := zapcore.InfoLevel
if strings.HasPrefix(msg, "WARNING:") {
level = zapcore.WarnLevel
} else if strings.HasPrefix(msg, "ERROR:") {
level = zapcore.ErrorLevel
}
if ce := w.Log.Check(level, string(b)); ce != nil {
ce.Write()
}
}