journal/journslog/journslog.go

/*
Package journslog implements an [slog.Handler] which writes output to the
[systemd journal].

[systemd journal]: https://www.freedesktop.org/software/systemd/man/latest/systemd-journald.service.html
*/
package journslog

import (
	"context"
	"log/slog"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"time"

	"src.lwithers.me.uk/go/journal"
)

// JournalHandler is an [slog.Handler] which writes output to the
// [systemd journal].
//
// Note that the handler discards the Time field from [slog.Record] by default,
// as there is no standardised way to override the timestamp recorded by the
// systemd-journal upon reception of a message. It is possible to emit a field
// using [WithAttrsFromTime].
//
// [systemd journal]: https://www.freedesktop.org/software/systemd/man/latest/systemd-journald.service.html
type JournalHandler struct {
	// addSource causes CODE_FILE, CODE_LINE and CODE_FUNC attrs to be added
	// to each entry.
	addSource bool

	// minLevel is the minimum level of messages we emit. Lower levels are
	// not emitted. Note it can be overridden by levelFromContext.
	minLevel slog.Leveler

	// replaceAttr is an optional function (may be nil) which the user can
	// use to map (resolved) [slog.Attr] values into other values, including
	// empty to skip.
	replaceAttr func(groups []string, a slog.Attr) slog.Attr

	// levelMapper turns [slog.Level] values into [journal.Priority] values.
	levelMapper LevelMapper

	// keyMapper turns [slog.Attr] key strings into [journal.AttrKey] values.
	keyMapper KeyMapper

	// conn is the underlying journal connection to use.
	conn *journal.Conn

	// levelFromContext allows overriding the minimum log level for a logger
	// based on the [context.Context] value. It may be nil.
	levelFromContext func(context.Context) (slog.Level, bool)

	// attrsFromContext may supply additional attributes to the logger from
	// the [context.Context] value. It may be nil.
	attrsFromContext func(context.Context) []slog.Attr

	// attrsFromTime may return zero or more attributes to represent the
	// Time from an [slog.Record]. It may be nil.
	attrsFromTime func(time.Time) []slog.Attr

	// groups are the ordered group names added by WithGroup. We need to
	// maintain the full list, not just a pre-rendered prefix string,
	// because of the API of replaceAttr.
	groups []string

	// group is the pre-rendered prefix string corresponding to groups.
	group groupPrefix

	// jattrs are the pre-rendered attributes added by WithAttrs.
	jattrs []journal.Attr
}

// NewHandler returns a new [JournalHandler]. It returns an error if a
// connection to the journal socket could not be established.
func NewHandler(opts ...HandlerOption) (*JournalHandler, error) {
	jh := &JournalHandler{
		minLevel:    slog.LevelInfo,
		levelMapper: DefaultLevelMapper(),
		keyMapper:   DefaultKeyMapper(),
	}
	for _, opt := range opts {
		opt(jh)
	}

	// For fields which require allocation of resources, initialise them
	// after processing options in case the user already supplied a non-nil
	// value. In particular this means we don't cause the global connection
	// to the journal to be opened if the user supplied something more
	// specific.
	if jh.levelMapper == nil {
		jh.levelMapper = DefaultLevelMapper()
	}
	if jh.keyMapper == nil {
		jh.keyMapper = DefaultKeyMapper()
	}
	if jh.conn == nil {
		var err error
		if jh.conn, err = journal.GetGlobalConn(); err != nil {
			return nil, err
		}
	}

	return jh, nil
}

// WithAttrs returns an [slog.Handler] containing some fixed attributes, which
// will be emitted on all future calls to Handle. The attribute values are
// resolved within this function call.
func (jh *JournalHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return jh
	}

	jattrs := make([]journal.Attr, 0, len(attrs))
	for _, attr := range attrs {
		jattrs = jh.renderAttribute(jattrs, jh.groups, jh.group, attr)
	}

	if len(jattrs) == 0 {
		return jh
	}

	nh := new(JournalHandler)
	*nh = *jh
	nh.jattrs = slices.Concat(jh.jattrs, jattrs)
	return nh
}

// WithGroup returns an [slog.Handler] where are future attributes are assigned
// to the given group.
func (jh *JournalHandler) WithGroup(group string) slog.Handler {
	nh := new(JournalHandler)
	*nh = *jh
	nh.groups = append(slices.Clone(jh.groups), group)
	nh.group = jh.group.Append(group)
	return nh
}

// Enabled returns true if the given log level is enabled, false otherwise.
func (jh *JournalHandler) Enabled(ctx context.Context, level slog.Level) bool {
	var (
		minLevel slog.Level
		ok       bool
	)
	// avoid call to jh.minLevel.Level() if we have a more specific value
	// from the context
	if jh.levelFromContext != nil {
		minLevel, ok = jh.levelFromContext(ctx)
	}
	if !ok {
		minLevel = jh.minLevel.Level()
	}
	return level >= minLevel
}

// Handle a log record, writing it to the journal similar to [journal.Entry].
func (jh *JournalHandler) Handle(ctx context.Context, rec slog.Record) error {
	pri := jh.levelMapper.LevelMap(rec.Level)

	var addAttrs []slog.Attr
	if jh.attrsFromTime != nil && !rec.Time.IsZero() {
		addAttrs = attrPoolGet()
		defer attrPoolPut(addAttrs)
		addAttrs = append(addAttrs, jh.attrsFromTime(rec.Time)...)
	}
	if jh.attrsFromContext != nil {
		if addAttrs == nil {
			addAttrs = attrPoolGet()
			defer attrPoolPut(addAttrs)
		}
		addAttrs = append(addAttrs, jh.attrsFromContext(ctx)...)
	}

	// 2 for journal.Conn.Entry (MESSAGE and PRIORITY),
	// 3 for SRC_* (optional but we always allocate space)
	numAttrsEstimate := 2 + len(jh.conn.Common) + 3 + len(addAttrs) + len(jh.jattrs) + rec.NumAttrs()
	jattrs := slices.Grow(jattrPoolGet(), numAttrsEstimate)
	defer jattrPoolPut(jattrs)

	// add SRC_* if requested
	if jh.addSource && rec.PC != 0 {
		jattrs = append(jattrs, jh.sourceAttrs(rec.PC)...)
	}

	// add time / context attributes, if any
	for i := range addAttrs {
		jattrs = jh.renderAttribute(jattrs, nil, "", addAttrs[i])
	}

	// add pre-rendered attributes from WithAttrs
	jattrs = append(jattrs, jh.jattrs...)

	// map [slog.Attr] values to [journal.Attr] values
	rec.Attrs(func(a slog.Attr) bool {
		jattrs = jh.renderAttribute(jattrs, jh.groups, jh.group, a)
		return true
	})

	return jh.conn.EntryErr(pri, rec.Message, jattrs)
}

// groupPrefix enables efficient building of attribute keys prepended by group
// names, separated by '.'. Invariant: always n × (group-name '.'), n ≥ 0.
type groupPrefix string

func (gp groupPrefix) Prepend(key string) string {
	if gp == "" { // avoid allocation in frequent path
		return key
	}
	return string(gp) + key
}

func (gp groupPrefix) Append(group string) groupPrefix {
	var b strings.Builder
	b.WriteString(string(gp))
	b.WriteString(group)
	b.WriteRune('.')
	return groupPrefix(b.String())
}

// renderAttribute turns an [slog.Attr] into some number of resolved,
// fully-qualified [journal.Attr] values. It calls replaceAttr if defined.
// New attributes are appended to the slice jattrs, and the new slice header
// is returned.
//
// If the attribute (after replacement) is a zero [slog.Attr] value, or it is
// of kind [slog.KindGroup] with zero child elements, no change is made to the
// jattrs slice.
//
// If the attribute (after replacement) is of kind [slog.Group], the function
// recurses, potentially adding multiple elements to the jattrs slice.
// list:[]journal.Attr{...}}.
//
// Otherwise, the function appends a single element to the jattrs slice.
func (jh *JournalHandler) renderAttribute(jattrs []journal.Attr, groups []string, gp groupPrefix, attr slog.Attr) []journal.Attr {
	val := attr.Value.Resolve()

	if jh.replaceAttr != nil && val.Kind() != slog.KindGroup {
		attr = jh.replaceAttr(groups, attr)
		val = attr.Value.Resolve()
	}

	if attr.Equal(slog.Attr{}) {
		return jattrs
	}

	// handle slog.KindGroup by recursing
	if val.Kind() == slog.KindGroup {
		groups = append(groups, attr.Key)
		gp = gp.Append(attr.Key)
		for _, attr := range val.Group() {
			jattrs = jh.renderAttribute(jattrs, groups, gp, attr)
		}
		return jattrs
	}

	key := gp.Prepend(attr.Key)
	return append(jattrs, journal.Attr{
		Key:   jh.keyMapper.KeyMap(key),
		Value: []byte(val.String()),
	})
}

var (
	srcAttrCodeFile = journal.MustAttrKey("CODE_FILE")
	srcAttrCodeLine = journal.MustAttrKey("CODE_LINE")
	srcAttrCodeFunc = journal.MustAttrKey("CODE_FUNC")
)

// sourceAttrs returns the attributes we emit for source location.
func (jh *JournalHandler) sourceAttrs(PC uintptr) []journal.Attr {
	fs := runtime.CallersFrames([]uintptr{PC})
	f, _ := fs.Next()

	return []journal.Attr{
		journal.Attr{
			Key:   srcAttrCodeFile,
			Value: []byte(f.File),
		},
		journal.Attr{
			Key:   srcAttrCodeLine,
			Value: []byte(strconv.Itoa(f.Line)),
		},
		journal.Attr{
			Key:   srcAttrCodeFunc,
			Value: []byte(f.Function),
		},
	}
}