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 30, 2022 · b8d4eff · b8d4eff
1 parent efd40f8
commit b8d4eff
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 2 deletions.
diff --git a/.../javaagent/instrumentation/jbosslogmanager/mdc/v1_1/JbossExtLogRecordInstrumentation.java b/.../javaagent/instrumentation/jbosslogmanager/mdc/v1_1/JbossExtLogRecordInstrumentation.java
@@ -20,6 +20,7 @@
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -33,13 +34,19 @@ 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 +88,41 @@ 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.getSpanId());
+      }
+
+      if (!value.containsKey(TRACE_FLAGS)) {
+        value.put(TRACE_FLAGS, spanContext.getTraceFlags().asHex());
+      }
+    }
+  }
 }
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:
@@ -76,11 +76,21 @@ class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
     then:
 
     logRecords.size() == 3
+
+    // The method getMdcCopy exists only in jboss-logmanager 1.3+
+    def hasGetMdcCopy = logRecords.get(0).metaClass.getMetaMethod("getMdcCopy") != null
+
     logRecords.get(2).message == "log message 1"
     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"
 
+    if (hasGetMdcCopy) {
+      assert logRecords.get(2).getMdcCopy().get("trace_id") == span1.spanContext.traceId
+      assert logRecords.get(2).getMdcCopy().get("span_id") == span1.spanContext.spanId
+      assert logRecords.get(2).getMdcCopy().get("trace_flags") == "01"
+    }
+
     logRecords.get(1).message == "log message 2"
     logRecords.get(1).getMdc("trace_id") == null
     logRecords.get(1).getMdc("span_id") == null
@@ -91,6 +101,12 @@ class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
     logRecords.get(0).getMdc("span_id") == span2.spanContext.spanId
     logRecords.get(0).getMdc("trace_flags") == "01"
 
+    if (hasGetMdcCopy) {
+      assert logRecords.get(0).getMdcCopy().get("trace_id") == span2.spanContext.traceId
+      assert logRecords.get(0).getMdcCopy().get("span_id") == span2.spanContext.spanId
+      assert logRecords.get(0).getMdcCopy().get("trace_flags") == "01"
+    }
+
     cleanup:
     logRecords.clear()
   }