[WIP] support DD_LOGS_INJECTION in logrus and slog contribs, in dd-trace-go and orchestrion

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)

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,25 +114,25 @@ 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) {
 	createLogger := func(w io.Writer) *slog.Logger {
 		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()

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,13 +27,28 @@ 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}
 }
 
 // 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

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")
+}