zen/core/zlog/zlog.go

174 lines
3.8 KiB
Go

package zlog
import (
"context"
"io"
"log/slog"
"sync"
"time"
"github.com/fatih/color"
"github.com/tidwall/pretty"
)
type WriteLocker interface {
io.Writer
sync.Locker
}
func WrapLocker(w io.Writer) WriteLocker {
if wl, ok := w.(WriteLocker); ok {
return wl
}
return &writeLocker{Writer: w}
}
type writeLocker struct {
io.Writer
sync.Mutex
}
type Log struct {
level slog.Level
bufPool *bufferPool
jsonPool *jsonHandlerPool
withAttrs []slog.Attr
withGroup []string
opts *slog.HandlerOptions
pretty bool
color bool
writer WriteLocker
}
// Clone clones the Logger, replacing the output writer with w.
//
// If w is nil, the clone shares the original Logger's output writer.
func (log *Log) Clone(w io.Writer) *Log {
var output WriteLocker
if w == nil {
output = log.writer
} else {
output = WrapLocker(w)
}
withAttrs := make([]slog.Attr, len(log.withAttrs))
copy(withAttrs, log.withAttrs)
withGroup := make([]string, len(log.withGroup))
copy(withGroup, log.withGroup)
return &Log{
level: log.level,
withAttrs: withAttrs,
withGroup: withGroup,
opts: log.opts,
pretty: log.pretty,
bufPool: newBufferPool(),
jsonPool: newJsonHandlerPool(log.opts),
writer: output,
}
}
// Enabled implements the slog.Handler interface.
func (lo *Log) Enabled(ctx context.Context, lvl slog.Level) bool {
return lvl >= lo.level
}
// Handle implements the slog.Handler interface.
func (lo *Log) Handle(ctx context.Context, record slog.Record) error {
if !lo.pretty {
j := lo.jsonPool.Get()
defer j.Put()
return j.Handle(ctx, record)
}
var levelColor *color.Color
switch {
case record.Level >= slog.LevelError:
levelColor = color.New(color.FgRed)
case record.Level >= slog.LevelWarn:
levelColor = color.New(color.FgYellow)
case record.Level >= slog.LevelInfo:
levelColor = color.New(color.FgGreen)
default:
levelColor = color.New(color.FgWhite)
}
buf := lo.bufPool.Get()
defer lo.bufPool.Put(buf)
if record.PC != 0 && lo.opts.AddSource {
f := frame{Caller(record.PC).Frame()}
levelColor.Fprint(buf, f.ShortFile())
levelColor.Fprint(buf, ":")
levelColor.Fprint(buf, f.Line)
levelColor.Fprint(buf, " -- ")
levelColor.Fprint(buf, f.ShortFunction())
buf.WriteByte('\n')
}
if !record.Time.IsZero() {
const format = `[` + time.DateTime + `] `
b := record.Time.AppendFormat(nil, format)
buf.Write(b)
}
buf.WriteByte('[')
levelColor.Add(color.Bold).Fprint(buf, record.Level.String())
buf.WriteString("] ")
if record.Message != "" {
buf.WriteString(record.Message)
}
buf.WriteByte('\n')
jHandler := lo.jsonPool.Get()
defer jHandler.Put()
_ = jHandler.Handle(ctx, record)
if jHandler.buf.Len() > 3 { // Skip empty objects like "{}\n"
jsonData := jHandler.buf.Bytes()
jsonData = pretty.Pretty(jsonData)
if lo.color {
jsonData = pretty.Color(jsonData, nil)
}
buf.Write(jsonData)
}
buf.WriteByte('\n')
lo.writer.Lock()
defer lo.writer.Unlock()
_, err := buf.WriteTo(lo.writer)
return err
}
// WithAttrs implements the slog.Handler interface.
func (lo *Log) WithAttrs(attrs []slog.Attr) slog.Handler {
l := lo.Clone(nil)
l.withAttrs = append(l.withAttrs, attrs...)
return l
}
// WithGroup implements the slog.Handler interface.
func (lo *Log) WithGroup(name string) slog.Handler {
l := lo.Clone(nil)
l.withGroup = append(l.withGroup, name)
return l
}
type replaceAttrFunc = func([]string, slog.Attr) slog.Attr
// wrapReplaceAttr disables adding Time, Level, Source, Message
// fields when pretty mode is enabled.
func (lo *Log) wrapReplaceAttr(f replaceAttrFunc) replaceAttrFunc {
return func(s []string, a slog.Attr) slog.Attr {
if !lo.pretty {
return f(s, a)
}
switch a.Key {
case slog.TimeKey, slog.LevelKey, slog.SourceKey, slog.MessageKey:
return slog.Attr{}
}
return f(s, a)
}
}