1 // Copyright (c) 2021 Uber Technologies, Inc. 2 // 3 // Permission is hereby granted, free of charge, to any person obtaining a copy 4 // of this software and associated documentation files (the "Software"), to deal 5 // in the Software without restriction, including without limitation the rights 6 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 // copies of the Software, and to permit persons to whom the Software is 8 // furnished to do so, subject to the following conditions: 9 // 10 // The above copyright notice and this permission notice shall be included in 11 // all copies or substantial portions of the Software. 12 // 13 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 // THE SOFTWARE. 20 21 // Package zapio provides tools for interacting with IO streams through Zap. 22 package zapio 23 24 import ( 25 "bytes" 26 "io" 27 28 "go.uber.org/zap" 29 "go.uber.org/zap/zapcore" 30 ) 31 32 // Writer is an io.Writer that writes to the provided Zap logger, splitting log 33 // messages on line boundaries. The Writer will buffer writes in memory until 34 // it encounters a newline, or the caller calls Sync or Close. 35 // 36 // Use the Writer with packages like os/exec where an io.Writer is required, 37 // and you want to log the output using your existing logger configuration. For 38 // example, 39 // 40 // writer := &zapio.Writer{Log: logger, Level: zap.DebugLevel} 41 // defer writer.Close() 42 // 43 // cmd := exec.CommandContext(ctx, ...) 44 // cmd.Stdout = writer 45 // cmd.Stderr = writer 46 // if err := cmd.Run(); err != nil { 47 // return err 48 // } 49 // 50 // Writer must be closed when finished to flush buffered data to the logger. 51 type Writer struct { 52 // Log specifies the logger to which the Writer will write messages. 53 // 54 // The Writer will panic if Log is unspecified. 55 Log *zap.Logger 56 57 // Log level for the messages written to the provided logger. 58 // 59 // If unspecified, defaults to Info. 60 Level zapcore.Level 61 62 buff bytes.Buffer 63 } 64 65 var ( 66 _ zapcore.WriteSyncer = (*Writer)(nil) 67 _ io.Closer = (*Writer)(nil) 68 ) 69 70 // Write writes the provided bytes to the underlying logger at the configured 71 // log level and returns the length of the bytes. 72 // 73 // Write will split the input on newlines and post each line as a new log entry 74 // to the logger. 75 func (w *Writer) Write(bs []byte) (n int, err error) { 76 // Skip all checks if the level isn't enabled. 77 if !w.Log.Core().Enabled(w.Level) { 78 return len(bs), nil 79 } 80 81 n = len(bs) 82 for len(bs) > 0 { 83 bs = w.writeLine(bs) 84 } 85 86 return n, nil 87 } 88 89 // writeLine writes a single line from the input, returning the remaining, 90 // unconsumed bytes. 91 func (w *Writer) writeLine(line []byte) (remaining []byte) { 92 idx := bytes.IndexByte(line, '\n') 93 if idx < 0 { 94 // If there are no newlines, buffer the entire string. 95 w.buff.Write(line) 96 return nil 97 } 98 99 // Split on the newline, buffer and flush the left. 100 line, remaining = line[:idx], line[idx+1:] 101 102 // Fast path: if we don't have a partial message from a previous write 103 // in the buffer, skip the buffer and log directly. 104 if w.buff.Len() == 0 { 105 w.log(line) 106 return 107 } 108 109 w.buff.Write(line) 110 111 // Log empty messages in the middle of the stream so that we don't lose 112 // information when the user writes "foo\n\nbar". 113 w.flush(true /* allowEmpty */) 114 115 return remaining 116 } 117 118 // Close closes the writer, flushing any buffered data in the process. 119 // 120 // Always call Close once you're done with the Writer to ensure that it flushes 121 // all data. 122 func (w *Writer) Close() error { 123 return w.Sync() 124 } 125 126 // Sync flushes buffered data to the logger as a new log entry even if it 127 // doesn't contain a newline. 128 func (w *Writer) Sync() error { 129 // Don't allow empty messages on explicit Sync calls or on Close 130 // because we don't want an extraneous empty message at the end of the 131 // stream -- it's common for files to end with a newline. 132 w.flush(false /* allowEmpty */) 133 return nil 134 } 135 136 // flush flushes the buffered data to the logger, allowing empty messages only 137 // if the bool is set. 138 func (w *Writer) flush(allowEmpty bool) { 139 if allowEmpty || w.buff.Len() > 0 { 140 w.log(w.buff.Bytes()) 141 } 142 w.buff.Reset() 143 } 144 145 func (w *Writer) log(b []byte) { 146 if ce := w.Log.Check(w.Level, string(b)); ce != nil { 147 ce.Write() 148 } 149 } 150