diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ActivationListener.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ActivationListener.java index 43e5f6bdb5..a728add3da 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ActivationListener.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ActivationListener.java @@ -35,7 +35,7 @@ public interface ActivationListener { * @param context the {@link TraceContextHolder} which is being activated * @throws Throwable if there was an error while calling this method */ - void onActivate(TraceContextHolder context) throws Throwable; + void beforeActivate(TraceContextHolder context) throws Throwable; /** * A callback for {@link TraceContextHolder#deactivate()} @@ -46,5 +46,5 @@ public interface ActivationListener { * * @throws Throwable if there was an error while calling this method */ - void onDeactivate() throws Throwable; + void afterDeactivate() throws Throwable; } diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContextHolder.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContextHolder.java index 81ae7d9b69..3a363b27ad 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContextHolder.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContextHolder.java @@ -63,17 +63,17 @@ protected TraceContextHolder(ElasticApmTracer tracer) { public abstract boolean isChildOf(TraceContextHolder other); public T activate() { - tracer.activate(this); List activationListeners = tracer.getActivationListeners(); for (int i = 0; i < activationListeners.size(); i++) { try { - activationListeners.get(i).onActivate(this); + activationListeners.get(i).beforeActivate(this); } catch (Error e) { throw e; } catch (Throwable t) { - logger.warn("Exception while calling {}#onActivate", activationListeners.get(i).getClass().getSimpleName(), t); + logger.warn("Exception while calling {}#beforeActivate", activationListeners.get(i).getClass().getSimpleName(), t); } } + tracer.activate(this); return (T) this; } @@ -82,11 +82,11 @@ public T deactivate() { List activationListeners = tracer.getActivationListeners(); for (int i = 0; i < activationListeners.size(); i++) { try { - activationListeners.get(i).onDeactivate(); + activationListeners.get(i).afterDeactivate(); } catch (Error e) { throw e; } catch (Throwable t) { - logger.warn("Exception while calling {}#onDeactivate", activationListeners.get(i).getClass().getSimpleName(), t); + logger.warn("Exception while calling {}#afterDeactivate", activationListeners.get(i).getClass().getSimpleName(), t); } } return (T) this; diff --git a/apm-agent-plugins/apm-slf4j-plugin/src/main/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListener.java b/apm-agent-plugins/apm-slf4j-plugin/src/main/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListener.java index 8061f8248a..552bbb51fd 100644 --- a/apm-agent-plugins/apm-slf4j-plugin/src/main/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListener.java +++ b/apm-agent-plugins/apm-slf4j-plugin/src/main/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListener.java @@ -22,7 +22,6 @@ import co.elastic.apm.agent.cache.WeakKeySoftValueLoadingCache; import co.elastic.apm.agent.impl.ActivationListener; import co.elastic.apm.agent.impl.ElasticApmTracer; -import co.elastic.apm.agent.impl.transaction.Span; import co.elastic.apm.agent.impl.transaction.TraceContext; import co.elastic.apm.agent.impl.transaction.TraceContextHolder; import co.elastic.apm.agent.logging.LoggingConfiguration; @@ -40,6 +39,7 @@ public class Slf4JMdcActivationListener implements ActivationListener { private static final String TRACE_ID = "trace.id"; private static final String SPAN_ID = "span.id"; private static final String TRANSACTION_ID = "transaction.id"; + private final WeakKeySoftValueLoadingCache mdcPutMethodHandleCache = new WeakKeySoftValueLoadingCache<>(new WeakKeySoftValueLoadingCache.ValueSupplier() { @Nullable @@ -70,35 +70,58 @@ public MethodHandle get(ClassLoader classLoader) { }); @Nullable private LoggingConfiguration config; + @Nullable + private ElasticApmTracer tracer; @Override public void init(ElasticApmTracer tracer) { + this.tracer = tracer; config = tracer.getConfig(LoggingConfiguration.class); } @Override - public void onActivate(TraceContextHolder context) throws Throwable { - if (config != null && config.isLogCorrelationEnabled() && context.isSampled()) { - MethodHandle put = mdcPutMethodHandleCache.get(Thread.currentThread().getContextClassLoader()); - TraceContext traceContext = context.getTraceContext(); + public void beforeActivate(TraceContextHolder context) throws Throwable { + if (config != null && config.isLogCorrelationEnabled()) { + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); + + MethodHandle put = mdcPutMethodHandleCache.get(contextClassLoader); if (put != null) { - put.invokeExact(TRACE_ID, traceContext.getTraceId().toString()); - String idKey = context instanceof Span ? SPAN_ID : TRANSACTION_ID; - put.invokeExact(idKey, traceContext.getId().toString()); + TraceContext traceContext = context.getTraceContext(); + if (tracer != null) { + put.invokeExact(SPAN_ID, traceContext.getId().toString()); + if (tracer.getActive() == null) { + put.invokeExact(TRACE_ID, traceContext.getTraceId().toString()); + put.invokeExact(TRANSACTION_ID, traceContext.getTransactionId().toString()); + } + } } } } @Override - public void onDeactivate() throws Throwable { + public void afterDeactivate() throws Throwable { if (config != null && config.isLogCorrelationEnabled()) { - MethodHandle remove = mdcRemoveMethodHandleCache.get(Thread.currentThread().getContextClassLoader()); - if (remove != null) { - remove.invokeExact(TRACE_ID); - remove.invokeExact(SPAN_ID); - remove.invokeExact(TRANSACTION_ID); + + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); + if (tracer != null) { + TraceContextHolder active = tracer.getActive(); + + MethodHandle remove = mdcRemoveMethodHandleCache.get(contextClassLoader); + if (remove != null) { + remove.invokeExact(SPAN_ID); + if (active == null) { + remove.invokeExact(TRACE_ID); + remove.invokeExact(TRANSACTION_ID); + } + } + + if (active != null) { + MethodHandle put = mdcPutMethodHandleCache.get(contextClassLoader); + if (put != null) { + put.invokeExact(SPAN_ID, active.getTraceContext().getId().toString()); + } + } } } } - } diff --git a/apm-agent-plugins/apm-slf4j-plugin/src/test/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListenerTest.java b/apm-agent-plugins/apm-slf4j-plugin/src/test/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListenerTest.java index e1a7c36bdc..5d09ad6851 100644 --- a/apm-agent-plugins/apm-slf4j-plugin/src/test/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListenerTest.java +++ b/apm-agent-plugins/apm-slf4j-plugin/src/test/java/co/elastic/apm/agent/slf4j/Slf4JMdcActivationListenerTest.java @@ -22,6 +22,8 @@ import co.elastic.apm.agent.AbstractInstrumentationTest; import co.elastic.apm.agent.configuration.SpyConfiguration; import co.elastic.apm.agent.impl.Scope; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; +import co.elastic.apm.agent.impl.transaction.Span; import co.elastic.apm.agent.impl.transaction.Transaction; import co.elastic.apm.agent.logging.LoggingConfiguration; import org.junit.jupiter.api.BeforeEach; @@ -51,13 +53,23 @@ void testMdcIntegration() { assertMdcIsEmpty(); try (Scope scope = transaction.activateInScope()) { assertMdcIsSet(transaction); + Span child = transaction.createSpan(); + try (Scope childScope = child.activateInScope()) { + assertMdcIsSet(child); + Span grandchild = child.createSpan(); + try (Scope grandchildScope = grandchild.activateInScope()) { + assertMdcIsSet(grandchild); + } + assertMdcIsSet(child); + } + assertMdcIsSet(transaction); } assertMdcIsEmpty(); transaction.end(); } @Test - void testMdcIntegrationScopeInDifferentThread() throws Exception { + void testMdcIntegrationTransactionScopeInDifferentThread() throws Exception { when(loggingConfiguration.isLogCorrelationEnabled()).thenReturn(true); Transaction transaction = tracer.startTransaction().withType("request").withName("test"); assertMdcIsEmpty(); @@ -74,6 +86,33 @@ void testMdcIntegrationScopeInDifferentThread() throws Exception { transaction.end(); } + @Test + void testMdcIntegrationContextScopeInDifferentThread() throws Exception { + when(loggingConfiguration.isLogCorrelationEnabled()).thenReturn(true); + final Transaction transaction = tracer.startTransaction().withType("request").withName("test"); + assertMdcIsEmpty(); + try (Scope scope = transaction.activateInScope()) { + assertMdcIsSet(transaction); + final Span child = transaction.createSpan(); + try (Scope childScope = child.activateInScope()) { + assertMdcIsSet(child); + Thread thread = new Thread(() -> { + assertMdcIsEmpty(); + try (Scope otherThreadScope = child.getTraceContext().activateInScope()) { + assertMdcIsSet(child); + } + assertMdcIsEmpty(); + }); + thread.start(); + thread.join(); + assertMdcIsSet(child); + } + assertMdcIsSet(transaction); + } + assertMdcIsEmpty(); + transaction.end(); + } + @Test void testDisablingAtRuntimeNotPossible() throws IOException { assertThat(loggingConfiguration.isLogCorrelationEnabled()).isFalse(); @@ -85,9 +124,10 @@ void testDisablingAtRuntimeNotPossible() throws IOException { assertThat(loggingConfiguration.isLogCorrelationEnabled()).isTrue(); } - private void assertMdcIsSet(Transaction transaction) { - assertThat(MDC.get("trace.id")).isEqualTo(transaction.getTraceContext().getTraceId().toString()); - assertThat(MDC.get("transaction.id")).isEqualTo(transaction.getTraceContext().getId().toString()); + private void assertMdcIsSet(AbstractSpan span) { + assertThat(MDC.get("trace.id")).isEqualTo(span.getTraceContext().getTraceId().toString()); + assertThat(MDC.get("transaction.id")).isEqualTo(span.getTraceContext().getTransactionId().toString()); + assertThat(MDC.get("span.id")).isEqualTo(span.getTraceContext().getId().toString()); } private void assertMdcIsEmpty() {