Log correlation for JBoss LogManager

apm-agent-plugins/apm-logging-plugin/apm-jboss-logging-plugin/pom.xml

@@ -27,6 +27,24 @@
             <version>3.4.2.Final</version>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jboss.logmanager</groupId>
+            <artifactId>jboss-logmanager</artifactId>
+            <version>2.1.18.Final</version>
+            <scope>provided</scope>
+        </dependency>
     </dependencies>
 
+    <build>
+        <plugins>
+            <plugin>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <configuration>
+                    <systemProperties>
+                        <java.util.logging.manager>org.jboss.logmanager.LogManager</java.util.logging.manager>
+                    </systemProperties>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
 </project>

apm-agent-plugins/apm-logging-plugin/apm-jboss-logging-plugin/src/main/java/co/elastic/apm/agent/jbosslogging/correlation/JBossLogManagerCorrelationHelper.java

@@ -0,0 +1,47 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ */
+package co.elastic.apm.agent.jbosslogging.correlation;
+
+import co.elastic.apm.agent.loginstr.correlation.AbstractLogCorrelationHelper;
+import org.jboss.logmanager.ExtLogRecord;
+
+public class JBossLogManagerCorrelationHelper extends AbstractLogCorrelationHelper.DefaultLogCorrelationHelper {
+
+    private final ThreadLocal<ExtLogRecord> cachedRecord = new ThreadLocal<>();
+
+    public boolean beforeLoggingEvent(ExtLogRecord record) {
+        cachedRecord.set(record);
+        return super.beforeLoggingEvent();
+    }
+
+    @Override
+    public void afterLoggingEvent(boolean addedToMdc) {
+        super.afterLoggingEvent(addedToMdc);
+        cachedRecord.remove();
+    }
+
+    @Override
+    protected void addToMdc(String key, String value) {
+        cachedRecord.get().putMdc(key, value);
+    }
+
+    @Override
+    protected void removeFromMdc(String key) {
+    }
+}

apm-agent-plugins/apm-logging-plugin/apm-jboss-logging-plugin/src/main/java/co/elastic/apm/agent/jbosslogging/correlation/JBossLogManagerCorrelationInstrumentation.java

@@ -0,0 +1,76 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ */
+package co.elastic.apm.agent.jbosslogging.correlation;
+
+import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.jboss.logmanager.ExtLogRecord;
+
+import java.util.Collection;
+import java.util.Collections;
+
+import static co.elastic.apm.agent.bci.bytebuddy.CustomElementMatchers.classLoaderCanLoadClass;
+import static net.bytebuddy.matcher.ElementMatchers.isBootstrapClassLoader;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+/**
+ * Instruments {@link org.jboss.logmanager.Logger#logRaw}
+ */
+public class JBossLogManagerCorrelationInstrumentation extends TracerAwareInstrumentation {
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Collections.singleton("jboss-logging-correlation");
+    }
+
+    @Override
+    public ElementMatcher.Junction<ClassLoader> getClassLoaderMatcher() {
+        return not(isBootstrapClassLoader())
+            .and(classLoaderCanLoadClass("org.jboss.logmanager.Logger"));
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return named("org.jboss.logmanager.Logger");
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("logRaw").and(takesArgument(0, named("org.jboss.logmanager.ExtLogRecord")));
+    }
+
+    public static class AdviceClass {
+        private static final JBossLogManagerCorrelationHelper helper = new JBossLogManagerCorrelationHelper();
+
+        @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+        public static boolean addToMdc(@Advice.Argument(0) ExtLogRecord record) {
+            return helper.beforeLoggingEvent(record);
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class, inline = false)
+        public static void removeFromMdc(@Advice.Enter boolean addedToMdc) {
+            helper.afterLoggingEvent(addedToMdc);
+        }
+    }
+}

apm-agent-plugins/apm-logging-plugin/apm-jboss-logging-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation

@@ -1 +1,2 @@
 co.elastic.apm.agent.jbosslogging.correlation.JBossLoggingCorrelationInstrumentation
+co.elastic.apm.agent.jbosslogging.correlation.JBossLogManagerCorrelationInstrumentation

apm-agent-plugins/apm-logging-plugin/apm-jboss-logging-plugin/src/test/java/co/elastic/apm/agent/jbosslogging/correlation/JBossLogManagerCorrelationInstrumentationTest.java

@@ -0,0 +1,137 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ */
+package co.elastic.apm.agent.jbosslogging.correlation;
+
+import co.elastic.apm.agent.AbstractInstrumentationTest;
+import co.elastic.apm.agent.impl.TextHeaderMapAccessor;
+import co.elastic.apm.agent.impl.error.ErrorCapture;
+import co.elastic.apm.agent.impl.transaction.TraceContext;
+import co.elastic.apm.agent.impl.transaction.Transaction;
+import co.elastic.apm.agent.loginstr.correlation.AbstractLogCorrelationHelper;
+import org.jboss.logmanager.ExtHandler;
+import org.jboss.logmanager.ExtLogRecord;
+import org.jboss.logmanager.Level;
+import org.jboss.logmanager.LogManager;
+import org.jboss.logmanager.Logger;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+import java.util.Objects;
+
+import static co.elastic.apm.agent.impl.transaction.TraceContext.W3C_TRACE_PARENT_TEXTUAL_HEADER_NAME;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class JBossLogManagerCorrelationInstrumentationTest extends AbstractInstrumentationTest {
+
+    public static final String LOGGER_NAME = "Test-Logger";
+
+    private static final LoggingCorrelationVerifier loggingCorrelationVerifier = new LoggingCorrelationVerifier();
+    private static Logger logger;
+
+    private Transaction transaction;
+
+    @BeforeAll
+    static void initializeLogger() {
+        logger = (org.jboss.logmanager.Logger) LogManager.getLogManager().getLogger(LOGGER_NAME);
+        logger.addHandler(new LoggingCorrelationVerifier());
+    }
+
+    @BeforeEach
+    public void setup() {
+        transaction = Objects.requireNonNull(tracer.startRootTransaction(null)).activate();
+        loggingCorrelationVerifier.reset();
+    }
+
+    @AfterEach
+    public void tearDown() {
+        transaction.deactivate().end();
+    }
+
+    @Test
+    public void testSimple() {
+        assertThat(loggingCorrelationVerifier.isVerified()).isFalse();
+        // TraceContext#toString() returns the text representation of the traceparent header
+        logger.info(transaction.getTraceContext().toString());
+        assertThat(loggingCorrelationVerifier.isVerified()).isTrue();
+    }
+
+    // todo: enable once support for jboss-logmanager error logging is added
+    @Test
+    @Disabled
+    public void testErrorLogging() {
+        assertThat(loggingCorrelationVerifier.isVerified()).isFalse();
+        // TraceContext#toString() returns the text representation of the traceparent header
+        logger.log(Level.ERROR, transaction.getTraceContext().toString(), new RuntimeException());
+        assertThat(loggingCorrelationVerifier.isVerified()).isTrue();
+    }
+
+    private static class LoggingCorrelationVerifier extends ExtHandler {
+        private boolean verified;
+
+        public void reset() {
+            verified = false;
+        }
+
+        public boolean isVerified() {
+            return verified;
+        }
+
+        public void setVerified() {
+            verified = true;
+        }
+
+        @Override
+        protected void doPublish(ExtLogRecord record) {
+            try {
+                Object traceId = record.getMdc(AbstractLogCorrelationHelper.TRACE_ID_MDC_KEY);
+                System.out.println("traceId = " + traceId);
+                assertThat(traceId).isNotNull();
+                Object transactionId = record.getMdc(AbstractLogCorrelationHelper.TRANSACTION_ID_MDC_KEY);
+                System.out.println("transactionId = " + transactionId);
+                assertThat(transactionId).isNotNull();
+                Object errorId = record.getMdc(AbstractLogCorrelationHelper.ERROR_ID_MDC_KEY);
+                System.out.println("errorId = " + transactionId);
+                boolean shouldContainErrorId = record.getLevel().getName().equals("ERROR");
+                String traceParent = record.getMessage();
+                assertThat(traceParent).isNotNull();
+                Map<String, String> textHeaderMap = Map.of(W3C_TRACE_PARENT_TEXTUAL_HEADER_NAME, traceParent);
+                TraceContext childTraceContext = TraceContext.with64BitId(tracer);
+                TraceContext.<Map<String, String>>getFromTraceContextTextHeaders().asChildOf(childTraceContext, textHeaderMap, TextHeaderMapAccessor.INSTANCE);
+                System.out.println("childTraceContext = " + childTraceContext);
+                assertThat(childTraceContext.getTraceId().toString()).isEqualTo(traceId.toString());
+                assertThat(childTraceContext.getParentId().toString()).isEqualTo(transactionId.toString());
+                if (shouldContainErrorId) {
+                    assertThat(errorId).isNotNull();
+                    ErrorCapture activeError = ErrorCapture.getActive();
+                    assertThat(activeError).isNotNull();
+                    assertThat(activeError.getTraceContext().getId().toString()).isEqualTo(errorId.toString());
+                } else {
+                    assertThat(errorId).isNull();
+                }
+                loggingCorrelationVerifier.setVerified();
+            } catch (Throwable throwable) {
+                throwable.printStackTrace();
+            }
+        }
+    }
+}