sorbet/pkg/log/logger.go

package log

import (
	"fmt"
	"io"
	"log"
	"log/slog"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
	"unsafe"
)

const (
	Ldate = 1 << iota
	Ltime
	Lmicroseconds
	Llongfile
	Lshortfile
	Lfields
	Lcolor
	LstdFlags = Ltime | Lmicroseconds | Lfields | Lcolor
)

type Options struct {
	Flags         int
	Level         Level
	IgnorePC      bool
	Timezone      *time.Location
	Colorer       Colorer
	PersistWriter io.Writer
	Writer        io.Writer
}

type Logger interface {
	Flags() int
	SetFlags(flags int)
	Level() Level
	SetLevel(Level)
	Timezone() *time.Location
	SetTimezone(loc *time.Location)
	IgnorePC() bool
	SetIgnorePC(ignore bool)
	Enabled(level Level) bool
	Writer() io.Writer
	SetWriter(w io.Writer)
	PersistWriter() io.Writer
	SetPersistWriter(w io.Writer)
	With(attrs ...Attr) Logger
	WithGroup(name string) Logger
	Log(level Level, msg string, args ...any)
	Trace(msg string, args ...any)
	Debug(msg string, args ...any)
	Info(msg string, args ...any)
	Warn(msg string, args ...any)
	Error(msg string, args ...any)
	Panic(msg string, args ...any)
	Fatal(msg string, args ...any)
}

var _ Logger = &logger{}

type logger struct {
	flags         int32
	level         int32
	ignorePC      int32
	unpersist     int32 // 忽略持久化写入
	discard       int32 // 忽略控制台输出
	timezone      unsafe.Pointer
	mu            *sync.Mutex
	persistWriter io.Writer
	writer        io.Writer
	handler       slog.Handler
}

func New(o *Options) Logger {
	if o.Flags == 0 {
		o.Flags = LstdFlags
	}
	if o.Timezone == nil {
		o.Timezone = time.FixedZone("CST", 8*3600) // 使用东八区时间
	}
	if o.Writer == nil {
		o.Writer = os.Stderr
	}
	if o.PersistWriter == nil {
		o.PersistWriter = io.Discard
	}
	if o.Colorer == nil {
		o.Colorer = NewColorer()
	}

	l := &logger{mu: &sync.Mutex{}}

	l.handler = &Handler{
		ctx:   unsafe.Pointer(l),
		color: o.Colorer,
		Handler: slog.NewJSONHandler(
			&mutablePersistWriter{l},
			&slog.HandlerOptions{
				AddSource:   true,
				Level:       &mutableLevel{l},
				ReplaceAttr: createAttrReplacer(l),
			},
		),
		l: log.New(&mutableWriter{l}, "", 0),
	}

	l.SetFlags(o.Flags)
	l.SetLevel(o.Level)
	l.SetIgnorePC(o.IgnorePC)
	l.SetTimezone(o.Timezone)
	l.SetPersistWriter(o.PersistWriter)
	l.SetWriter(o.Writer)

	return l
}

func (l *logger) Flags() int {
	return int(atomic.LoadInt32(&l.flags))
}

func (l *logger) SetFlags(flags int) {
	atomic.StoreInt32(&l.flags, int32(flags))
}

func (l *logger) Level() Level {
	return Level(int(atomic.LoadInt32(&l.level)))
}

func (l *logger) SetLevel(level Level) {
	atomic.StoreInt32(&l.level, int32(level))
}

func (l *logger) Timezone() *time.Location {
	// 参考 atomic.Pointer#Load 实现
	return (*time.Location)(atomic.LoadPointer(&l.timezone))
}

func (l *logger) SetTimezone(loc *time.Location) {
	// 参考 atomic.Pointer#Store 实现
	atomic.StorePointer(&l.timezone, unsafe.Pointer(loc))
}

func (l *logger) IgnorePC() bool {
	return atomic.LoadInt32(&l.ignorePC) == 1
}

func (l *logger) SetIgnorePC(ignore bool) {
	atomic.StoreInt32(&l.ignorePC, bool2int32(ignore))
}

func (l *logger) Enabled(level Level) bool {
	// todo panic 和 fatal 用于开启，取消 off 等级
	return l.handler.Enabled(nil, level.slog().Level())
}

func (l *logger) SetWriter(w io.Writer) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.writer = w
	atomic.StoreInt32(&l.discard, bool2int32(w == io.Discard))
}

func (l *logger) Writer() io.Writer {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.writer
}

func (l *logger) SetPersistWriter(w io.Writer) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.persistWriter = w
	atomic.StoreInt32(&l.unpersist, bool2int32(w == io.Discard))
}

func (l *logger) PersistWriter() io.Writer {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.persistWriter
}

func (l *logger) With(attrs ...Attr) Logger {
	if len(attrs) == 0 {
		return l
	}
	c := l.clone()
	c.handler = l.handler.WithAttrs(attrs).(*Handler)
	return c
}

func (l *logger) WithGroup(name string) Logger {
	if name == "" {
		return l
	}
	c := l.clone()
	c.handler = l.handler.WithGroup(name).(*Handler)
	return c
}

func (l *logger) clone() *logger {
	c := *l
	// TODO(hupeh): 测试 clone 是否报错
	//c.writer = l.writer
	//c.persistWriter = l.persistWriter
	return &c
}

// Log logs at level.
func (l *logger) Log(level Level, msg string, args ...any) {
	l.log(level, msg, args...)
}

// Trace logs at LevelTrace.
func (l *logger) Trace(msg string, args ...any) {
	l.log(LevelTrace, msg, args...)
}

// Debug logs at LevelDebug.
func (l *logger) Debug(msg string, args ...any) {
	l.log(LevelDebug, msg, args...)
}

// Info logs at LevelInfo.
func (l *logger) Info(msg string, args ...any) {
	l.log(LevelInfo, msg, args...)
}

// Warn logs at LevelWarn.
func (l *logger) Warn(msg string, args ...any) {
	l.log(LevelWarn, msg, args...)
}

// Error logs at LevelError.
func (l *logger) Error(msg string, args ...any) {
	l.log(LevelError, msg, args...)
}

// Panic logs at PanicError.
func (l *logger) Panic(msg string, args ...any) {
	l.log(LevelPanic, msg, args...)
}

// Fatal logs at LevelFatal.
func (l *logger) Fatal(msg string, args ...any) {
	l.log(LevelFatal, msg, args...)
}

func (l *logger) log(level Level, msg string, args ...any) {
	if !l.Enabled(level) {
		return
	}
	var pc uintptr
	if atomic.LoadInt32(&l.ignorePC) != 1 {
		var pcs [1]uintptr
		// skip [runtime.Callers, this function, this function's caller]
		runtime.Callers(3, pcs[:])
		pc = pcs[0]
	}
	r := slog.NewRecord(time.Now(), level.slog().Level(), msg, pc)
	if len(args) > 0 {
		var sprintfArgs []any
		for _, arg := range args {
			switch v := arg.(type) {
			case Attr:
				r.AddAttrs(v)
			default:
				sprintfArgs = append(sprintfArgs, arg)
			}
		}
		if len(sprintfArgs) > 0 {
			msg = fmt.Sprintf(msg, sprintfArgs...)
		}
		r.Message = msg
	}
	_ = l.handler.Handle(nil, r)
}