sorbet/pkg/logs/logger.go

package logs

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

var (
	// 使用东八区时间
	// https://cloud.tencent.com/developer/article/1805859
	cstZone        = time.FixedZone("CST", 8*3600)
	childLoggerKey = "sorbet/log:ChildLogger"
)

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

type Logger interface {
	SetFlags(flags int)
	Flags() int
	SetTimezone(loc *time.Location)
	Timezone() *time.Location
	SetLevel(level Level)
	Level() Level
	SetPersistWriter(w io.Writer)
	SetWriter(w io.Writer)
	With(args ...Attr) Logger
	WithGroup(name string) Logger
	Enabled(level Level) bool
	Log(level Level, msg string, args ...any)
	ForkLevel(level Level, msg string, args ...any) ChildLogger
	Trace(msg string, args ...any)
	ForkTrace(msg string, args ...any) ChildLogger
	Debug(msg string, args ...any)
	ForkDebug(msg string, args ...any) ChildLogger
	Info(msg string, args ...any)
	ForkInfo(msg string, args ...any) ChildLogger
	Warn(msg string, args ...any)
	ForkWarn(msg string, args ...any) ChildLogger
	Error(msg string, args ...any)
	ForkError(msg string, args ...any) ChildLogger
	Fatal(msg string, args ...any)
	ForkFatal(msg string, args ...any) ChildLogger
}

type ChildLogger interface {
	Print(msg string, args ...any)
	Finish()
}

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

type logger struct {
	parent  *logger
	isChild int32
	indent  int32

	level    int32
	flags    int32
	timezone unsafe.Pointer

	outMu            *sync.Mutex
	isPersistDiscard int32
	isDiscard        int32
	persistWriter    io.Writer
	writer           io.Writer

	handler slog.Handler
	l       *log.Logger
}

func New(opts *Options) Logger {
	if opts.Flags == 0 {
		opts.Flags = LstdFlags
	}
	if opts.Timezone == nil {
		opts.Timezone = cstZone
	}
	if opts.PersistWriter == nil {
		opts.PersistWriter = io.Discard
	}
	if opts.Writer == nil {
		opts.Writer = os.Stderr
	}

	var l *logger
	l = &logger{
		outMu:         &sync.Mutex{},
		persistWriter: opts.PersistWriter,
		writer:        opts.Writer,
		l:             log.New(opts.Writer, "", 0),
	}

	l.SetLevel(opts.Level)
	l.SetFlags(opts.Flags)
	l.SetTimezone(opts.Timezone)
	l.SetPersistWriter(opts.PersistWriter)
	l.SetWriter(opts.Writer)

	l.handler = slog.NewJSONHandler(opts.PersistWriter, &slog.HandlerOptions{
		AddSource:   true,
		Level:       opts.Level,
		ReplaceAttr: l.onAttr,
	})

	return l
}

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

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

func (l *logger) SetTimezone(loc *time.Location) {
	// FIXME(hupeh): 如何原子化储存结构体实例
	atomic.StorePointer(&l.timezone, unsafe.Pointer(loc))
}

func (l *logger) Timezone() *time.Location {
	return (*time.Location)(atomic.LoadPointer(&l.timezone))
}

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

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

func (l *logger) SetPersistWriter(w io.Writer) {
	l.outMu.Lock()
	defer l.outMu.Unlock()
	l.persistWriter = w
	atomic.StoreInt32(&l.isPersistDiscard, discard(w))
}

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

func discard(w io.Writer) int32 {
	if w == io.Discard {
		return 1
	}
	return 0
}

func (l *logger) onAttr(_ []string, a slog.Attr) slog.Attr {
	switch a.Key {
	case slog.LevelKey:
		level := a.Value.Any().(slog.Level)
		levelLabel := parseSlogLevel(level).String()
		a.Value = slog.StringValue(levelLabel)
	case slog.TimeKey:
		t := a.Value.Any().(time.Time)
		a.Value = slog.TimeValue(t.In(l.Timezone()))
	case slog.SourceKey:
		s := a.Value.Any().(*slog.Source)
		var as []Attr
		if s.Function != "" {
			as = append(as, String("func", s.Function))
		}
		if s.File != "" {
			as = append(as, String("file", s.File))
		}
		if s.Line != 0 {
			as = append(as, Int("line", s.Line))
		}
		a.Value = slog.GroupValue(as...)
	}
	return a
}

func (l *logger) Handle(ctx context.Context, r slog.Record) error {
	if atomic.LoadInt32(&l.isDiscard) == 0 {
		child, ok := ctx.Value(childLoggerKey).(*childLogger)
		indent, err := l.println(child, r)
		if err != nil {
			return err
		}
		if ok && indent > 0 {
			atomic.StoreInt32(&child.indent, indent)
		}
	}
	if atomic.LoadInt32(&l.isPersistDiscard) == 0 {
		return l.handler.Handle(ctx, r)
	}
	return nil
}

func (l *logger) println(child *childLogger, r slog.Record) (int32, error) {
	var output string
	var sep string
	var indent int32

	write := func(s string) {
		if sep == "" {
			sep = " "
		} else {
			output += sep
		}
		output += s
	}

	flags := l.Flags()
	colorful := flags&Lcolor != 0
	msg := r.Message
	level := parseSlogLevel(r.Level)
	levelStr := level.String()
	withChild := child != nil

	if withChild {
		indent = atomic.LoadInt32(&child.indent)
		withChild = indent > 0
	}

	if withChild {
		write(strings.Repeat(" ", int(indent)))
		indent = 0
	} else {
		if flags&(Ldate|Ltime|Lmicroseconds) != 0 {
			t := r.Time.In(l.Timezone())
			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"))
				}
			}
		}

		// 保存缩进
		indent += int32(len(output) + len(levelStr) + 3)

		if colorful {
			switch level {
			case LevelDebug:
				levelStr = FgCyan.Wrap(levelStr)
				msg = FgCyan.Wrap(msg)
			case LevelInfo:
				levelStr = FgBlue.Wrap(levelStr) + " "
				msg = FgBlue.Wrap(msg)
				indent += 1
			case LevelWarn:
				levelStr = FgYellow.Wrap(levelStr) + " "
				msg = FgYellow.Wrap(msg)
				indent += 1
			case LevelError:
				levelStr = FgRed.Wrap(levelStr)
				msg = FgRed.Wrap(msg)
			case LevelFatal:
				levelStr = FgMagenta.Wrap(levelStr)
				msg = FgMagenta.Wrap(msg)
			}
			levelStr = FgHiBlack.Wrap("[") + levelStr + FgHiBlack.Wrap("]")
		} else {
			levelStr = "[" + r.Level.String() + "]"
		}

		write(levelStr)
	}

	write(msg)

	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 = fgGreenItalic.Wrap(fileStr)
		}
		write(fileStr)
	}

	if numAttrs := r.NumAttrs(); flags&Lfields != 0 && numAttrs > 0 {
		fields := make(map[string]any, numAttrs)
		r.Attrs(func(a Attr) bool {
			fields[a.Key] = a.Value.Any()
			return true
		})
		b, err := json.Marshal(fields)
		if err != nil {
			return 0, err
		}
		fieldsStr := string(bytes.TrimSpace(b))
		if fieldsStr != "" {
			if colorful {
				fieldsStr = FgHiBlack.Wrap(fieldsStr)
			}
			write(fieldsStr)
		}
	}

	l.l.Println(output)

	return indent, nil
}

func (l *logger) clone() *logger {
	c := *l
	return &c
}

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

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

func (l *logger) Enabled(level Level) bool {
	return l.handler.Enabled(nil, level.slog().Level())
}

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

func (l *logger) ForkLevel(level Level, msg string, args ...any) ChildLogger {
	c := &childLogger{
		parent:  l,
		level:   level,
		indent:  0,
		records: make([]slog.Record, 0),
		closed:  make(chan struct{}),
	}
	c.Print(msg, args...)
	return c
}

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

func (l *logger) ForkTrace(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelTrace, msg, args...)
}

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

func (l *logger) ForkDebug(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelDebug, msg, args...)
}

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

func (l *logger) ForkInfo(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelInfo, msg, args...)
}

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

func (l *logger) ForkWarn(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelWarn, msg, args...)
}

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

func (l *logger) ForkError(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelError, msg, args...)
}

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

func (l *logger) ForkFatal(msg string, args ...any) ChildLogger {
	return l.ForkLevel(LevelFatal, msg, args...)
}

func (l *logger) log(ctx context.Context, level Level, msg string, args ...any) {
	if !l.Enabled(level) {
		return
	}
	if ctx == nil {
		ctx = context.Background()
	}
	_ = l.Handle(ctx, newRecord(level, msg, args))
}

func newRecord(level Level, msg string, args []any) slog.Record {
	//var pc uintptr
	//if !internal.IgnorePC {
	//	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)
	r := slog.NewRecord(time.Now(), level.slog().Level(), msg, 0)
	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 {
			r.Message = fmt.Sprintf(msg, sprintfArgs...)
		}
	}
	return r
}

type childLogger struct {
	parent  *logger
	level   Level
	indent  int32
	begin   slog.Record
	finish  slog.Record
	records []slog.Record
	closed  chan struct{}
}

func (c *childLogger) Print(msg string, args ...any) {
	select {
	case <-c.closed:
	default:
		c.records = append(
			c.records,
			newRecord(c.level, msg, args),
		)
	}
}

func (c *childLogger) Finish() {
	select {
	case <-c.closed:
		return
	default:
		close(c.closed)
	}

	ctx := context.Background()
	ctx = context.WithValue(ctx, childLoggerKey, c)
	for _, record := range c.records {
		_ = c.parent.Handle(ctx, record)
	}
}