Add Baggage to logback MDC controlled by flag (#7892)

The intention here is to allow users of the java agent to set a VM flag
in order to be able to add values in the current Baggage context to MDC
for logback. It seemed unwise to turn this on by default - if the
application is configured to print all MDC contents (as it often the
case with JSON output) then baggage would be logged out by default which
may either bloat the logs or result in sensitive data being exposed
unitentionally.

Addresses #1391 and #6708

Note that this is my first contribution to this repo, I've done my best
to follow the existing approaches to things like testing but would
appreciate any feedback.

---------

Co-authored-by: Mateusz Rzeszutek <mrzeszutek@splunk.com>

instrumentation/logback/logback-mdc-1.0/javaagent/build.gradle.kts

@@ -10,6 +10,42 @@ muzzle {
   }
 }
 
+testing {
+  suites {
+    val addBaggageTest by registering(JvmTestSuite::class) {
+      targets {
+        all {
+          testTask.configure {
+            jvmArgs("-Dotel.instrumentation.logback-mdc.add-baggage=true")
+          }
+        }
+      }
+    }
+
+    withType(JvmTestSuite::class) {
+      dependencies {
+        if (findProperty("testLatestDeps") as Boolean) {
+          implementation("ch.qos.logback:logback-classic:+")
+        } else {
+          implementation("ch.qos.logback:logback-classic") {
+            version {
+              strictly("1.0.0")
+            }
+          }
+          implementation("org.slf4j:slf4j-api") {
+            version {
+              strictly("1.6.4")
+            }
+          }
+        }
+
+        implementation(project(":instrumentation:logback:logback-mdc-1.0:testing"))
+        implementation(project(":instrumentation:logback:logback-mdc-1.0:javaagent"))
+      }
+    }
+  }
+}
+
 dependencies {
   implementation(project(":instrumentation:logback:logback-mdc-1.0:library"))
 
@@ -24,21 +60,10 @@ dependencies {
       strictly("1.6.4")
     }
   }
+}
 
-  if (findProperty("testLatestDeps") as Boolean) {
-    testImplementation("ch.qos.logback:logback-classic:+")
-  } else {
-    testImplementation("ch.qos.logback:logback-classic") {
-      version {
-        strictly("1.0.0")
-      }
-    }
-    testImplementation("org.slf4j:slf4j-api") {
-      version {
-        strictly("1.6.4")
-      }
-    }
+tasks {
+  named("check") {
+    dependsOn(testing.suites)
   }
-
-  testImplementation(project(":instrumentation:logback:logback-mdc-1.0:testing"))
 }

instrumentation/logback/logback-mdc-1.0/javaagent/src/addBaggageTest/groovy/io/opentelemetry/javaagent/instrumentation/logback/v1_0/LogbackWithBaggageTest.groovy

@@ -0,0 +1,12 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.logback.v1_0
+
+import io.opentelemetry.instrumentation.logback.mdc.v1_0.AbstractLogbackWithBaggageTest
+import io.opentelemetry.instrumentation.test.AgentTestTrait
+
+class LogbackWithBaggageTest extends AbstractLogbackWithBaggageTest implements AgentTestTrait {
+}

instrumentation/logback/logback-mdc-1.0/javaagent/src/addBaggageTest/resources/logback.xml

@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<configuration>
+  <appender name="LIST" class="ch.qos.logback.core.read.ListAppender" />
+
+  <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+    <encoder>
+      <pattern>%coloredLevel %logger{15} - %message%n%xException{10}</pattern>
+    </encoder>
+  </appender>
+
+  <logger name="test">
+    <level value="info" />
+    <appender-ref ref="LIST" />
+  </logger>
+
+  <root level="debug">
+    <appender-ref ref="STDOUT" />
+  </root>
+</configuration>

instrumentation/logback/logback-mdc-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/logback/mdc/v1_0/LogbackSingletons.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.logback.mdc.v1_0;
+
+import io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil;
+
+public final class LogbackSingletons {
+  private static final boolean ADD_BAGGAGE =
+      ConfigPropertiesUtil.getBoolean("otel.instrumentation.logback-mdc.add-baggage", false);
+
+  public static boolean addBaggage() {
+    return ADD_BAGGAGE;
+  }
+
+  private LogbackSingletons() {}
+}

instrumentation/logback/logback-mdc-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/logback/mdc/v1_0/LoggingEventInstrumentation.java

@@ -17,6 +17,8 @@ import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import ch.qos.logback.classic.spi.ILoggingEvent;
+import io.opentelemetry.api.baggage.Baggage;
+import io.opentelemetry.api.baggage.BaggageEntry;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
@@ -54,11 +56,11 @@ public class LoggingEventInstrumentation implements TypeInstrumentation {
 
   @SuppressWarnings("unused")
   public static class GetMdcAdvice {
-
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(
         @Advice.This ILoggingEvent event,
         @Advice.Return(typing = Typing.DYNAMIC, readOnly = false) Map<String, String> contextData) {
+
       if (contextData != null && contextData.containsKey(TRACE_ID)) {
         // Assume already instrumented event if traceId is present.
         return;
@@ -69,15 +71,28 @@ public class LoggingEventInstrumentation implements TypeInstrumentation {
         return;
       }
 
+      Map<String, String> spanContextData = new HashMap<>();
+
       SpanContext spanContext = Java8BytecodeBridge.spanFromContext(context).getSpanContext();
-      if (!spanContext.isValid()) {
-        return;
+
+      if (spanContext.isValid()) {
+        spanContextData.put(TRACE_ID, spanContext.getTraceId());
+        spanContextData.put(SPAN_ID, spanContext.getSpanId());
+        spanContextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
       }
 
-      Map<String, String> spanContextData = new HashMap<>();
-      spanContextData.put(TRACE_ID, spanContext.getTraceId());
-      spanContextData.put(SPAN_ID, spanContext.getSpanId());
-      spanContextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
+      if (LogbackSingletons.addBaggage()) {
+        Baggage baggage = Java8BytecodeBridge.baggageFromContext(context);
+
+        // using a lambda here does not play nicely with instrumentation bytecode process
+        // (Java 6 related errors are observed) so relying on for loop instead
+        for (Map.Entry<String, BaggageEntry> entry : baggage.asMap().entrySet()) {
+          spanContextData.put(
+              entry.getKey(),
+              // prefix all baggage values to avoid clashes with existing context
+              "baggage." + entry.getValue().getValue());
+        }
+      }
 
       if (contextData == null) {
         contextData = spanContextData;

instrumentation/logback/logback-mdc-1.0/library/build.gradle.kts

@@ -2,7 +2,44 @@ plugins {
   id("otel.library-instrumentation")
 }
 
+testing {
+  suites {
+    val addBaggageTest by registering(JvmTestSuite::class) {
+      targets {
+        all {
+          testTask.configure {
+            jvmArgs("-Dotel.instrumentation.logback-mdc.add-baggage=true")
+          }
+        }
+      }
+    }
+
+    withType(JvmTestSuite::class) {
+      dependencies {
+        if (findProperty("testLatestDeps") as Boolean) {
+          implementation("ch.qos.logback:logback-classic:+")
+        } else {
+          implementation("ch.qos.logback:logback-classic") {
+            version {
+              strictly("1.0.0")
+            }
+          }
+          implementation("org.slf4j:slf4j-api") {
+            version {
+              strictly("1.6.4")
+            }
+          }
+        }
+
+        implementation(project(":instrumentation:logback:logback-mdc-1.0:testing"))
+        implementation(project(":instrumentation:logback:logback-mdc-1.0:library"))
+      }
+    }
+  }
+}
+
 dependencies {
+
   // pin the version strictly to avoid overriding by dependencyManagement versions
   compileOnly("ch.qos.logback:logback-classic") {
     version {
@@ -14,21 +51,10 @@ dependencies {
       strictly("1.6.4")
     }
   }
+}
 
-  if (findProperty("testLatestDeps") as Boolean) {
-    testImplementation("ch.qos.logback:logback-classic:+")
-  } else {
-    testImplementation("ch.qos.logback:logback-classic") {
-      version {
-        strictly("1.0.0")
-      }
-    }
-    testImplementation("org.slf4j:slf4j-api") {
-      version {
-        strictly("1.6.4")
-      }
-    }
+tasks {
+  named("check") {
+    dependsOn(testing.suites)
   }
-
-  testImplementation(project(":instrumentation:logback:logback-mdc-1.0:testing"))
 }

instrumentation/logback/logback-mdc-1.0/library/src/addBaggageTest/groovy/io/opentelemetry/instrumentation/logback/mdc/v1_0/LogbackWithBaggageTest.groovy

@@ -0,0 +1,12 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.logback.mdc.v1_0
+
+
+import io.opentelemetry.instrumentation.test.LibraryTestTrait
+
+class LogbackWithBaggageTest extends AbstractLogbackWithBaggageTest implements LibraryTestTrait {
+}

instrumentation/logback/logback-mdc-1.0/library/src/addBaggageTest/resources/logback.xml

@@ -0,0 +1,31 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright The OpenTelemetry Authors
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~     http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<configuration>
+  <appender name="LIST" class="ch.qos.logback.core.read.ListAppender"/>
+
+  <appender name="OTEL"
+    class="io.opentelemetry.instrumentation.logback.mdc.v1_0.OpenTelemetryAppender">
+    <addBaggage>true</addBaggage>
+    <appender-ref ref="LIST"/>
+  </appender>
+
+  <logger name="test">
+    <level value="info"/>
+    <appender-ref ref="OTEL"/>
+  </logger>
+</configuration>

instrumentation/logback/logback-mdc-1.0/library/src/main/java/io/opentelemetry/instrumentation/logback/mdc/v1_0/OpenTelemetryAppender.java

@@ -15,8 +15,10 @@ import ch.qos.logback.core.Appender;
 import ch.qos.logback.core.UnsynchronizedAppenderBase;
 import ch.qos.logback.core.spi.AppenderAttachable;
 import ch.qos.logback.core.spi.AppenderAttachableImpl;
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.logback.mdc.v1_0.internal.UnionMap;
 import java.lang.reflect.Proxy;
 import java.util.HashMap;
@@ -25,15 +27,22 @@ import java.util.Map;
 
 public class OpenTelemetryAppender extends UnsynchronizedAppenderBase<ILoggingEvent>
     implements AppenderAttachable<ILoggingEvent> {
+  private volatile boolean addBaggage;
 
   private final AppenderAttachableImpl<ILoggingEvent> aai = new AppenderAttachableImpl<>();
 
-  public static ILoggingEvent wrapEvent(ILoggingEvent event) {
-    Span currentSpan = Span.current();
-    if (!currentSpan.getSpanContext().isValid()) {
-      return event;
-    }
+  /**
+   * When set to true this will enable addition of all baggage entries to MDC. This can be done by
+   * adding the following to the logback.xml config for this appender. {@code
+   * <addBaggage>true</addBaggage>}
+   *
+   * @param addBaggage True if baggage should be added to MDC
+   */
+  public void setAddBaggage(boolean addBaggage) {
+    this.addBaggage = addBaggage;
+  }
 
+  public ILoggingEvent wrapEvent(ILoggingEvent event) {
     Map<String, String> eventContext = event.getMDCPropertyMap();
     if (eventContext != null && eventContext.containsKey(TRACE_ID)) {
       // Assume already instrumented event if traceId is present.
@@ -41,10 +50,25 @@ public class OpenTelemetryAppender extends UnsynchronizedAppenderBase<ILoggingEv
     }
 
     Map<String, String> contextData = new HashMap<>();
-    SpanContext spanContext = currentSpan.getSpanContext();
-    contextData.put(TRACE_ID, spanContext.getTraceId());
-    contextData.put(SPAN_ID, spanContext.getSpanId());
-    contextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
+    Context context = Context.current();
+    Span currentSpan = Span.fromContext(context);
+
+    if (currentSpan.getSpanContext().isValid()) {
+      SpanContext spanContext = currentSpan.getSpanContext();
+      contextData.put(TRACE_ID, spanContext.getTraceId());
+      contextData.put(SPAN_ID, spanContext.getSpanId());
+      contextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
+    }
+
+    if (addBaggage) {
+      Baggage baggage = Baggage.fromContext(context);
+      baggage.forEach(
+          (key, value) ->
+              contextData.put(
+                  key,
+                  // prefix all baggage values to avoid clashes with existing context
+                  "baggage." + value.getValue()));
+    }
 
     if (eventContext == null) {
       eventContext = contextData;

instrumentation/logback/logback-mdc-1.0/testing/src/main/groovy/io/opentelemetry/instrumentation/logback/mdc/v1_0/AbstractLogbackTest.groovy

@@ -7,6 +7,7 @@ package io.opentelemetry.instrumentation.logback.mdc.v1_0
 
 import ch.qos.logback.classic.spi.ILoggingEvent
 import ch.qos.logback.core.read.ListAppender
+import io.opentelemetry.api.baggage.Baggage
 import io.opentelemetry.api.trace.Span
 import io.opentelemetry.instrumentation.test.InstrumentationSpecification
 import org.slf4j.Logger
@@ -39,8 +40,12 @@ abstract class AbstractLogbackTest extends InstrumentationSpecification {
 
   def "no ids when no span"() {
     when:
-    logger.info("log message 1")
-    logger.info("log message 2")
+    Baggage baggage = Baggage.empty().toBuilder().put("baggage_key", "baggage_value").build()
+
+    runWithBaggage(baggage) {
+      AbstractLogbackTest.logger.info("log message 1")
+      AbstractLogbackTest.logger.info("log message 2")
+    }
 
     def events = listAppender.list
 
@@ -50,25 +55,27 @@ abstract class AbstractLogbackTest extends InstrumentationSpecification {
     events[0].getMDCPropertyMap().get("trace_id") == null
     events[0].getMDCPropertyMap().get("span_id") == null
     events[0].getMDCPropertyMap().get("trace_flags") == null
+    events[0].getMDCPropertyMap().get("baggage_key") == (expectBaggage() ? "baggage.baggage_value" : null)
 
     events[1].message == "log message 2"
     events[1].getMDCPropertyMap().get("trace_id") == null
     events[1].getMDCPropertyMap().get("span_id") == null
     events[1].getMDCPropertyMap().get("trace_flags") == null
+    events[1].getMDCPropertyMap().get("baggage_key") == (expectBaggage() ? "baggage.baggage_value" : null)
   }
 
   def "ids when span"() {
     when:
-    Span span1 = runWithSpan("test") {
+    Baggage baggage = Baggage.empty().toBuilder().put("baggage_key", "baggage_value").build()
+
+    Span span1 = runWithSpanAndBaggage("test", baggage) {
       AbstractLogbackTest.logger.info("log message 1")
-      Span.current()
     }
 
     logger.info("log message 2")
 
-    Span span2 = runWithSpan("test 2") {
+    Span span2 = runWithSpanAndBaggage("test 2", baggage) {
       AbstractLogbackTest.logger.info("log message 3")
-      Span.current()
     }
 
     def events = listAppender.list
@@ -79,15 +86,35 @@ abstract class AbstractLogbackTest extends InstrumentationSpecification {
     events[0].getMDCPropertyMap().get("trace_id") == span1.spanContext.traceId
     events[0].getMDCPropertyMap().get("span_id") == span1.spanContext.spanId
     events[0].getMDCPropertyMap().get("trace_flags") == "01"
+    events[0].getMDCPropertyMap().get("baggage_key") == (expectBaggage() ? "baggage.baggage_value" : null)
 
     events[1].message == "log message 2"
     events[1].getMDCPropertyMap().get("trace_id") == null
     events[1].getMDCPropertyMap().get("span_id") == null
     events[1].getMDCPropertyMap().get("trace_flags") == null
+    events[1].getMDCPropertyMap().get("baggage_key") == null
 
     events[2].message == "log message 3"
     events[2].getMDCPropertyMap().get("trace_id") == span2.spanContext.traceId
     events[2].getMDCPropertyMap().get("span_id") == span2.spanContext.spanId
     events[2].getMDCPropertyMap().get("trace_flags") == "01"
+    events[2].getMDCPropertyMap().get("baggage_key") == (expectBaggage() ? "baggage.baggage_value" : null)
+  }
+
+  Span runWithSpanAndBaggage(String spanName, Baggage baggage, Closure callback) {
+    return runWithSpan(spanName) {
+      runWithBaggage(baggage, callback)
+      Span.current()
+    }
+  }
+
+  void runWithBaggage(Baggage baggage, Closure callback) {
+    try (var unusedScope = baggage.makeCurrent()) {
+      callback.call()
+    }
+  }
+
+  boolean expectBaggage() {
+    return false
   }
 }

instrumentation/logback/logback-mdc-1.0/testing/src/main/groovy/io/opentelemetry/instrumentation/logback/mdc/v1_0/AbstractLogbackWithBaggageTest.groovy

@@ -0,0 +1,13 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.logback.mdc.v1_0
+
+abstract class AbstractLogbackWithBaggageTest extends AbstractLogbackTest {
+  @Override
+  boolean expectBaggage() {
+    return true
+  }
+}

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/Java8BytecodeBridge.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.bootstrap;
 
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 
@@ -37,5 +38,9 @@ public final class Java8BytecodeBridge {
     return Span.fromContext(context);
   }
 
+  public static Baggage baggageFromContext(Context context) {
+    return Baggage.fromContext(context);
+  }
+
   private Java8BytecodeBridge() {}
 }