logger.go

  1package log
  2
  3import (
  4	"bytes"
  5	"errors"
  6	"fmt"
  7	"io"
  8	"os"
  9	"runtime"
 10	"strings"
 11	"sync"
 12	"sync/atomic"
 13	"time"
 14
 15	"github.com/charmbracelet/colorprofile"
 16)
 17
 18// ErrMissingValue is returned when a key is missing a value.
 19var ErrMissingValue = fmt.Errorf("missing value")
 20
 21// LoggerOption is an option for a logger.
 22type LoggerOption = func(*Logger)
 23
 24// Logger is a Logger that implements Logger.
 25type Logger struct {
 26	w  colorprofile.Writer
 27	b  bytes.Buffer
 28	mu *sync.RWMutex
 29
 30	isDiscard uint32
 31
 32	level           int64
 33	prefix          string
 34	timeFunc        TimeFunction
 35	timeFormat      string
 36	callerOffset    int
 37	callerFormatter CallerFormatter
 38	formatter       Formatter
 39
 40	reportCaller    bool
 41	reportTimestamp bool
 42
 43	fields []interface{}
 44
 45	helpers *sync.Map
 46	styles  *Styles
 47}
 48
 49// Logf logs a message with formatting.
 50func (l *Logger) Logf(level Level, format string, args ...interface{}) {
 51	l.Log(level, fmt.Sprintf(format, args...))
 52}
 53
 54// Log logs the given message with the given keyvals for the given level.
 55func (l *Logger) Log(level Level, msg interface{}, keyvals ...interface{}) {
 56	if atomic.LoadUint32(&l.isDiscard) != 0 {
 57		return
 58	}
 59
 60	// check if the level is allowed
 61	if atomic.LoadInt64(&l.level) > int64(level) {
 62		return
 63	}
 64
 65	var frame runtime.Frame
 66	if l.reportCaller {
 67		// Skip log.log, the caller, and any offset added.
 68		frames := l.frames(l.callerOffset + 2)
 69		for {
 70			f, more := frames.Next()
 71			_, helper := l.helpers.Load(f.Function)
 72			if !helper || !more {
 73				// Found a frame that wasn't a helper function.
 74				// Or we ran out of frames to check.
 75				frame = f
 76				break
 77			}
 78		}
 79	}
 80	l.handle(level, l.timeFunc(time.Now()), []runtime.Frame{frame}, msg, keyvals...)
 81}
 82
 83func (l *Logger) handle(level Level, ts time.Time, frames []runtime.Frame, msg interface{}, keyvals ...interface{}) {
 84	var kvs []interface{}
 85	if l.reportTimestamp && !ts.IsZero() {
 86		kvs = append(kvs, TimestampKey, ts)
 87	}
 88
 89	_, ok := l.styles.Levels[level]
 90	if ok {
 91		kvs = append(kvs, LevelKey, level)
 92	}
 93
 94	if l.reportCaller && len(frames) > 0 && frames[0].PC != 0 {
 95		file, line, fn := l.location(frames)
 96		if file != "" {
 97			caller := l.callerFormatter(file, line, fn)
 98			kvs = append(kvs, CallerKey, caller)
 99		}
100	}
101
102	if l.prefix != "" {
103		kvs = append(kvs, PrefixKey, l.prefix)
104	}
105
106	if msg != nil {
107		if m := fmt.Sprint(msg); m != "" {
108			kvs = append(kvs, MessageKey, m)
109		}
110	}
111
112	// append logger fields
113	kvs = append(kvs, l.fields...)
114	if len(l.fields)%2 != 0 {
115		kvs = append(kvs, ErrMissingValue)
116	}
117
118	// append the rest
119	kvs = append(kvs, keyvals...)
120	if len(keyvals)%2 != 0 {
121		kvs = append(kvs, ErrMissingValue)
122	}
123
124	l.mu.Lock()
125	defer l.mu.Unlock()
126	switch l.formatter {
127	case LogfmtFormatter:
128		l.logfmtFormatter(kvs...)
129	case JSONFormatter:
130		l.jsonFormatter(kvs...)
131	default:
132		l.textFormatter(kvs...)
133	}
134
135	// WriteTo will reset the buffer
136	if _, err := l.b.WriteTo(&l.w); err != nil {
137		if errors.Is(err, io.ErrShortWrite) {
138			// Reset the buffer even if the lengths don't match up. If we're
139			// using colorprofile's Writer, it will strip the ansi sequences based on
140			// the color profile which can cause this error.
141			l.b.Reset()
142		}
143	}
144}
145
146// Helper marks the calling function as a helper
147// and skips it for source location information.
148// It's the equivalent of testing.TB.Helper().
149func (l *Logger) Helper() {
150	l.helper(1)
151}
152
153func (l *Logger) helper(skip int) {
154	var pcs [1]uintptr
155	// Skip runtime.Callers, and l.helper
156	n := runtime.Callers(skip+2, pcs[:])
157	frames := runtime.CallersFrames(pcs[:n])
158	frame, _ := frames.Next()
159	l.helpers.LoadOrStore(frame.Function, struct{}{})
160}
161
162// frames returns the runtime.Frames for the caller.
163func (l *Logger) frames(skip int) *runtime.Frames {
164	// Copied from testing.T
165	const maxStackLen = 50
166	var pc [maxStackLen]uintptr
167
168	// Skip runtime.Callers, and l.frame
169	n := runtime.Callers(skip+2, pc[:])
170	frames := runtime.CallersFrames(pc[:n])
171	return frames
172}
173
174func (l *Logger) location(frames []runtime.Frame) (file string, line int, fn string) {
175	if len(frames) == 0 {
176		return "", 0, ""
177	}
178	f := frames[0]
179	return f.File, f.Line, f.Function
180}
181
182// Cleanup a path by returning the last n segments of the path only.
183func trimCallerPath(path string, n int) string {
184	// lovely borrowed from zap
185	// nb. To make sure we trim the path correctly on Windows too, we
186	// counter-intuitively need to use '/' and *not* os.PathSeparator here,
187	// because the path given originates from Go stdlib, specifically
188	// runtime.Caller() which (as of Mar/17) returns forward slashes even on
189	// Windows.
190	//
191	// See https://github.com/golang/go/issues/3335
192	// and https://github.com/golang/go/issues/18151
193	//
194	// for discussion on the issue on Go side.
195
196	// Return the full path if n is 0.
197	if n <= 0 {
198		return path
199	}
200
201	// Find the last separator.
202	idx := strings.LastIndexByte(path, '/')
203	if idx == -1 {
204		return path
205	}
206
207	for i := 0; i < n-1; i++ {
208		// Find the penultimate separator.
209		idx = strings.LastIndexByte(path[:idx], '/')
210		if idx == -1 {
211			return path
212		}
213	}
214
215	return path[idx+1:]
216}
217
218// SetReportTimestamp sets whether the timestamp should be reported.
219func (l *Logger) SetReportTimestamp(report bool) {
220	l.mu.Lock()
221	defer l.mu.Unlock()
222	l.reportTimestamp = report
223}
224
225// SetReportCaller sets whether the caller location should be reported.
226func (l *Logger) SetReportCaller(report bool) {
227	l.mu.Lock()
228	defer l.mu.Unlock()
229	l.reportCaller = report
230}
231
232// GetLevel returns the current level.
233func (l *Logger) GetLevel() Level {
234	l.mu.RLock()
235	defer l.mu.RUnlock()
236	return Level(l.level)
237}
238
239// SetLevel sets the current level.
240func (l *Logger) SetLevel(level Level) {
241	l.mu.Lock()
242	defer l.mu.Unlock()
243	atomic.StoreInt64(&l.level, int64(level))
244}
245
246// GetPrefix returns the current prefix.
247func (l *Logger) GetPrefix() string {
248	l.mu.RLock()
249	defer l.mu.RUnlock()
250	return l.prefix
251}
252
253// SetPrefix sets the current prefix.
254func (l *Logger) SetPrefix(prefix string) {
255	l.mu.Lock()
256	defer l.mu.Unlock()
257	l.prefix = prefix
258}
259
260// SetTimeFormat sets the time format.
261func (l *Logger) SetTimeFormat(format string) {
262	l.mu.Lock()
263	defer l.mu.Unlock()
264	l.timeFormat = format
265}
266
267// SetTimeFunction sets the time function.
268func (l *Logger) SetTimeFunction(f TimeFunction) {
269	l.mu.Lock()
270	defer l.mu.Unlock()
271	l.timeFunc = f
272}
273
274// SetOutput sets the output destination.
275func (l *Logger) SetOutput(w io.Writer) {
276	l.mu.Lock()
277	defer l.mu.Unlock()
278	if w == nil {
279		w = os.Stderr
280	}
281	l.w.Forward = w
282	var isDiscard uint32
283	if w == io.Discard {
284		isDiscard = 1
285	}
286	atomic.StoreUint32(&l.isDiscard, isDiscard)
287}
288
289// SetColorProfile force sets the underlying color profile for the
290// TextFormatter.
291func (l *Logger) SetColorProfile(profile colorprofile.Profile) {
292	l.mu.Lock()
293	defer l.mu.Unlock()
294	l.w.Profile = profile
295}
296
297// SetFormatter sets the formatter.
298func (l *Logger) SetFormatter(f Formatter) {
299	l.mu.Lock()
300	defer l.mu.Unlock()
301	l.formatter = f
302}
303
304// SetCallerFormatter sets the caller formatter.
305func (l *Logger) SetCallerFormatter(f CallerFormatter) {
306	l.mu.Lock()
307	defer l.mu.Unlock()
308	l.callerFormatter = f
309}
310
311// SetCallerOffset sets the caller offset.
312func (l *Logger) SetCallerOffset(offset int) {
313	l.mu.Lock()
314	defer l.mu.Unlock()
315	l.callerOffset = offset
316}
317
318// SetStyles sets the logger styles for the TextFormatter.
319func (l *Logger) SetStyles(s *Styles) {
320	if s == nil {
321		s = DefaultStyles()
322	}
323	l.mu.Lock()
324	defer l.mu.Unlock()
325	l.styles = s
326}
327
328// With returns a new logger with the given keyvals added.
329func (l *Logger) With(keyvals ...interface{}) *Logger {
330	var st Styles
331	l.mu.Lock()
332	sl := *l
333	st = *l.styles
334	l.mu.Unlock()
335	sl.b = bytes.Buffer{}
336	sl.mu = &sync.RWMutex{}
337	sl.helpers = &sync.Map{}
338	sl.fields = append(make([]interface{}, 0, len(l.fields)+len(keyvals)), l.fields...)
339	sl.fields = append(sl.fields, keyvals...)
340	sl.styles = &st
341	return &sl
342}
343
344// WithPrefix returns a new logger with the given prefix.
345func (l *Logger) WithPrefix(prefix string) *Logger {
346	sl := l.With()
347	sl.SetPrefix(prefix)
348	return sl
349}
350
351// Debug prints a debug message.
352func (l *Logger) Debug(msg interface{}, keyvals ...interface{}) {
353	l.Log(DebugLevel, msg, keyvals...)
354}
355
356// Info prints an info message.
357func (l *Logger) Info(msg interface{}, keyvals ...interface{}) {
358	l.Log(InfoLevel, msg, keyvals...)
359}
360
361// Warn prints a warning message.
362func (l *Logger) Warn(msg interface{}, keyvals ...interface{}) {
363	l.Log(WarnLevel, msg, keyvals...)
364}
365
366// Error prints an error message.
367func (l *Logger) Error(msg interface{}, keyvals ...interface{}) {
368	l.Log(ErrorLevel, msg, keyvals...)
369}
370
371// Fatal prints a fatal message and exits.
372func (l *Logger) Fatal(msg interface{}, keyvals ...interface{}) {
373	l.Log(FatalLevel, msg, keyvals...)
374	os.Exit(1)
375}
376
377// Print prints a message with no level.
378func (l *Logger) Print(msg interface{}, keyvals ...interface{}) {
379	l.Log(noLevel, msg, keyvals...)
380}
381
382// Debugf prints a debug message with formatting.
383func (l *Logger) Debugf(format string, args ...interface{}) {
384	l.Log(DebugLevel, fmt.Sprintf(format, args...))
385}
386
387// Infof prints an info message with formatting.
388func (l *Logger) Infof(format string, args ...interface{}) {
389	l.Log(InfoLevel, fmt.Sprintf(format, args...))
390}
391
392// Warnf prints a warning message with formatting.
393func (l *Logger) Warnf(format string, args ...interface{}) {
394	l.Log(WarnLevel, fmt.Sprintf(format, args...))
395}
396
397// Errorf prints an error message with formatting.
398func (l *Logger) Errorf(format string, args ...interface{}) {
399	l.Log(ErrorLevel, fmt.Sprintf(format, args...))
400}
401
402// Fatalf prints a fatal message with formatting and exits.
403func (l *Logger) Fatalf(format string, args ...interface{}) {
404	l.Log(FatalLevel, fmt.Sprintf(format, args...))
405	os.Exit(1)
406}
407
408// Printf prints a message with no level and formatting.
409func (l *Logger) Printf(format string, args ...interface{}) {
410	l.Log(noLevel, fmt.Sprintf(format, args...))
411}