Support for setting the application class loader for a transaction via the public api

CHANGELOG.asciidoc

@@ -54,6 +54,7 @@ endif::[]
 * Added support for setting the service name on Log4j2's EcsLayout - {pull}2296[#2296]
 * Print the used instrumentation groups when the application stops - {pull}2448[#2448]
 * Add `elastic.apm.start_async` property that makes the agent start on a non-premain/main thread - {pull}2454[#2454]
+* Added support for setting the application class loader for a transaction via the public api - {pull}2444[#2444]
 
 [float]
 ===== Bug fixes

apm-agent-api/src/main/java/co/elastic/apm/api/NoopTransaction.java

@@ -226,4 +226,10 @@ public Span setDestinationAddress(@Nullable String address, int port) {
     public Span setDestinationService(@Nullable String resource) {
         return this;
     }
+
+    @Nonnull
+    @Override
+    public Transaction setApplicationClassLoader(@Nullable ClassLoader classLoader) {
+        return this;
+    }
 }

apm-agent-api/src/main/java/co/elastic/apm/api/Transaction.java

@@ -331,4 +331,16 @@ public interface Transaction extends Span {
     @Override
     Scope activate();
 
+    /**
+     * Sets the class loader of the application that started the transaction.
+     * <p>
+     * The class loader is used to determine the service name and version of the transaction. If this method is called
+     * after child spans are already created, they may be associated with the wrong service name and version.
+     * </p>
+     *
+     * @param classLoader the class loader of the application that started the transaction
+     * @return this
+     */
+    @Nonnull
+    Transaction setApplicationClassLoader(ClassLoader classLoader);
 }

apm-agent-api/src/main/java/co/elastic/apm/api/TransactionImpl.java

@@ -186,4 +186,10 @@ public Transaction setDestinationAddress(@Nullable String address, int port) {
     public Transaction setDestinationService(@Nullable String resource) {
         throw new UnsupportedOperationException();
     }
+
+    @Nonnull
+    public Transaction setApplicationClassLoader(ClassLoader classLoader) {
+        // co.elastic.apm.agent.pluginapi.TransactionInstrumentation.SetClassLoaderInstrumentation
+        return this;
+    }
 }

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java

@@ -174,7 +174,7 @@ public Transaction startRootTransaction(Sampler sampler, long epochMicros, @Null
         Transaction transaction = null;
         if (isRunning()) {
             transaction = createTransaction().start(TraceContext.asRoot(), null, epochMicros, sampler, initiatingClassLoader);
-            afterTransactionStart(initiatingClassLoader, transaction);
+            afterTransactionStart(transaction);
         }
         return transaction;
     }
@@ -199,7 +199,7 @@ public <C> Transaction startChildTransaction(@Nullable C headerCarrier, TextHead
         if (isRunning()) {
             transaction = createTransaction().start(TraceContext.<C>getFromTraceContextTextHeaders(), headerCarrier,
                 textHeadersGetter, epochMicros, sampler, initiatingClassLoader);
-            afterTransactionStart(initiatingClassLoader, transaction);
+            afterTransactionStart(transaction);
         }
         return transaction;
     }
@@ -218,24 +218,19 @@ public <C> Transaction startChildTransaction(@Nullable C headerCarrier, BinaryHe
         if (isRunning()) {
             transaction = createTransaction().start(TraceContext.<C>getFromTraceContextBinaryHeaders(), headerCarrier,
                 binaryHeadersGetter, epochMicros, sampler, initiatingClassLoader);
-            afterTransactionStart(initiatingClassLoader, transaction);
+            afterTransactionStart(transaction);
         }
         return transaction;
     }
 
-    private void afterTransactionStart(@Nullable ClassLoader initiatingClassLoader, Transaction transaction) {
+    private void afterTransactionStart(Transaction transaction) {
         if (logger.isDebugEnabled()) {
             logger.debug("startTransaction {}", transaction);
             if (logger.isTraceEnabled()) {
                 logger.trace("starting transaction at",
                     new RuntimeException("this exception is just used to record where the transaction has been started from"));
             }
         }
-        final ServiceInfo serviceInfo = getServiceInfo(initiatingClassLoader);
-        if (serviceInfo != null) {
-            transaction.getTraceContext().setServiceName(serviceInfo.getServiceName());
-            transaction.getTraceContext().setServiceVersion(serviceInfo.getServiceVersion());
-        }
     }
 
     public Transaction noopTransaction() {
@@ -344,11 +339,7 @@ private ErrorCapture captureException(long epochMicros, @Nullable Throwable e, @
                 parent.setNonDiscardable();
             } else {
                 error.getTraceContext().getId().setToRandomValue();
-                ServiceInfo serviceInfo = getServiceInfo(initiatingClassLoader);
-                if (serviceInfo != null) {
-                    error.getTraceContext().setServiceName(serviceInfo.getServiceName());
-                    error.getTraceContext().setServiceVersion(serviceInfo.getServiceVersion());
-                }
+                error.getTraceContext().setApplicationClassLoader(initiatingClassLoader);
             }
             return error;
         }

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java

@@ -19,17 +19,18 @@
 package co.elastic.apm.agent.impl.transaction;
 
 import co.elastic.apm.agent.configuration.CoreConfiguration;
+import co.elastic.apm.agent.configuration.ServiceInfo;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.Tracer;
 import co.elastic.apm.agent.impl.sampling.Sampler;
 import co.elastic.apm.agent.objectpool.Recyclable;
+import co.elastic.apm.agent.sdk.logging.Logger;
+import co.elastic.apm.agent.sdk.logging.LoggerFactory;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
 import co.elastic.apm.agent.util.ByteUtils;
 import co.elastic.apm.agent.util.ClassLoaderUtils;
 import co.elastic.apm.agent.util.HexUtils;
-import co.elastic.apm.agent.sdk.logging.Logger;
-import co.elastic.apm.agent.sdk.logging.LoggerFactory;
 
 import javax.annotation.Nullable;
 import java.lang.ref.WeakReference;
@@ -732,8 +733,8 @@ public int hashCode() {
         return Objects.hash(traceId, id, parentId, flags);
     }
 
-    void setApplicationClassLoader(@Nullable ClassLoader classLoader) {
-        if (ClassLoaderUtils.isBootstrapClassLoader(classLoader) || ClassLoaderUtils.isAgentClassLoader(classLoader)) {
+    public void setApplicationClassLoader(@Nullable ClassLoader classLoader) {
+        if (ClassLoaderUtils.isBootstrapClassLoader(classLoader)) {
             return;
         }
         WeakReference<ClassLoader> local = classLoaderWeakReferenceCache.get(classLoader);
@@ -742,6 +743,11 @@ void setApplicationClassLoader(@Nullable ClassLoader classLoader) {
             classLoaderWeakReferenceCache.putIfAbsent(classLoader, local);
         }
         applicationClassLoader = local;
+        final ServiceInfo serviceInfo = tracer.getServiceInfo(classLoader);
+        if (serviceInfo != null) {
+            setServiceName(serviceInfo.getServiceName());
+            setServiceVersion(serviceInfo.getServiceVersion());
+        }
     }
 
     @Nullable

apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/pluginapi/TransactionInstrumentation.java

@@ -147,4 +147,22 @@ public static void addCustomContext(@Advice.FieldValue(value = "span", typing =
             }
         }
     }
+
+    public static class SetClassLoaderInstrumentation extends TransactionInstrumentation {
+        public SetClassLoaderInstrumentation() {
+            super(named("setApplicationClassLoader"));
+        }
+
+        public static class AdviceClass {
+
+            @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+            public static void setClassLoader(@Advice.FieldValue(value = "span", typing = Assigner.Typing.DYNAMIC) Object transaction,
+                                              @Advice.Argument(0) ClassLoader classLoader) {
+                if (transaction instanceof Transaction) {
+                    TraceContext traceContext = ((Transaction) transaction).getTraceContext();
+                    traceContext.setApplicationClassLoader(classLoader);
+                }
+            }
+        }
+    }
 }

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

@@ -7,6 +7,7 @@ co.elastic.apm.agent.pluginapi.TransactionInstrumentation$SetFrameworkNameInstru
 co.elastic.apm.agent.pluginapi.TransactionInstrumentation$SetUserInstrumentation
 co.elastic.apm.agent.pluginapi.TransactionInstrumentation$EnsureParentIdInstrumentation
 co.elastic.apm.agent.pluginapi.TransactionInstrumentation$SetResultInstrumentation
+co.elastic.apm.agent.pluginapi.TransactionInstrumentation$SetClassLoaderInstrumentation
 co.elastic.apm.agent.pluginapi.TransactionInstrumentation$AddCustomContextInstrumentation
 co.elastic.apm.agent.pluginapi.AbstractSpanInstrumentation$SetNameInstrumentation
 co.elastic.apm.agent.pluginapi.AbstractSpanInstrumentation$SetTypeInstrumentation

apm-agent-plugins/apm-api-plugin/src/test/java/co/elastic/apm/agent/pluginapi/TransactionInstrumentationTest.java

@@ -19,10 +19,12 @@
 package co.elastic.apm.agent.pluginapi;
 
 import co.elastic.apm.AbstractApiTest;
+import co.elastic.apm.agent.configuration.ServiceInfo;
 import co.elastic.apm.agent.impl.TracerInternalApiUtils;
 import co.elastic.apm.api.AbstractSpanImplAccessor;
 import co.elastic.apm.api.ElasticApm;
 import co.elastic.apm.api.Outcome;
+import co.elastic.apm.api.Scope;
 import co.elastic.apm.api.Span;
 import co.elastic.apm.api.Transaction;
 import org.apache.commons.lang3.RandomStringUtils;
@@ -32,6 +34,8 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 import javax.annotation.Nullable;
+import java.net.URL;
+import java.net.URLClassLoader;
 import java.security.SecureRandom;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -296,6 +300,24 @@ void setOutcome_success() {
         testSetOutcome(Outcome.SUCCESS);
     }
 
+    @Test
+    void setClassLoader() {
+        ClassLoader classLoader = new URLClassLoader(new URL[0]);
+
+        tracer.overrideServiceInfoForClassLoader(classLoader, ServiceInfo.of("My Service", "My Version"));
+
+        transaction.setApplicationClassLoader(classLoader);
+        try (Scope scope = transaction.activate()) {
+            assertThat(tracer.getActive().getTraceContext().getApplicationClassLoader()).isSameAs(classLoader);
+        } finally {
+            endTransaction();
+        }
+
+        co.elastic.apm.agent.impl.transaction.Transaction reportedTransaction = reporter.getFirstTransaction();
+        assertThat(reportedTransaction.getTraceContext().getServiceName()).isEqualTo("My Service");
+        assertThat(reportedTransaction.getTraceContext().getServiceVersion()).isEqualTo("My Version");
+    }
+
     private void testSetOutcome(Outcome outcome) {
         // set it first to a different value than the expected one
         Outcome[] values = Outcome.values();

apm-agent-plugins/apm-log-correlation-plugin/src/main/java/co/elastic/apm/agent/mdc/MdcActivationListener.java

@@ -27,6 +27,7 @@
 import co.elastic.apm.agent.logging.LoggingConfiguration;
 import co.elastic.apm.agent.sdk.logging.Logger;
 import co.elastic.apm.agent.sdk.logging.LoggerFactory;
+import co.elastic.apm.agent.util.ClassLoaderUtils;
 
 import javax.annotation.Nullable;
 import java.lang.invoke.MethodHandle;
@@ -227,10 +228,10 @@ public void after(TraceContext deactivatedContext, boolean isError) throws Throw
      */
     private ClassLoader getApplicationClassLoader(TraceContext context) {
         ClassLoader applicationClassLoader = context.getApplicationClassLoader();
-        if (applicationClassLoader != null) {
-            return applicationClassLoader;
-        } else {
+        if (applicationClassLoader == null || ClassLoaderUtils.isAgentClassLoader(applicationClassLoader)) {
             return getFallbackClassLoader();
+        } else {
+            return applicationClassLoader;
         }
     }
     private ClassLoader getFallbackClassLoader() {

docs/public-api.asciidoc

@@ -551,6 +551,15 @@ use `UNKNOWN` when the outcome can't be properly known.
 Outcome is used to compute error rates between services, using `UNKNOWN` will not alter those rates.
 The value set through API will have higher priority over the value that might be set by auto-instrumentation.
 
+[float]
+[[api-transaction-set-class-loader]]
+==== `Transaction setApplicationClassLoader(ClassLoader classLoader)`
+Sets the class loader of the application that started the transaction.
+
+* `classLoader` the class loader of the application that started the transaction
+
+NOTE: The class loader is used to determine the service name and version of the transaction. If this method is called after
+child spans are already created, they may be associated with the wrong service name and version.
 
 [float]
 [[api-transaction-set-start-timestamp]]