diamond-orm/internal/logging/custom_handler.go

package logging

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"runtime"
	"slices"
	"strings"
	"sync"
	"text/template"
	"time"
)

type FormattedHandler struct {
	mu       *sync.Mutex
	out      io.Writer
	opts     Options
	attrs    map[string]slog.Value
	groups   []string
	groupLvl int
}

type Options struct {
	Level       slog.Leveler
	Format      string
	ReplaceAttr func(groups []string, attr slog.Attr) slog.Attr
}
type locData struct {
	FileName string
	Function string
	Line     int
}

func NewFormattedHandler(out io.Writer, options Options) *FormattedHandler {
	h := &FormattedHandler{
		opts:   options,
		out:    out,
		mu:     &sync.Mutex{},
		groups: make([]string, 0),
	}
	if h.opts.Format == "" {
		h.opts.Format = "{{.Time}} [{{.Level}}]"
	}
	if h.opts.Level == nil {
		h.opts.Level = slog.LevelInfo
	}
	return h
}

func (f *FormattedHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return level >= f.opts.Level.Level()
}

func (f *FormattedHandler) Handle(ctx context.Context, r slog.Record) error {
	bufp := allocBuf()
	buf := *bufp
	defer func() {
		*bufp = buf
		freeBuf(bufp)
	}()
	rep := f.opts.ReplaceAttr
	key := slog.LevelKey
	val := r.Level
	if rep == nil {
		r.AddAttrs(slog.String(key, val.String()))
	} else {
		nattr := slog.Any(key, val)
		nattr.Value = rep(f.groups, nattr).Value
		r.AddAttrs(nattr)
	}

	f.mu.Lock()
	defer f.mu.Unlock()
	tctx, tmpl := f.newFmtCtx(r)
	wr := bytes.NewBuffer(buf)
	parsed, err := tmpl.Parse(f.opts.Format)
	if err != nil {
		return err
	}
	err = parsed.Execute(wr, tctx)
	if err != nil {
		return err
	}

	wr.WriteByte('\n')
	_, err = f.out.Write(wr.Bytes())
	return err
}

func (f *FormattedHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return f
	}
	nf := f.clone()
	bufp := allocBuf()
	buf := *bufp
	defer func() {
		*bufp = buf
		freeBuf(bufp)
	}()
	s := f.newState(bytes.NewBuffer(buf))
	defer s.free()
	pos := s.buf.Len()
	s.startGroups()
	if !s.appendAttrs(attrs) {
		s.buf.Truncate(pos)
	} else {
		nf.groupLvl = len(nf.groups)
	}
	return nf
}

func (f *FormattedHandler) WithGroup(name string) slog.Handler {
	if name == "" {
		return f
	}
	f2 := f.clone()
	f2.groups = append(f2.groups, name)
	return f2
}

func (f *FormattedHandler) clone() *FormattedHandler {
	return &FormattedHandler{
		opts:     f.opts,
		groups:   slices.Clip(f.groups),
		out:      f.out,
		mu:       f.mu,
		groupLvl: f.groupLvl,
	}
}

type tmplData struct {
	Level    string
	Message  string
	RawTime  time.Time
	Time     string
	PC       uintptr
	Location locData
	Record   slog.Record
}

func hasBuiltInKey(a slog.Attr) bool {
	return a.Key == slog.MessageKey ||
		a.Key == slog.TimeKey ||
		a.Key == slog.SourceKey
}

func (f *FormattedHandler) newFmtCtx(r slog.Record) (ctx *tmplData, tmpl *template.Template) {
	tmpl = template.New("log")
	ctx = &tmplData{
		Message:  r.Message,
		RawTime:  r.Time,
		PC:       r.PC,
		Location: locData{},
	}
	if !r.Time.IsZero() {
		ctx.Time = r.Time.Format(time.RFC3339Nano)
	}
	r.Attrs(func(a slog.Attr) bool {
		if a.Key == slog.LevelKey {
			str := strings.ToUpper(a.Value.String())
			if rep := f.opts.ReplaceAttr; rep != nil {
				str = strings.ToUpper(a.Value.String())
			}
			ctx.Level = str
		}
		return true
	})
	if r.PC != 0 {
		frames := runtime.CallersFrames([]uintptr{r.PC})
		frame, _ := frames.Next()
		ctx.Location.FileName = frame.File
		ctx.Location.Function = frame.Function
		ctx.Location.Line = frame.Line
	}
	fm := make(map[string]any)
	fm["rest"] = func() string {
		bb := new(bytes.Buffer)
		s := f.newState(bb)
		defer s.free()
		s.begin(r)
		return s.buf.String()
	}
	tmpl = tmpl.Funcs(fm)
	return
}

var bufPool = sync.Pool{
	New: func() any {
		b := make([]byte, 0, 4096)
		return &b
	},
}

func allocBuf() *[]byte {
	return bufPool.Get().(*[]byte)
}

func freeBuf(b *[]byte) {
	const maxBufferSize = 16 << 10
	if cap(*b) <= maxBufferSize {
		*b = (*b)[:0]
		bufPool.Put(b)
	}
}