From 0d580ef9b62628b03c0bd67321d6b1dbf2822ef8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rodrigo=20Arg=C3=BCello?= Date: Tue, 22 Oct 2024 13:09:46 +0200 Subject: [PATCH] contrib/log/slog: fix WithAttrs and WithGroup implementation (#2857) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Dario Castañé --- contrib/log/slog/slog.go | 73 ++++++++++++++++--- contrib/log/slog/slog_test.go | 131 +++++++++++++++++++++++++++++----- 2 files changed, 177 insertions(+), 27 deletions(-) diff --git a/contrib/log/slog/slog.go b/contrib/log/slog/slog.go index 5cb0fcbdc2..a29866f294 100644 --- a/contrib/log/slog/slog.go +++ b/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) } diff --git a/contrib/log/slog/slog_test.go b/contrib/log/slog/slog_test.go index 86254dde88..0d25702cb2 100644 --- a/contrib/log/slog/slog_test.go +++ b/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