contrib/log/slog: fix WithAttrs and WithGroup implementation (#2857)

Co-authored-by: Dario Castañé <dario.castane@datadoghq.com>

contrib/log/slog/slog.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"io"
 	"log/slog"
+	"strconv"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
@@ -23,6 +24,13 @@ func init() {
 	tracer.MarkIntegrationImported("log/slog")
 }
 
+var _ slog.Handler = (*handler)(nil)
+
+type group struct {
+	name  string
+	attrs []slog.Attr
+}
+
 // NewJSONHandler is a convenience function that returns a *slog.JSONHandler logger enhanced with
 // tracing information.
 func NewJSONHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
@@ -31,22 +39,71 @@ func NewJSONHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
 
 // WrapHandler enhances the given logger handler attaching tracing information to logs.
 func WrapHandler(h slog.Handler) slog.Handler {
-	return &handler{h}
+	return &handler{wrapped: h}
 }
 
 type handler struct {
-	slog.Handler
+	wrapped slog.Handler
+	groups  []group
+}
+
+// Enabled calls the wrapped handler Enabled method.
+func (h *handler) Enabled(ctx context.Context, level slog.Level) bool {
+	return h.wrapped.Enabled(ctx, level)
 }
 
 // Handle handles the given Record, attaching tracing information if found.
 func (h *handler) Handle(ctx context.Context, rec slog.Record) error {
+	reqHandler := h.wrapped
+
+	// We need to ensure the trace id and span id keys are set at the root level:
+	// https://docs.datadoghq.com/tracing/other_telemetry/connect_logs_and_traces/
+	// In case the user has created group loggers, we ignore those and
+	// set them at the root level.
 	span, ok := tracer.SpanFromContext(ctx)
 	if ok {
-		rec = rec.Clone()
-		rec.Add(
-			slog.Uint64(ext.LogKeyTraceID, span.Context().TraceID()),
-			slog.Uint64(ext.LogKeySpanID, span.Context().SpanID()),
-		)
+		traceID := strconv.FormatUint(span.Context().TraceID(), 10)
+		spanID := strconv.FormatUint(span.Context().SpanID(), 10)
+
+		attrs := []slog.Attr{
+			slog.String(ext.LogKeyTraceID, traceID),
+			slog.String(ext.LogKeySpanID, spanID),
+		}
+		reqHandler = reqHandler.WithAttrs(attrs)
+	}
+	for _, g := range h.groups {
+		reqHandler = reqHandler.WithGroup(g.name)
+		if len(g.attrs) > 0 {
+			reqHandler = reqHandler.WithAttrs(g.attrs)
+		}
+	}
+	return reqHandler.Handle(ctx, rec)
+}
+
+// WithAttrs saves the provided attributes associated to the current Group.
+// If Group was not called for the logger, we just call WithAttrs for the wrapped handler.
+func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	if len(h.groups) == 0 {
+		return &handler{
+			wrapped: h.wrapped.WithAttrs(attrs),
+			groups:  h.groups,
+		}
+	}
+	groups := append([]group{}, h.groups...)
+	curGroup := groups[len(groups)-1]
+	curGroup.attrs = append(curGroup.attrs, attrs...)
+	groups[len(groups)-1] = curGroup
+
+	return &handler{
+		wrapped: h.wrapped,
+		groups:  groups,
+	}
+}
+
+// WithGroup saves the provided group to be used later in the Handle method.
+func (h *handler) WithGroup(name string) slog.Handler {
+	return &handler{
+		wrapped: h.wrapped,
+		groups:  append(h.groups, group{name: name}),
 	}
-	return h.Handler.Handle(ctx, rec)
 }

contrib/log/slog/slog_test.go

@@ -11,6 +11,7 @@ import (
 	"encoding/json"
 	"io"
 	"log/slog"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -22,29 +23,37 @@ import (
 	internallog "gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 )
 
-func assertLogEntry(t *testing.T, rawEntry, wantMsg, wantLevel string) {
+func assertLogEntry(t *testing.T, rawEntry, wantMsg, wantLevel string, span tracer.Span, assertExtra func(t *testing.T, entry map[string]interface{})) {
 	t.Helper()
 
-	var data map[string]interface{}
-	err := json.Unmarshal([]byte(rawEntry), &data)
+	t.Log(rawEntry)
+
+	var entry map[string]interface{}
+	err := json.Unmarshal([]byte(rawEntry), &entry)
 	require.NoError(t, err)
-	require.NotEmpty(t, data)
+	require.NotEmpty(t, entry)
+
+	assert.Equal(t, wantMsg, entry["msg"])
+	assert.Equal(t, wantLevel, entry["level"])
+	assert.NotEmpty(t, entry["time"])
+
+	traceID := strconv.FormatUint(span.Context().TraceID(), 10)
+	spanID := strconv.FormatUint(span.Context().SpanID(), 10)
+	assert.Equal(t, traceID, entry[ext.LogKeyTraceID], "trace id not found")
+	assert.Equal(t, spanID, entry[ext.LogKeySpanID], "span id not found")
 
-	assert.Equal(t, wantMsg, data["msg"])
-	assert.Equal(t, wantLevel, data["level"])
-	assert.NotEmpty(t, data["time"])
-	assert.NotEmpty(t, data[ext.LogKeyTraceID])
-	assert.NotEmpty(t, data[ext.LogKeySpanID])
+	if assertExtra != nil {
+		assertExtra(t, entry)
+	}
 }
 
-func testLogger(t *testing.T, createHandler func(b *bytes.Buffer) slog.Handler) {
+func testLogger(t *testing.T, createLogger func(b io.Writer) *slog.Logger, assertExtra func(t *testing.T, entry map[string]interface{})) {
 	tracer.Start(tracer.WithLogger(internallog.DiscardLogger{}))
 	defer tracer.Stop()
 
 	// create the application logger
 	var b bytes.Buffer
-	h := createHandler(&b)
-	logger := slog.New(h)
+	logger := createLogger(&b)
 
 	// start a new span
 	span, ctx := tracer.StartSpanFromContext(context.Background(), "test")
@@ -60,19 +69,89 @@ func testLogger(t *testing.T, createHandler func(b *bytes.Buffer) slog.Handler)
 	)
 	// assert log entries contain trace information
 	require.Len(t, logs, 2)
-	assertLogEntry(t, logs[0], "this is an info log with tracing information", "INFO")
-	assertLogEntry(t, logs[1], "this is an error log with tracing information", "ERROR")
+	assertLogEntry(t, logs[0], "this is an info log with tracing information", "INFO", span, assertExtra)
+	assertLogEntry(t, logs[1], "this is an error log with tracing information", "ERROR", span, assertExtra)
 }
 
 func TestNewJSONHandler(t *testing.T) {
-	testLogger(t, func(b *bytes.Buffer) slog.Handler {
-		return NewJSONHandler(b, nil)
-	})
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(NewJSONHandler(w, nil))
+		},
+		nil,
+	)
 }
 
 func TestWrapHandler(t *testing.T) {
-	testLogger(t, func(b *bytes.Buffer) slog.Handler {
-		return WrapHandler(slog.NewJSONHandler(b, nil))
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(WrapHandler(slog.NewJSONHandler(w, nil)))
+		},
+		nil,
+	)
+}
+
+func TestHandlerWithAttrs(t *testing.T) {
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(NewJSONHandler(w, nil)).
+				With("key1", "val1").
+				With(ext.LogKeyTraceID, "trace-id").
+				With(ext.LogKeySpanID, "span-id")
+		},
+		nil,
+	)
+}
+
+func TestHandlerWithGroup(t *testing.T) {
+	t.Run("simple", func(t *testing.T) {
+		testLogger(
+			t,
+			func(w io.Writer) *slog.Logger {
+				return slog.New(NewJSONHandler(w, nil)).
+					WithGroup("some-group").
+					With("key1", "val1")
+			},
+			func(t *testing.T, entry map[string]interface{}) {
+				assert.Equal(t, map[string]interface{}{
+					"key1": "val1",
+				}, entry["some-group"], "group entry not found")
+			},
+		)
+	})
+
+	t.Run("nested groups", func(t *testing.T) {
+		testLogger(
+			t,
+			func(w io.Writer) *slog.Logger {
+				return slog.New(NewJSONHandler(w, nil)).
+					With("key0", "val0").
+					WithGroup("group1").
+					With("key1", "val1").
+					WithGroup("group1"). // repeat same key again
+					With("key1", "val1").
+					WithGroup("group2").
+					With("key2", "val2").
+					With("key3", "val3")
+			},
+			func(t *testing.T, entry map[string]interface{}) {
+				groupKeys := map[string]interface{}{
+					"key1": "val1",
+					"group1": map[string]interface{}{
+						"key1": "val1",
+						"group2": map[string]interface{}{
+							"key2": "val2",
+							"key3": "val3",
+						},
+					},
+				}
+				assert.Equal(t, "val0", entry["key0"], "root level key not found")
+				assert.Equal(t, groupKeys, entry["group1"], "nested group entries not found")
+			},
+		)
 	})
 }
 
@@ -106,6 +185,20 @@ func TestRecordClone(t *testing.T) {
 	assert.True(t, foundSentinel)
 }
 
+func BenchmarkHandler(b *testing.B) {
+	span, ctx := tracer.StartSpanFromContext(context.Background(), "test")
+	defer span.Finish()
+
+	// create a logger with a bunch of nested groups and fields
+	logger := slog.New(NewJSONHandler(io.Discard, nil))
+	logger = logger.With("attr1", "val1").WithGroup("group1").With("attr2", "val2").WithGroup("group3").With("attr3", "val3")
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		logger.InfoContext(ctx, "some message")
+	}
+}
+
 // handlerGate calls a gate function before calling the underlying handler. This
 // allows simulating a concurrent modification of the record that happens after
 // Handle is called (and the record has been copied), but before the back array