sorbet/pkg/log/handler.go

package log

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"log/slog"
	"runtime"
	"strings"
	"sync/atomic"
	"time"
	"unsafe"
)

type Handler struct {
	ctx   unsafe.Pointer // 结构体 logger
	color Colorer
	slog.Handler
	attrs []Attr
	l     *log.Logger
}

func (h *Handler) WithAttrs(attrs []Attr) slog.Handler {
	return &Handler{
		ctx:     h.ctx,
		color:   h.color,
		Handler: h.Handler.WithAttrs(attrs),
		attrs:   append(h.attrs, attrs...),
		l:       h.l,
	}
}

func (h *Handler) WithGroup(name string) slog.Handler {
	return &Handler{
		ctx:     h.ctx,
		color:   h.color,
		Handler: h.Handler.WithGroup(name),
		l:       h.l,
	}
}

func (h *Handler) Handle(ctx context.Context, r slog.Record) error {
	// 参考 atomic.Pointer#Load 实现
	l := (*logger)(atomic.LoadPointer(&h.ctx))
	if atomic.LoadInt32(&l.discard) != 1 {
		if len(h.attrs) > 0 {
			x := slog.NewRecord(r.Time, r.Level, r.Message, r.PC)
			r.Attrs(func(attr slog.Attr) bool {
				x.AddAttrs(attr)
				return true
			})
			x.AddAttrs(h.attrs...)
			if err := h.Print(l, x); err != nil {
				return err
			}
		} else if err := h.Print(l, r); err != nil {
			return err
		}
	}
	if atomic.LoadInt32(&l.unpersist) != 1 {
		return h.Handler.Handle(ctx, r)
	}
	return nil
}

func (h *Handler) Print(l *logger, r slog.Record) error {
	flags := l.Flags()
	colorful := flags&Lcolor != 0
	level := parseSlogLevel(r.Level)
	levelStr := level.String()
	var fields map[string]any
	if numAttrs := r.NumAttrs(); numAttrs > 0 {
		fields = make(map[string]any, numAttrs)
		r.Attrs(func(a Attr) bool {
			switch a.Key {
			case rawTimeKey:
				r.Time = a.Value.Any().(time.Time)
			case rawLevelKey:
				levelStr = a.Value.Any().(string)
			default:
				fields[a.Key] = a.Value.Any()
			}
			return true
		})
	}
	var output string
	var sep string
	write := func(s string) {
		if sep == "" {
			sep = " "
		} else {
			output += sep
		}
		output += s
	}
	if flags&(Ldate|Ltime|Lmicroseconds) != 0 {
		t := r.Time.In(l.Timezone()) // todo 原子操作
		if flags&Ldate != 0 {
			write(t.Format("2006/01/02"))
		}
		if flags&(Ltime|Lmicroseconds) != 0 {
			if flags&Lmicroseconds != 0 {
				write(t.Format("15:04:05.000"))
			} else {
				write(t.Format("15:04:05"))
			}
		}
	}
	if colorful {
		// TODO(hupeh): 重新设计不同颜色
		colorize := identify
		switch level {
		case LevelDebug:
			colorize = h.color.Cyan
		case LevelInfo:
			colorize = h.color.Blue
		case LevelWarn:
			colorize = h.color.Yellow
		case LevelError:
			colorize = h.color.Red
		case LevelFatal, LevelPanic:
			colorize = h.color.Magenta
		}
		write(h.color.Grey("[") + colorize(levelStr) + h.color.Grey("]"))
		write(colorize(r.Message))
	} else {
		write("[" + levelStr + "]")
		write(r.Message)
	}
	if flags&(Lshortfile|Llongfile) != 0 && r.PC > 0 {
		var fileStr string
		fs := runtime.CallersFrames([]uintptr{r.PC})
		f, _ := fs.Next()
		file := f.File
		if flags&Lshortfile != 0 {
			i := strings.LastIndexAny(file, "\\/")
			if i > -1 {
				file = file[i+1:]
			}
		}
		fileStr = fmt.Sprintf("%s:%s:%d", f.Function, file, f.Line)
		if colorful {
			fileStr = h.color.Green(fileStr)
		}
		write(fileStr)
	}
	if flags&Lfields != 0 && len(fields) > 0 {
		b, err := json.Marshal(fields)
		if err != nil {
			return err
		}
		fieldsStr := string(bytes.TrimSpace(b))
		if fieldsStr != "" {
			if colorful {
				fieldsStr = h.color.White(fieldsStr)
			}
			write(fieldsStr)
		}
	}
	h.l.Println(strings.TrimSpace(output))
	return nil
}