zen/core/zlog/zlog.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)
	}
}