Merge pull request #1 from adamleantech/add-mdc-baggage

Add mdc baggage

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

@@ -17,8 +17,11 @@
 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.internal.ConfigPropertiesUtil;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.instrumentation.logback.mdc.v1_0.internal.UnionMap;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
@@ -54,11 +57,16 @@ public void transform(TypeTransformer transformer) {
 
   @SuppressWarnings("unused")
   public static class GetMdcAdvice {
+    // this has to be public otherwise the Advice cannot access it
+    // making this non-final helps greatly with testing
+    public static boolean addBaggage =
+        ConfigPropertiesUtil.getBoolean("otel.instrumentation.logback.mdc.add-baggage", false);
 
     @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;
@@ -79,6 +87,16 @@ public static void onExit(
       spanContextData.put(SPAN_ID, spanContext.getSpanId());
       spanContextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
 
+      if (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(), entry.getValue().getValue());
+        }
+      }
+
       if (contextData == null) {
         contextData = spanContextData;
       } else {

instrumentation/logback/logback-mdc-1.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/logback/v1_0/LogbackTest.groovy

@@ -7,6 +7,11 @@ package io.opentelemetry.javaagent.instrumentation.logback.v1_0
 
 import io.opentelemetry.instrumentation.logback.mdc.v1_0.AbstractLogbackTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.javaagent.instrumentation.logback.mdc.v1_0.LoggingEventInstrumentation
 
 class LogbackTest extends AbstractLogbackTest implements AgentTestTrait {
+  @Override
+  void setBaggageFlag() {
+    LoggingEventInstrumentation.GetMdcAdvice.addBaggage = false
+  }
 }

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

@@ -0,0 +1,18 @@
+/*
+ * 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
+import io.opentelemetry.javaagent.instrumentation.logback.mdc.v1_0.LoggingEventInstrumentation
+
+class LogbackWithBaggageTest extends AbstractLogbackWithBaggageTest implements AgentTestTrait {
+
+  @Override
+  void setBaggageFlag() {
+    LoggingEventInstrumentation.GetMdcAdvice.addBaggage = true
+  }
+}

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

@@ -14,15 +14,24 @@
 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.instrumentation.api.internal.ConfigPropertiesUtil;
 import io.opentelemetry.instrumentation.logback.mdc.v1_0.internal.UnionMap;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
 
 public class OpenTelemetryAppender extends UnsynchronizedAppenderBase<ILoggingEvent>
     implements AppenderAttachable<ILoggingEvent> {
+  private static boolean addBaggage =
+      ConfigPropertiesUtil.getBoolean("otel.instrumentation.logback.mdc.add-baggage", false);
+
+  // to aid testing
+  public static void setAddBaggage(boolean addBaggage) {
+    OpenTelemetryAppender.addBaggage = addBaggage;
+  }
 
   private final AppenderAttachableImpl<ILoggingEvent> aai = new AppenderAttachableImpl<>();
 
@@ -44,6 +53,11 @@ public static ILoggingEvent wrapEvent(ILoggingEvent event) {
     contextData.put(SPAN_ID, spanContext.getSpanId());
     contextData.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
 
+    if (addBaggage) {
+      Baggage baggage = Baggage.current();
+      baggage.forEach((key, value) -> contextData.put(key, value.getValue()));
+    }
+
     if (eventContext == null) {
       eventContext = contextData;
     } else {

instrumentation/logback/logback-mdc-1.0/library/src/test/groovy/io/opentelemetry/instrumentation/logback/mdc/v1_0/LogbackTest.groovy

@@ -8,4 +8,8 @@ package io.opentelemetry.instrumentation.logback.mdc.v1_0
 import io.opentelemetry.instrumentation.test.LibraryTestTrait
 
 class LogbackTest extends AbstractLogbackTest implements LibraryTestTrait {
+  @Override
+  void setBaggageFlag() {
+    OpenTelemetryAppender.addBaggage = false
+  }
 }

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

@@ -0,0 +1,16 @@
+/*
+ * 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 {
+  @Override
+  void setBaggageFlag() {
+    OpenTelemetryAppender.addBaggage = true
+  }
+}

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
@@ -35,6 +36,7 @@ abstract class AbstractLogbackTest extends InstrumentationSpecification {
 
   def setup() {
     listAppender.list.clear()
+    setBaggageFlag()
   }
 
   def "no ids when no span"() {
@@ -59,16 +61,16 @@ abstract class AbstractLogbackTest extends InstrumentationSpecification {
 
   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 +81,34 @@ 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_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_value" : null)
+  }
+
+  Span runWithSpanAndBaggage(String spanName, Baggage baggage, Closure callback) {
+    return runWithSpan(spanName) {
+      try (var unusedScope = baggage.makeCurrent()) {
+        callback.call()
+      }
+
+      Span.current()
+    }
+  }
+
+  abstract void setBaggageFlag()
+
+  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 static Span spanFromContext(Context context) {
     return Span.fromContext(context);
   }
 
+  public static Baggage baggageFromContext(Context context) {
+    return Baggage.fromContext(context);
+  }
+
   private Java8BytecodeBridge() {}
 }