From b5011198c0bb2c4e90da8e48749d6c04607f8207 Mon Sep 17 00:00:00 2001 From: Mikayla Toffler Date: Tue, 28 Jan 2025 16:04:48 -0500 Subject: [PATCH] support DD_LOGS_INJECTION=false in logrus and slog contribs --- contrib/log/slog/slog.go | 52 +++++++++++++++++--------- contrib/log/slog/slog_test.go | 32 ++++++++++++---- contrib/sirupsen/logrus/logrus.go | 18 +++++++++ contrib/sirupsen/logrus/logrus_test.go | 19 ++++++++++ 4 files changed, 95 insertions(+), 26 deletions(-) diff --git a/contrib/log/slog/slog.go b/contrib/log/slog/slog.go index ccb9800ee4..afe5bf174c 100644 --- a/contrib/log/slog/slog.go +++ b/contrib/log/slog/slog.go @@ -10,6 +10,7 @@ import ( "context" "io" "log/slog" + "os" "strconv" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" @@ -18,6 +19,9 @@ import ( ) const componentName = "log/slog" +const logInjection = "DD_LOGS_INJECTION" + +var cfg = newConfig() func init() { telemetry.LoadIntegration(componentName) @@ -31,6 +35,16 @@ type group struct { attrs []slog.Attr } +type config struct { + enabled bool +} + +func newConfig() *config { + return &config{ + enabled: os.Getenv(logInjection) != "false", + } +} + // 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 { @@ -56,25 +70,27 @@ func (h *handler) Enabled(ctx context.Context, level slog.Level) bool { 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 && span.Context().TraceID() != 0 { - 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), + if cfg.enabled { + // 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 && span.Context().TraceID() != 0 { + 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) } - reqHandler = reqHandler.WithAttrs(attrs) - } - for _, g := range h.groups { - reqHandler = reqHandler.WithGroup(g.name) - if len(g.attrs) > 0 { - reqHandler = reqHandler.WithAttrs(g.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) diff --git a/contrib/log/slog/slog_test.go b/contrib/log/slog/slog_test.go index 36a5960db4..fbadb867c7 100644 --- a/contrib/log/slog/slog_test.go +++ b/contrib/log/slog/slog_test.go @@ -95,13 +95,16 @@ func testLogger(t *testing.T, createLogger func(b io.Writer) *slog.Logger, asser assertLogEntry(t, logs[1], "this is an error log with tracing information", "ERROR", span, assertExtra) } -func testLoggerNoTrace(t *testing.T, createLogger func(b io.Writer) *slog.Logger) { +func testLoggerTraceDisabled(t *testing.T, createLogger func(b io.Writer) *slog.Logger) { tracer.Start( tracer.WithTraceEnabled(false), tracer.WithLogger(internallog.DiscardLogger{}), ) defer tracer.Stop() + testLoggerNoTrace(t, createLogger) +} +func testLoggerNoTrace(t *testing.T, createLogger func(b io.Writer) *slog.Logger) { // create the application logger var b bytes.Buffer logger := createLogger(&b) @@ -111,17 +114,17 @@ func testLoggerNoTrace(t *testing.T, createLogger func(b io.Writer) *slog.Logger defer span.Finish() // log a message using the context containing span information - logger.Log(ctx, slog.LevelInfo, "this is an info log with tracing information") - logger.Log(ctx, slog.LevelError, "this is an error log with tracing information") + logger.Log(ctx, slog.LevelInfo, "this is an info log without tracing information") + logger.Log(ctx, slog.LevelError, "this is an error log without tracing information") logs := strings.Split( strings.TrimRight(b.String(), "\n"), "\n", ) - // assert log entries contain trace information + // assert log entries do not contain trace information require.Len(t, logs, 2) - assertLogEntryNoTrace(t, logs[0], "this is an info log with tracing information", "INFO") - assertLogEntryNoTrace(t, logs[1], "this is an error log with tracing information", "ERROR") + assertLogEntryNoTrace(t, logs[0], "this is an info log without tracing information", "INFO") + assertLogEntryNoTrace(t, logs[1], "this is an error log without tracing information", "ERROR") } func TestNewJSONHandler(t *testing.T) { @@ -129,7 +132,7 @@ func TestNewJSONHandler(t *testing.T) { return slog.New(NewJSONHandler(w, nil)) } testLogger(t, createLogger, nil) - testLoggerNoTrace(t, createLogger) + testLoggerTraceDisabled(t, createLogger) } func TestWrapHandler(t *testing.T) { @@ -138,7 +141,7 @@ func TestWrapHandler(t *testing.T) { return slog.New(WrapHandler(slog.NewJSONHandler(w, nil))) } testLogger(t, createLogger, nil) - testLoggerNoTrace(t, createLogger) + testLoggerTraceDisabled(t, createLogger) }) t.Run("slogtest", func(t *testing.T) { @@ -252,6 +255,19 @@ func TestRecordClone(t *testing.T) { assert.True(t, foundSentinel) } +func TestDDLogInjectionDisabled(t *testing.T) { + t.Setenv(logInjection, "false") + createLogger := func(w io.Writer) *slog.Logger { + return slog.New(NewJSONHandler(w, nil)) + } + tracer.Start() + defer tracer.Stop() + // Re-initialize to account for race condition between setting env var in the test and reading it in the contrib + cfg = newConfig() + testLoggerNoTrace(t, createLogger) + +} + func BenchmarkHandler(b *testing.B) { span, ctx := tracer.StartSpanFromContext(context.Background(), "test") defer span.Finish() diff --git a/contrib/sirupsen/logrus/logrus.go b/contrib/sirupsen/logrus/logrus.go index d0e379d796..2281ae2aea 100644 --- a/contrib/sirupsen/logrus/logrus.go +++ b/contrib/sirupsen/logrus/logrus.go @@ -7,6 +7,8 @@ package logrus import ( + "os" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" "gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry" @@ -15,6 +17,7 @@ import ( ) const componentName = "sirupsen/logrus" +const logInjection = "DD_LOGS_INJECTION" func init() { telemetry.LoadIntegration(componentName) @@ -24,6 +27,18 @@ func init() { // DDContextLogHook ensures that any span in the log context is correlated to log output. type DDContextLogHook struct{} +type config struct { + enabled bool +} + +var cfg = newConfig() + +func newConfig() *config { + return &config{ + enabled: os.Getenv(logInjection) != "false", + } +} + // Levels implements logrus.Hook interface, this hook applies to all defined levels func (d *DDContextLogHook) Levels() []logrus.Level { return []logrus.Level{logrus.PanicLevel, logrus.FatalLevel, logrus.ErrorLevel, logrus.WarnLevel, logrus.InfoLevel, logrus.DebugLevel, logrus.TraceLevel} @@ -31,6 +46,9 @@ func (d *DDContextLogHook) Levels() []logrus.Level { // Fire implements logrus.Hook interface, attaches trace and span details found in entry context func (d *DDContextLogHook) Fire(e *logrus.Entry) error { + if !cfg.enabled { + return nil + } span, found := tracer.SpanFromContext(e.Context) if !found { return nil diff --git a/contrib/sirupsen/logrus/logrus_test.go b/contrib/sirupsen/logrus/logrus_test.go index be23324b4d..91d092505a 100644 --- a/contrib/sirupsen/logrus/logrus_test.go +++ b/contrib/sirupsen/logrus/logrus_test.go @@ -29,3 +29,22 @@ func TestFire(t *testing.T) { assert.Equal(t, uint64(1234), e.Data["dd.trace_id"]) assert.Equal(t, uint64(1234), e.Data["dd.span_id"]) } + +func TestDDLogInjectionDisabled(t *testing.T) { + t.Setenv(logInjection, "false") + // Re-initialize to account for race condition between setting env var in the test and reading it in the contrib + cfg = newConfig() + + tracer.Start() + defer tracer.Stop() + _, sctx := tracer.StartSpanFromContext(context.Background(), "testSpan", tracer.WithSpanID(1234)) + + hook := &DDContextLogHook{} + e := logrus.NewEntry(logrus.New()) + e.Context = sctx + err := hook.Fire(e) + + assert.NoError(t, err) + assert.NotContains(t, e.Data, "dd.trace_id") + assert.NotContains(t, e.Data, "dd.span_id") +}