Add jboss-logmanager mdc support (open-telemetry#5842)

* add support for jboss-logmanager-mdc

Signed-off-by: Cuichen Li <[email protected]>

* update based on feedback

Signed-off-by: Cuichen Li <[email protected]>

* update

* check spanContext is valid

instrumentation/jboss-logmanager-mdc-1.1/javaagent/build.gradle.kts

@@ -0,0 +1,16 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("org.jboss.logmanager")
+    module.set("jboss-logmanager")
+    versions.set("[1.1.0.GA,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("org.jboss.logmanager:jboss-logmanager:1.1.0.GA")
+}

instrumentation/jboss-logmanager-mdc-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jbosslogmanager/v1_1/JbossExtLogRecordInstrumentation.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jbosslogmanager.v1_1;
+
+import static io.opentelemetry.instrumentation.api.log.LoggingContextConstants.SPAN_ID;
+import static io.opentelemetry.instrumentation.api.log.LoggingContextConstants.TRACE_FLAGS;
+import static io.opentelemetry.instrumentation.api.log.LoggingContextConstants.TRACE_ID;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.jboss.logmanager.ExtLogRecord;
+
+public class JbossExtLogRecordInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.jboss.logmanager.ExtLogRecord");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("getMdc"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class)),
+        JbossExtLogRecordInstrumentation.class.getName() + "$GetMdcAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetMdcAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This ExtLogRecord record,
+        @Advice.Argument(0) String key,
+        @Advice.Return(readOnly = false) String value) {
+      if (TRACE_ID.equals(key) || SPAN_ID.equals(key) || TRACE_FLAGS.equals(key)) {
+        if (value != null) {
+          // Assume already instrumented event if traceId/spanId/sampled is present.
+          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;
+        }
+
+        switch (key) {
+          case TRACE_ID:
+            value = spanContext.getTraceId();
+            break;
+          case SPAN_ID:
+            value = spanContext.getSpanId();
+            break;
+          case TRACE_FLAGS:
+            value = spanContext.getTraceFlags().asHex();
+            break;
+          default:
+            // do nothing
+        }
+      }
+    }
+  }
+}

instrumentation/jboss-logmanager-mdc-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jbosslogmanager/v1_1/JbossLoggerInstrumentation.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jbosslogmanager.v1_1;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.jboss.logmanager.ExtLogRecord;
+
+public class JbossLoggerInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.jboss.logmanager.Logger");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("logRaw"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("org.jboss.logmanager.ExtLogRecord"))),
+        JbossLoggerInstrumentation.class.getName() + "$CallAppendersAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class CallAppendersAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(0) ExtLogRecord record) {
+      VirtualField.find(ExtLogRecord.class, Context.class)
+          .set(record, Java8BytecodeBridge.currentContext());
+    }
+  }
+}

instrumentation/jboss-logmanager-mdc-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jbosslogmanager/v1_1/JbossLogmanagerInstrumentationModule.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jbosslogmanager.v1_1;
+
+import static java.util.Arrays.asList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class JbossLogmanagerInstrumentationModule extends InstrumentationModule {
+
+  public JbossLogmanagerInstrumentationModule() {
+    super("jboss-logmanager-mdc", "jboss-logmanager-mdc-1.1");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(new JbossLoggerInstrumentation(), new JbossExtLogRecordInstrumentation());
+  }
+}

instrumentation/jboss-logmanager-mdc-1.1/javaagent/src/test/groovy/JbossLogmanagerMdcTest.groovy

@@ -0,0 +1,97 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import io.opentelemetry.api.trace.Span
+import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
+import org.jboss.logmanager.ExtLogRecord
+import org.jboss.logmanager.Level
+import org.jboss.logmanager.LogContext
+
+import java.util.logging.Handler
+import java.util.logging.LogRecord
+
+class JbossLogmanagerMdcTest extends AgentInstrumentationSpecification {
+  class LogHandler extends Handler {
+    public List<ExtLogRecord> logRecords
+
+    LogHandler(LinkedList<ExtLogRecord> logRecords) {
+      this.logRecords = logRecords
+    }
+    @Override
+    void publish(LogRecord record) {
+      logRecords.push(record as ExtLogRecord)
+    }
+
+    @Override
+    void flush() {
+
+    }
+
+    @Override
+    void close() throws SecurityException {
+    }   
+  }
+  def "no ids when no span"() {
+    given:
+    LinkedList<ExtLogRecord> logRecords = []
+    def logger = LogContext.getLogContext().getLogger('TestLogger')
+    logger.setLevel(Level.INFO)
+    logger.addHandler(new LogHandler(logRecords))
+
+    when:
+    logger.info("log message 1")
+
+    then:
+    logRecords.size() == 1
+    logRecords.get(0).message == "log message 1"
+    logRecords.get(0).getMdc("trace_id") == null
+    logRecords.get(0).getMdc("span_id") == null
+    logRecords.get(0).getMdc("trace_flags") == null
+
+    cleanup:
+    logRecords.clear()
+  }
+
+  def "ids when span"() {
+    given:
+    def logger = LogContext.getLogContext().getLogger('TestLogger')
+    logger.setLevel(Level.DEBUG)
+    LinkedList<ExtLogRecord> logRecords = []
+    logger.addHandler(new LogHandler(logRecords))
+    when:
+    def span1 = runWithSpan("test") {
+      logger.info("log message 1")
+      Span.current()
+    }
+
+    logger.info("log message 2")
+
+    def span2 = runWithSpan("test 2") {
+      logger.info("log message 3")
+      Span.current()
+    }
+
+    then:
+
+    logRecords.size() == 3
+    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"
+
+    logRecords.get(1).message == "log message 2"
+    logRecords.get(1).getMdc("trace_id") == null
+    logRecords.get(1).getMdc("span_id") == null
+    logRecords.get(1).getMdc("trace_flags") == null
+
+    logRecords.get(0).message == "log message 3"
+    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"
+
+    cleanup:
+    logRecords.clear()
+  }
+}

settings.gradle.kts

@@ -254,6 +254,7 @@ include(":instrumentation:jaxws:jaxws-2.0-wildfly-testing")
 include(":instrumentation:jaxws:jaxws-common:library")
 include(":instrumentation:jaxws:jaxws-jws-api-1.1:javaagent")
 include(":instrumentation:jboss-logmanager-1.1:javaagent")
+include(":instrumentation:jboss-logmanager-mdc-1.1:javaagent")
 include(":instrumentation:jdbc:javaagent")
 include(":instrumentation:jdbc:library")
 include(":instrumentation:jdbc:testing")