Instrumentation for jboss-logmanager getMdcCopy() (#6111)

This allows, for example, in JSON logs to have the MDC to contain span and trace information.
open-telemetry · May 28, 2022 · 576c8ff · 576c8ff
1 parent efd40f8
commit 576c8ff
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 3 deletions.
diff --git a/instrumentation/jboss-logmanager/jboss-logmanager-mdc-1.1/javaagent/build.gradle.kts b/instrumentation/jboss-logmanager/jboss-logmanager-mdc-1.1/javaagent/build.gradle.kts
@@ -12,5 +12,5 @@ muzzle {
 }
 
 dependencies {
-  library("org.jboss.logmanager:jboss-logmanager:1.1.0.GA")
+  library("org.jboss.logmanager:jboss-logmanager:1.3.0.Final")
 }
diff --git a/.../javaagent/instrumentation/jbosslogmanager/mdc/v1_1/JbossExtLogRecordInstrumentation.java b/.../javaagent/instrumentation/jbosslogmanager/mdc/v1_1/JbossExtLogRecordInstrumentation.java
@@ -24,6 +24,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.jboss.logmanager.ExtLogRecord;
+import java.util.Map;
 
 public class JbossExtLogRecordInstrumentation implements TypeInstrumentation {
   @Override
@@ -33,13 +34,22 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
   @Override
   public void transform(TypeTransformer transformer) {
+    // available since jboss-logmanager 1.1
     transformer.applyAdviceToMethod(
         isMethod()
             .and(isPublic())
             .and(named("getMdc"))
             .and(takesArguments(1))
             .and(takesArgument(0, String.class)),
         JbossExtLogRecordInstrumentation.class.getName() + "$GetMdcAdvice");
+
+    // available since jboss-logmanager 1.3
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(takesArguments(0))
+            .and(named("getMdcCopy")),
+        JbossExtLogRecordInstrumentation.class.getName() + "$GetMdcCopyAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -81,4 +91,40 @@ public static void onExit(
       }
     }
   }
+
+  public static class GetMdcCopyAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This ExtLogRecord record,
+        @Advice.Return(readOnly = false) Map<String, String> value) {
+
+      if (value.containsKey(TRACE_ID) && value.containsKey(SPAN_ID) && value.containsKey(TRACE_FLAGS)) {
+        return;
+      }
+
+      Context context = VirtualField.find(ExtLogRecord.class, Context.class).get(record);
+      if (context == null) {
+        return;
+      }
+
+      SpanContext spanContext = Java8BytecodeBridge.spanFromContext(context).getSpanContext();
+      if (!spanContext.isValid()) {
+        return;
+      }
+
+      if (!value.containsKey(TRACE_ID)) {
+        value.put(TRACE_ID, spanContext.getTraceId());
+      }
+
+      if (!value.containsKey(SPAN_ID)) {
+        value.put(SPAN_ID, spanContext.getTraceId());
+      }
+
+      if (!value.containsKey(TRACE_FLAGS)) {
+        value.put(TRACE_FLAGS, spanContext.getTraceId());
+      }
+    }
+  }
+
 }
diff --git a/...gmanager/jboss-logmanager-mdc-1.1/javaagent/src/test/groovy/JbossLogmanagerMdcTest.groovy b/...gmanager/jboss-logmanager-mdc-1.1/javaagent/src/test/groovy/JbossLogmanagerMdcTest.groovy
@@ -15,7 +15,7 @@ import java.util.logging.LogRecord
 class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
   class LogHandler extends Handler {
     public List<ExtLogRecord> logRecords
-    
+
     LogHandler(LinkedList<ExtLogRecord> logRecords) {
       this.logRecords = logRecords
     }
@@ -31,7 +31,7 @@ class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
 
     @Override
     void close() throws SecurityException {
-    }   
+    }
   }
   def "no ids when no span"() {
     given:
@@ -80,6 +80,10 @@ class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
     logRecords.get(2).getMdc("trace_id") == span1.spanContext.traceId
     logRecords.get(2).getMdc("span_id") == span1.spanContext.spanId
     logRecords.get(2).getMdc("trace_flags") == "01"
+    def mdcCopy2 = logRecords.get(2).getMdcCopy()
+    mdcCopy2.get("trace_id") != null
+    mdcCopy2.get("span_id") != null
+    mdcCopy2.get("trace_flags") != null
 
     logRecords.get(1).message == "log message 2"
     logRecords.get(1).getMdc("trace_id") == null
@@ -90,6 +94,10 @@ class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
     logRecords.get(0).getMdc("trace_id") == span2.spanContext.traceId
     logRecords.get(0).getMdc("span_id") == span2.spanContext.spanId
     logRecords.get(0).getMdc("trace_flags") == "01"
+    def mdcCopy0 = logRecords.get(2).getMdcCopy()
+    mdcCopy0.get("trace_id") != null
+    mdcCopy0.get("span_id") != null
+    mdcCopy0.get("trace_flags") != null
 
     cleanup:
     logRecords.clear()