go-omni/logx/logx.go

// Package logx provides the application logger used by GoOmni and business
// packages. Its public API is deliberately independent of Gin so the same
// logger can be used from HTTP handlers, jobs and command line programs.
package logx

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"strings"
	"sync"
	"sync/atomic"
)

// Config controls a logger created by New or Setup.
//
// Writer is useful for tests and embedding. When it is nil, Output selects a
// standard stream. File takes precedence over Output and appends to the named
// file.
type Config struct {
	Level       string
	Format      string
	Output      string
	File        string
	Service     string
	Environment string
	AddSource   bool
	Writer      io.Writer
}

// Default returns the default production-friendly logger configuration.
func Default() Config {
	return Config{
		Level:  "info",
		Format: "json",
		Output: "stderr",
	}
}

// Logger is a structured logger. Child loggers returned by With and WithGroup
// share the root logger's lifecycle and must not be closed independently.
type Logger struct {
	logger *slog.Logger
	closer *closeState
}

type closeState struct {
	once sync.Once
	fn   func() error
	err  error
}

var (
	defaultLogger atomic.Pointer[Logger]
	setupMu       sync.Mutex
)

func init() {
	// Keep package-level logging useful without requiring every service to
	// repeat setup. Applications can replace this logger with Setup.
	defaultLogger.Store(&Logger{logger: slog.New(slog.NewJSONHandler(os.Stderr, nil))})
}

// New creates an independent logger. The caller owns its lifecycle and should
// call Close when the process is shutting down.
func New(cfg Config) (*Logger, error) {
	cfg = normalizeConfig(cfg)
	level, err := parseLevel(cfg.Level)
	if err != nil {
		return nil, err
	}

	writer, closeWriter, err := openWriter(cfg)
	if err != nil {
		return nil, err
	}

	opts := &slog.HandlerOptions{Level: level, AddSource: cfg.AddSource}
	var handler slog.Handler
	if cfg.Format == "text" {
		handler = slog.NewTextHandler(writer, opts)
	} else {
		handler = slog.NewJSONHandler(writer, opts)
	}
	logger := slog.New(handler)
	if cfg.Service != "" {
		logger = logger.With("service", cfg.Service)
	}
	if cfg.Environment != "" {
		logger = logger.With("environment", cfg.Environment)
	}

	var closer *closeState
	if closeWriter != nil {
		closer = &closeState{fn: closeWriter}
	}
	return &Logger{logger: logger, closer: closer}, nil
}

// Setup replaces the process-wide logger used by the package-level helpers.
// The returned function closes the newly configured logger. Calling Close is
// usually more convenient for application shutdown.
func Setup(cfg Config) (func() error, error) {
	logger, err := New(cfg)
	if err != nil {
		return nil, err
	}
	setupMu.Lock()
	previous := defaultLogger.Swap(logger)
	setupMu.Unlock()
	if previous != nil {
		_ = previous.Close()
	}
	return logger.Close, nil
}

// SetDefault installs an existing logger as the package-wide logger.
func SetDefault(logger *Logger) {
	if logger == nil {
		logger = &Logger{logger: slog.New(slog.NewJSONHandler(io.Discard, nil))}
	}
	setupMu.Lock()
	previous := defaultLogger.Swap(logger)
	setupMu.Unlock()
	if previous != nil && previous != logger {
		_ = previous.Close()
	}
}

// Close closes the current process-wide logger and disables further output.
func Close() error {
	discard := &Logger{logger: slog.New(slog.NewJSONHandler(io.Discard, nil))}
	setupMu.Lock()
	previous := defaultLogger.Swap(discard)
	setupMu.Unlock()
	if previous == nil {
		return nil
	}
	return previous.Close()
}

// With returns a child logger carrying stable attributes.
func With(args ...any) *Logger { return defaultLogger.Load().With(args...) }

// WithGroup returns a child logger writing attributes under a group.
func WithGroup(name string) *Logger { return defaultLogger.Load().WithGroup(name) }

func (l *Logger) With(args ...any) *Logger {
	if l == nil || l.logger == nil {
		return defaultLogger.Load().With(args...)
	}
	return &Logger{logger: l.logger.With(args...), closer: l.closer}
}

func (l *Logger) WithGroup(name string) *Logger {
	if l == nil || l.logger == nil {
		return defaultLogger.Load().WithGroup(name)
	}
	return &Logger{logger: l.logger.WithGroup(name), closer: l.closer}
}

// Close releases resources owned by the root logger. It is safe to call more
// than once and is a no-op for loggers writing to standard streams.
func (l *Logger) Close() error {
	if l == nil || l.closer == nil {
		return nil
	}
	l.closer.once.Do(func() { l.closer.err = l.closer.fn() })
	return l.closer.err
}

func (l *Logger) DebugContext(ctx context.Context, msg string, args ...any) {
	ctx = normalizeContext(ctx)
	l.logger.Log(ctx, slog.LevelDebug, msg, contextAttrs(ctx, args)...)
}

func (l *Logger) InfoContext(ctx context.Context, msg string, args ...any) {
	ctx = normalizeContext(ctx)
	l.logger.Log(ctx, slog.LevelInfo, msg, contextAttrs(ctx, args)...)
}

func (l *Logger) WarnContext(ctx context.Context, msg string, args ...any) {
	ctx = normalizeContext(ctx)
	l.logger.Log(ctx, slog.LevelWarn, msg, contextAttrs(ctx, args)...)
}

func (l *Logger) ErrorContext(ctx context.Context, msg string, args ...any) {
	ctx = normalizeContext(ctx)
	l.logger.Log(ctx, slog.LevelError, msg, contextAttrs(ctx, args)...)
}

func (l *Logger) Debug(msg string, args ...any) { l.DebugContext(context.Background(), msg, args...) }
func (l *Logger) Info(msg string, args ...any)  { l.InfoContext(context.Background(), msg, args...) }
func (l *Logger) Warn(msg string, args ...any)  { l.WarnContext(context.Background(), msg, args...) }
func (l *Logger) Error(msg string, args ...any) { l.ErrorContext(context.Background(), msg, args...) }

func DebugContext(ctx context.Context, msg string, args ...any) {
	defaultLogger.Load().DebugContext(ctx, msg, args...)
}

func InfoContext(ctx context.Context, msg string, args ...any) {
	defaultLogger.Load().InfoContext(ctx, msg, args...)
}

func WarnContext(ctx context.Context, msg string, args ...any) {
	defaultLogger.Load().WarnContext(ctx, msg, args...)
}

func ErrorContext(ctx context.Context, msg string, args ...any) {
	defaultLogger.Load().ErrorContext(ctx, msg, args...)
}

func Debug(msg string, args ...any) { DebugContext(context.Background(), msg, args...) }
func Info(msg string, args ...any)  { InfoContext(context.Background(), msg, args...) }
func Warn(msg string, args ...any)  { WarnContext(context.Background(), msg, args...) }
func Error(msg string, args ...any) { ErrorContext(context.Background(), msg, args...) }

type contextKey uint8

const (
	requestIDKey contextKey = iota + 1
	traceIDKey
)

// WithRequestID adds a request correlation ID to a context.
func WithRequestID(ctx context.Context, requestID string) context.Context {
	return context.WithValue(ctx, requestIDKey, requestID)
}

// RequestID returns the request correlation ID stored in ctx.
func RequestID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	value, _ := ctx.Value(requestIDKey).(string)
	return value
}

// WithTraceID adds a trace correlation ID to a context.
func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, traceIDKey, traceID)
}

// TraceID returns the trace correlation ID stored in ctx.
func TraceID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	value, _ := ctx.Value(traceIDKey).(string)
	return value
}

func contextAttrs(ctx context.Context, args []any) []any {
	ctx = normalizeContext(ctx)
	attrs := make([]any, 0, len(args)+4)
	if requestID := RequestID(ctx); requestID != "" {
		attrs = append(attrs, "requestId", requestID)
	}
	if traceID := TraceID(ctx); traceID != "" {
		attrs = append(attrs, "traceId", traceID)
	}
	return append(attrs, args...)
}

func normalizeContext(ctx context.Context) context.Context {
	if ctx == nil {
		return context.Background()
	}
	return ctx
}

func normalizeConfig(cfg Config) Config {
	if strings.TrimSpace(cfg.Level) == "" {
		cfg.Level = "info"
	}
	cfg.Level = strings.ToLower(strings.TrimSpace(cfg.Level))
	if strings.TrimSpace(cfg.Format) == "" {
		cfg.Format = "json"
	}
	cfg.Format = strings.ToLower(strings.TrimSpace(cfg.Format))
	if strings.TrimSpace(cfg.Output) == "" {
		cfg.Output = "stderr"
	}
	cfg.Output = strings.ToLower(strings.TrimSpace(cfg.Output))
	return cfg
}

func parseLevel(value string) (slog.Level, error) {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "debug":
		return slog.LevelDebug, nil
	case "info":
		return slog.LevelInfo, nil
	case "warn", "warning":
		return slog.LevelWarn, nil
	case "error":
		return slog.LevelError, nil
	default:
		return 0, fmt.Errorf("logx: unsupported level %q", value)
	}
}

func openWriter(cfg Config) (io.Writer, func() error, error) {
	if cfg.Writer != nil {
		return cfg.Writer, nil, nil
	}
	if cfg.File != "" {
		file, err := os.OpenFile(cfg.File, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
		if err != nil {
			return nil, nil, err
		}
		return file, file.Close, nil
	}
	switch cfg.Output {
	case "stdout":
		return os.Stdout, nil, nil
	case "stderr":
		return os.Stderr, nil, nil
	case "discard", "none":
		return io.Discard, nil, nil
	default:
		return nil, nil, errors.New("logx: output must be stdout, stderr, discard, or a file must be set")
	}
}