laic-backend/logger/logger.go

package logger

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
)

const (
	Level_T Level = iota
	Level_D
	Level_I
	Level_W
	Level_E
)

const (
	Reset   = "\033[0m"
	Red     = "\033[31m"
	Green   = "\033[32m"
	Yellow  = "\033[33m"
	Blue    = "\033[34m"
	Magenta = "\033[35m"
	Cyan    = "\033[36m"
	White   = "\033[37m"
)

// LogConfig 日志配置
type LogConfig struct {
	LogPath    string
	MaxSize    int
	MaxBackups int
	MaxAge     int
	Compress   bool
}

var defaultConfig = LogConfig{
	LogPath:    "./log",
	MaxSize:    50,
	MaxBackups: 20,
	MaxAge:     30,
	Compress:   true,
}

type Level int

// TraceContext 存储跟踪信息
type TraceContext struct {
	TraceID string
	SpanID  string
}

// goroutineID 到 traceID 的映射
var traceMap sync.Map

// 获取当前 goroutine 的 ID
func getGoroutineID() uint64 {
	b := make([]byte, 64)
	b = b[:runtime.Stack(b, false)]
	b = bytes.TrimPrefix(b, []byte("goroutine "))
	b = b[:bytes.IndexByte(b, ' ')]
	n, _ := strconv.ParseUint(string(b), 10, 64)
	return n
}

// SetTraceID 设置当前 goroutine 的 traceID
func SetTraceID(traceID string) {
	traceMap.Store(getGoroutineID(), traceID)
}

// GetTraceID 获取当前 goroutine 的 traceID
func GetTraceID() string {
	if id, ok := traceMap.Load(getGoroutineID()); ok {
		return id.(string)
	}
	return ""
}

// CleanTraceID 清理当前 goroutine 的 traceID
func CleanTraceID() {
	traceMap.Delete(getGoroutineID())
}

type customLogWriter struct {
	prefix string
	output io.Writer
}

var DefaultWriter io.Writer

var (
	loglevel Level
	_trace   *log.Logger
	_debug   *log.Logger
	_info    *log.Logger
	_warn    *log.Logger
	_error   *log.Logger
)

func SetLevel(level Level) {
	loglevel = level
}

func GetLevel() string {
	return loglevel.String()
}

func formatLogWithTrace(data ...any) string {
	logData := formatLogData(data...)
	return fmt.Sprintf("===> %v", logData)
}

func formatLogData(data ...any) string {
	var formattedData []string
	for _, item := range data {
		switch v := item.(type) {
		case string:
			formattedData = append(formattedData, v)
		case int, int64, float64, bool:
			formattedData = append(formattedData, fmt.Sprintf("%v", v))
		case []interface{}:
			for _, elem := range v {
				formattedData = append(formattedData, fmt.Sprintf("%+v", elem))
			}
		default:
			formattedData = append(formattedData, fmt.Sprintf("%+v", v))
		}
	}
	return strings.Join(formattedData, " ")
}

func TRACE(data ...any) {
	if loglevel <= Level_T {
		_trace.Println(formatLogWithTrace(data...))
	}
}

func INFO(data ...any) {
	if loglevel <= Level_I {
		_info.Println(formatLogWithTrace(data...))
	}
}

func DEBUG(data ...any) {
	if loglevel <= Level_D {
		_debug.Println(formatLogWithTrace(data...))
	}
}

func ERROR(data any, e error) {
	if loglevel <= Level_E {
		_error.Println(formatLogWithTrace(data), e)
	}
}

func WARN(data ...any) {
	if loglevel <= Level_W {
		_warn.Println(formatLogWithTrace(data...))
	}
}

// AsyncWriter 异步写入器
type AsyncWriter struct {
	output io.Writer
	ch     chan []byte
	wg     sync.WaitGroup
}

func NewAsyncWriter(output io.Writer) *AsyncWriter {
	w := &AsyncWriter{
		output: output,
		ch:     make(chan []byte, 1000),
	}
	w.wg.Add(1)
	go w.run()
	return w
}

func NewAsyncWriterWithSize(output io.Writer, size int) *AsyncWriter {
	w := &AsyncWriter{
		output: output,
		ch:     make(chan []byte, size),
	}
	w.wg.Add(1)
	go w.run()
	return w
}

func (w *AsyncWriter) run() {
	defer w.wg.Done()
	for data := range w.ch {
		w.output.Write(data)
	}
}

func (w *AsyncWriter) Write(p []byte) (n int, err error) {
	w.ch <- append([]byte(nil), p...)
	return len(p), nil
}

func (w *AsyncWriter) Close() {
	close(w.ch)
	w.wg.Wait()
}

func InitCustomLog(level Level, config *LogConfig, logName string) {
	if config == nil {
		config = &defaultConfig
	}

	SetLevel(level)

	if err := os.MkdirAll(config.LogPath, 0o755); err != nil {
		panic(err)
	}

	infoLogger := newLevelLogger(fmt.Sprintf("%s-info", logName), config)
	errorLogger := newLevelLogger(fmt.Sprintf("%s-error", logName), config)
	env := os.Getenv("LAIC_ENV")

	var output io.Writer
	if env != "prod" && loglevel <= Level_D {
		output = io.MultiWriter(os.Stdout, infoLogger)
	} else {
		output = infoLogger
	}

	DefaultWriter = NewAsyncWriter(output)

	var errorOutput io.Writer
	if loglevel <= Level_D {
		errorOutput = NewAsyncWriter(io.MultiWriter(os.Stdout, errorLogger, infoLogger))
	} else {
		errorOutput = NewAsyncWriter(io.MultiWriter(errorLogger, infoLogger))
	}

	_trace = newLogger(Level_T, Magenta, DefaultWriter)
	_debug = newLogger(Level_D, Blue, DefaultWriter)
	_info = newLogger(Level_I, Green, DefaultWriter)
	_warn = newLogger(Level_W, Yellow, DefaultWriter)
	_error = newLogger(Level_E, Red, errorOutput)

	log.SetOutput(DefaultWriter)
	log.SetFlags(0)
}

func newLevelLogger(level string, config *LogConfig) io.Writer {
	return &lumberjack.Logger{
		Filename:   filepath.Join(config.LogPath, level+".log"),
		MaxSize:    config.MaxSize,
		MaxBackups: config.MaxBackups,
		MaxAge:     config.MaxAge,
		Compress:   config.Compress,
	}
}

func newLogger(level Level, color string, output io.Writer) *log.Logger {
	logger := log.New(output, "", 0)
	logger.SetOutput(&customLogWriter{
		prefix: fmt.Sprintf("%s[%s]: %s", color, strings.ToUpper(level.String()), Reset),
		output: output,
	})
	return logger
}

// String 方法
func (l Level) String() string {
	switch l {
	case Level_T:
		return "TRACE"
	case Level_D:
		return "DEBUG"
	case Level_I:
		return "INFO"
	case Level_W:
		return "WARN"
	case Level_E:
		return "ERROR"
	default:
		return "UNKNOWN"
	}
}

func (c *customLogWriter) Write(p []byte) (n int, err error) {
	buf := &bytes.Buffer{}

	traceID := GetTraceID()

	_, file, line, ok := runtime.Caller(4)
	if !ok {
		file = "unknown"
		line = 0
	}
	file = file[strings.LastIndex(file, "/")+1:]

	fmt.Fprintf(buf, "[%s] ", time.Now().Format("2006-01-02 15:04:05.999"))
	fmt.Fprintf(buf, "%s", c.prefix)
	if traceID != "" {
		fmt.Fprintf(buf, " (%s) ", traceID)
	}
	fmt.Fprintf(buf, "[%s:%d] %s", file, line, p)

	return c.output.Write(buf.Bytes())
}

func GetLogWriter() io.Writer {
	return DefaultWriter
}