kra-oa/pkg/logging/zap.go

package logging

import (
	"context"
	"log/slog"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/go-kratos/kratos/contrib/otel/v3/tracing"
	kratoslog "github.com/go-kratos/kratos/v3/log"
	"go.uber.org/zap"
	"go.uber.org/zap/exp/zapslog"
	"go.uber.org/zap/zapcore"
)

type Options struct {
	Level, Format, EncodeLevel, Prefix, StacktraceKey string
	LogInConsole, ShowLine                            bool
	RetentionDay                                      int
	FileOnlyModules                                   []string
}

type handlerOperation struct {
	attrs []slog.Attr
	group string
}

type reloadableHandlerState struct {
	mu      sync.RWMutex
	handler slog.Handler
	cleanup func()
}

type reloadableHandler struct {
	state *reloadableHandlerState
	ops   []handlerOperation
}

func (h *reloadableHandler) resolved() slog.Handler {
	current := h.state.handler
	for _, operation := range h.ops {
		if operation.group != "" {
			current = current.WithGroup(operation.group)
		} else {
			current = current.WithAttrs(operation.attrs)
		}
	}
	return current
}

func (h *reloadableHandler) Enabled(ctx context.Context, level slog.Level) bool {
	h.state.mu.RLock()
	defer h.state.mu.RUnlock()
	return h.resolved().Enabled(ctx, level)
}

func (h *reloadableHandler) Handle(ctx context.Context, record slog.Record) error {
	h.state.mu.RLock()
	defer h.state.mu.RUnlock()
	return h.resolved().Handle(ctx, record)
}

func (h *reloadableHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	next := append([]handlerOperation(nil), h.ops...)
	next = append(next, handlerOperation{attrs: append([]slog.Attr(nil), attrs...)})
	return &reloadableHandler{state: h.state, ops: next}
}

func (h *reloadableHandler) WithGroup(name string) slog.Handler {
	next := append([]handlerOperation(nil), h.ops...)
	next = append(next, handlerOperation{group: name})
	return &reloadableHandler{state: h.state, ops: next}
}

type ReloadableLogger struct {
	state    *reloadableHandlerState
	filename string
}

func (l *ReloadableLogger) Reload(root string, options Options) {
	handler, cleanup := newZapHandler(root, l.filename, options)
	l.state.mu.Lock()
	previous := l.state.cleanup
	l.state.handler = handler
	l.state.cleanup = cleanup
	l.state.mu.Unlock()
	if previous != nil {
		previous()
	}
}

func (l *ReloadableLogger) Close() {
	l.state.mu.Lock()
	cleanup := l.state.cleanup
	l.state.cleanup = nil
	l.state.mu.Unlock()
	if cleanup != nil {
		cleanup()
	}
}

type moduleFilterCore struct {
	zapcore.Core
	fileOnly map[string]struct{}
	mod      string
}

func (c *moduleFilterCore) With(fields []zapcore.Field) zapcore.Core {
	mod := c.mod
	if value := moduleField(fields); value != "" {
		mod = value
	}
	return &moduleFilterCore{Core: c.Core.With(fields), fileOnly: c.fileOnly, mod: mod}
}

func (c *moduleFilterCore) Check(entry zapcore.Entry, checked *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(entry.Level) {
		return checked.AddCore(entry, c)
	}
	return checked
}

func (c *moduleFilterCore) Write(entry zapcore.Entry, fields []zapcore.Field) error {
	mod := c.mod
	if value := moduleField(fields); value != "" {
		mod = value
	}
	if _, excluded := c.fileOnly[mod]; excluded {
		return nil
	}
	return c.Core.Write(entry, fields)
}

func moduleField(fields []zapcore.Field) string {
	for _, field := range fields {
		if field.Key == "mod" {
			return field.String
		}
	}
	return ""
}

// NewZapLogger adapts a Zap core to the slog logger used by Kratos v3.
// The file layout remains compatible with the administration log viewer.
func newZapHandler(root, filename string, options Options) (slog.Handler, func()) {
	file := NewDailyWriter(root, filename, options.RetentionDay)
	encoder := zap.NewProductionEncoderConfig()
	encoder.EncodeTime = zapcore.RFC3339NanoTimeEncoder
	if options.StacktraceKey != "" {
		encoder.StacktraceKey = options.StacktraceKey
	}
	switch options.EncodeLevel {
	case "CapitalLevelEncoder":
		encoder.EncodeLevel = zapcore.CapitalLevelEncoder
	case "CapitalColorLevelEncoder":
		encoder.EncodeLevel = zapcore.CapitalColorLevelEncoder
	case "LowercaseColorLevelEncoder":
		encoder.EncodeLevel = zapcore.LowercaseColorLevelEncoder
	default:
		encoder.EncodeLevel = zapcore.LowercaseLevelEncoder
	}
	var outputEncoder zapcore.Encoder = zapcore.NewJSONEncoder(encoder)
	if options.Format != "" && options.Format != "json" {
		if options.Prefix != "" {
			encoder.EncodeTime = func(value time.Time, output zapcore.PrimitiveArrayEncoder) {
				output.AppendString(options.Prefix + value.Format("2006-01-02 15:04:05.000"))
			}
		}
		outputEncoder = zapcore.NewConsoleEncoder(encoder)
	}
	level := zap.InfoLevel
	if parsed := level.Set(strings.ToLower(options.Level)); parsed != nil {
		level = zap.DebugLevel
	}
	levelEnabler := zap.NewAtomicLevelAt(level)
	fileCore := zapcore.NewCore(
		outputEncoder.Clone(),
		zapcore.AddSync(file),
		levelEnabler,
	)
	core := zapcore.Core(fileCore)
	if options.LogInConsole {
		fileOnly := make(map[string]struct{}, len(options.FileOnlyModules))
		for _, module := range options.FileOnlyModules {
			fileOnly[module] = struct{}{}
		}
		consoleCore := zapcore.NewCore(outputEncoder.Clone(), zapcore.AddSync(os.Stdout), levelEnabler)
		core = zapcore.NewTee(fileCore, &moduleFilterCore{Core: consoleCore, fileOnly: fileOnly})
	}
	zapLogger := zap.New(core)
	handlerOptions := []zapslog.HandlerOption{zapslog.AddStacktraceAt(slog.LevelError)}
	if options.ShowLine {
		handlerOptions = append(handlerOptions, zapslog.WithCaller(true))
	}
	handler := zapslog.NewHandler(zapLogger.Core(), handlerOptions...)
	cleanup := func() {
		_ = zapLogger.Sync()
		_ = file.Close()
	}
	return handler, cleanup
}

// NewReloadableZapLogger keeps the slog/Kratos adapter stable while replacing
// the underlying Zap core when the runtime configuration changes.
func NewReloadableZapLogger(root, filename string, options Options, attrs ...any) (*slog.Logger, *ReloadableLogger) {
	handler, cleanup := newZapHandler(root, filename, options)
	state := &reloadableHandlerState{handler: handler, cleanup: cleanup}
	control := &ReloadableLogger{state: state, filename: filename}
	logger := kratoslog.NewLogger(&reloadableHandler{state: state}, kratoslog.WithExtractor(tracing.TraceAttrs)).With(attrs...)
	return logger, control
}

func NewZapLogger(root, filename string, options Options, attrs ...any) (*slog.Logger, func()) {
	logger, control := NewReloadableZapLogger(root, filename, options, attrs...)
	return logger, control.Close
}