common/logger/logger.go

package logger

import (
	"context"
	"log"
	"log/slog"
	"os"
	"strconv"
	"sync"

	slogtraceid "github.com/remychantenay/slog-otel"
	slogmulti "github.com/samber/slog-multi"
	"go.opentelemetry.io/contrib/bridges/otelslog"
)

var ConfigPrefix = ""

var textLogger *slog.Logger
var otlpLogger *slog.Logger
var multiLogger *slog.Logger

var setupText sync.Once  // this sets the default
var setupOtlp sync.Once  // this never sets the default
var setupMulti sync.Once // this sets the default, and will always run after the others
var mu sync.Mutex

func setupStdErrHandler() slog.Handler {
	var programLevel = new(slog.LevelVar) // Info by default

	envVar := "DEBUG"
	if len(ConfigPrefix) > 0 {
		envVar = ConfigPrefix + "_" + envVar
	}

	if opt := os.Getenv(envVar); len(opt) > 0 {
		if debug, _ := strconv.ParseBool(opt); debug {
			programLevel.Set(slog.LevelDebug)
		}
	}

	logOptions := &slog.HandlerOptions{Level: programLevel}

	if len(os.Getenv("INVOCATION_ID")) > 0 {
		// don't add timestamps when running under systemd
		log.Default().SetFlags(0)

		logOptions.ReplaceAttr = logRemoveTime
	}

	logHandler := slogtraceid.OtelHandler{
		Next: slog.NewTextHandler(os.Stderr, logOptions),
	}

	return logHandler
}

func setupOtlpLogger() *slog.Logger {
	setupOtlp.Do(func() {
		otlpLogger = slog.New(
			newLogFmtHandler(otelslog.NewHandler("common")),
		)
	})
	return otlpLogger
}

// SetupMultiLogger will setup and make default a logger that
// logs as described in Setup() as well as an OLTP logger.
// The "multi logger" is made the default the first time
// this function is called
func SetupMultiLogger() *slog.Logger {
	setupMulti.Do(func() {
		textHandler := Setup().Handler()
		otlpHandler := setupOtlpLogger().Handler()

		multiHandler := slogmulti.Fanout(
			textHandler,
			otlpHandler,
		)
		mu.Lock()
		defer mu.Unlock()
		multiLogger = slog.New(multiHandler)
		slog.SetDefault(multiLogger)
	})

	return multiLogger
}

// SetupOLTP configures and returns a logger sending logs
// via OpenTelemetry (configured via the tracing package).
//
// This was made to work with Loki + Grafana that makes it
// hard to view the log attributes in the UI, so the log
// message is formatted similarly to the text logger. The
// attributes are duplicated as OLTP attributes in the
// log messages. https://github.com/grafana/loki/issues/14788
func SetupOLTP() *slog.Logger {
	return setupOtlpLogger()
}

// Setup returns an slog.Logger configured for text formatting
// to stderr.
// OpenTelemetry trace_id and span_id's are logged as attributes
// when available.
// When the application is running under systemd timestamps are
// omitted. On first call the slog default logger is set to this
// logger as well.
//
// If SetupMultiLogger has been called Setup() will return
// the "multi logger"
func Setup() *slog.Logger {
	setupText.Do(func() {
		h := setupStdErrHandler()
		textLogger = slog.New(h)
		slog.SetDefault(textLogger)
	})

	mu.Lock()
	defer mu.Unlock()

	if multiLogger != nil {
		return multiLogger
	}
	return textLogger
}

type loggerKey struct{}

// NewContext adds the logger to the context. Use this
// to for example make a request specific logger available
// to other functions through the context
func NewContext(ctx context.Context, l *slog.Logger) context.Context {
	return context.WithValue(ctx, loggerKey{}, l)
}

// FromContext retrieves a logger from the context. If there is none,
// it returns the default logger
func FromContext(ctx context.Context) *slog.Logger {
	if l, ok := ctx.Value(loggerKey{}).(*slog.Logger); ok {
		return l
	}
	return Setup()
}

func logRemoveTime(groups []string, a slog.Attr) slog.Attr {
	// Remove time
	if a.Key == slog.TimeKey && len(groups) == 0 {
		return slog.Attr{}
	}
	return a
}