[NH-86979] Add slog.Handler to associate trace context with logs (#129)

swo/log.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"github.com/solarwinds/apm-go/internal/state"
 	"go.opentelemetry.io/otel/trace"
+	"log/slog"
 )
 
 type LoggableTraceContext struct {
@@ -59,3 +60,47 @@ func LoggableTraceFromSpanContext(ctx trace.SpanContext) LoggableTraceContext {
 		ServiceName: state.GetServiceName(),
 	}
 }
+
+// LogHandler is a custom slog handler that adds a trace ID from the trace context to each log entry
+type LogHandler struct {
+	wrapped slog.Handler
+}
+
+var _ slog.Handler = &LogHandler{}
+
+// NewLogHandler creates a new LogHandler
+func NewLogHandler(wrapped slog.Handler) *LogHandler {
+	return &LogHandler{wrapped: wrapped}
+}
+
+// Enabled calls slog Enabled
+func (h *LogHandler) Enabled(ctx context.Context, level slog.Level) bool {
+	return h.wrapped.Enabled(ctx, level)
+}
+
+// Handle adds trace context to the record, in the format that allows SWO to
+// associate log lines with traces
+func (h *LogHandler) Handle(ctx context.Context, record slog.Record) error {
+	traceContext := LoggableTrace(ctx)
+	if traceContext.IsValid() {
+		record.AddAttrs(
+			slog.String("trace_id", traceContext.TraceID.String()),
+			slog.String("span_id", traceContext.SpanID.String()),
+			slog.String("trace_flags", traceContext.TraceFlags.String()),
+		)
+	}
+	record.AddAttrs(
+		slog.String("resource.service.name", traceContext.ServiceName),
+	)
+	return h.wrapped.Handle(ctx, record)
+}
+
+// WithAttrs calls slog WithAttrs
+func (h *LogHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	return &LogHandler{wrapped: h.wrapped.WithAttrs(attrs)}
+}
+
+// WithGroup calls slog WithGroup
+func (h *LogHandler) WithGroup(name string) slog.Handler {
+	return &LogHandler{wrapped: h.wrapped.WithGroup(name)}
+}

swo/log_test.go

@@ -15,11 +15,15 @@
 package swo
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"github.com/solarwinds/apm-go/internal/state"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/trace"
+	"log/slog"
 	"testing"
+	"time"
 )
 
 func TestLoggableTraceIDFromContext(t *testing.T) {
@@ -75,3 +79,76 @@ func TestLoggableTraceIDFromContext(t *testing.T) {
 		"trace_id=33000000000000000000000000000000 span_id=aa00000000000000 trace_flags=01 resource.service.name=test-service",
 		lt.String())
 }
+
+func TestHandle(t *testing.T) {
+	prevServiceName := state.GetServiceName()
+	defer state.SetServiceName(prevServiceName)
+	serviceName := "test-service"
+	state.SetServiceName(serviceName)
+
+	type logLine struct {
+		Time        time.Time `json:"time"`
+		Level       string    `json:"level"`
+		Msg         string    `json:"msg"`
+		ServiceName string    `json:"resource.service.name"`
+		TraceID     string    `json:"trace_id"`
+		SpanID      string    `json:"span_id"`
+		TraceFlags  string    `json:"trace_flags"`
+	}
+
+	tests := []struct {
+		name     string
+		traceCtx LoggableTraceContext
+	}{
+		{"info level with trace", LoggableTraceContext{
+			TraceID:    trace.TraceID{0x11},
+			SpanID:     trace.SpanID{0x22},
+			TraceFlags: trace.FlagsSampled,
+		}},
+		{"empty trace", LoggableTraceContext{}},
+	}
+
+	now := time.Now().Truncate(time.Second).UTC()
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// arrange
+			writer := bytes.NewBuffer(make([]byte, 0))
+			baseHandler := slog.NewJSONHandler(writer, &slog.HandlerOptions{})
+			handler := NewLogHandler(baseHandler)
+
+			spanContext := trace.NewSpanContext(trace.SpanContextConfig{
+				TraceID:    tt.traceCtx.TraceID,
+				SpanID:     tt.traceCtx.SpanID,
+				TraceFlags: tt.traceCtx.TraceFlags,
+			})
+			ctx := trace.ContextWithSpanContext(context.Background(), spanContext)
+			record := slog.NewRecord(now, slog.LevelInfo, "test message", 0)
+
+			// act
+			err := handler.Handle(ctx, record)
+			require.NoError(t, err)
+
+			// assert
+			logOutput := writer.String()
+
+			var line logLine
+			err = json.Unmarshal([]byte(logOutput), &line)
+			require.NoError(t, err)
+
+			if tt.traceCtx.IsValid() {
+				require.Equal(t, tt.traceCtx.TraceID.String(), line.TraceID)
+				require.Equal(t, tt.traceCtx.SpanID.String(), line.SpanID)
+				require.Equal(t, tt.traceCtx.TraceFlags.String(), line.TraceFlags)
+			} else {
+				require.Empty(t, line.TraceID)
+				require.Empty(t, line.SpanID)
+				require.Empty(t, line.TraceFlags)
+			}
+			require.Equal(t, "test message", line.Msg)
+			require.Equal(t, serviceName, line.ServiceName)
+			require.Equal(t, now, line.Time)
+
+		})
+	}
+}