Wrap Runnable in AsyncContext.start (#388)

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

@@ -7,9 +7,9 @@
  * Licensed 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.

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

@@ -63,6 +63,7 @@ public class ElasticApmTracer {
     private final ObjectPool<Transaction> transactionPool;
     private final ObjectPool<Span> spanPool;
     private final ObjectPool<ErrorCapture> errorPool;
+    private final ObjectPool<InScopeRunnableWrapper> runnableWrapperObjectPool;
     private final Reporter reporter;
     // Maintains a stack of all the activated spans
     // This way its easy to retrieve the bottom of the stack (the transaction)
@@ -107,6 +108,13 @@ public ErrorCapture createInstance() {
                     return new ErrorCapture(ElasticApmTracer.this);
                 }
             });
+        runnableWrapperObjectPool = QueueBasedObjectPool.ofRecyclable(AtomicQueueFactory.<InScopeRunnableWrapper>newQueue(createBoundedMpmc(maxPooledElements)), false,
+            new Allocator<InScopeRunnableWrapper>() {
+                @Override
+                public InScopeRunnableWrapper createInstance() {
+                    return new InScopeRunnableWrapper(ElasticApmTracer.this);
+                }
+            });
         sampler = ProbabilitySampler.of(coreConfiguration.getSampleRate().get());
         coreConfiguration.getSampleRate().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
             @Override
@@ -296,6 +304,14 @@ public void recycle(ErrorCapture error) {
         errorPool.recycle(error);
     }
 
+    public Runnable wrapRunnable(Runnable delegate, AbstractSpan<?> span) {
+        return runnableWrapperObjectPool.createInstance().wrap(delegate, span);
+    }
+
+    public void recycle(InScopeRunnableWrapper wrapper) {
+        runnableWrapperObjectPool.recycle(wrapper);
+    }
+
     /**
      * Called when the container shuts down.
      * Cleans up thread pools and other resources.

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

@@ -0,0 +1,90 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 Elastic and contributors
+ * %%
+ * Licensed 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.
+ * #L%
+ */
+package co.elastic.apm.agent.impl;
+
+
+import co.elastic.apm.agent.bci.VisibleForAdvice;
+import co.elastic.apm.agent.impl.transaction.AbstractSpan;
+import co.elastic.apm.agent.objectpool.Recyclable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+@VisibleForAdvice
+public class InScopeRunnableWrapper implements Runnable, Recyclable {
+    private final Logger logger = LoggerFactory.getLogger(InScopeRunnableWrapper.class);
+
+    @Nullable
+    private Runnable delegate;
+    @Nullable
+    private AbstractSpan<?> span;
+
+    private final ElasticApmTracer tracer;
+
+    InScopeRunnableWrapper(ElasticApmTracer tracer) {
+        this.tracer = tracer;
+    }
+
+    public InScopeRunnableWrapper wrap(Runnable delegate, AbstractSpan<?> span) {
+        this.delegate = delegate;
+        this.span = span;
+        return this;
+    }
+
+    @Override
+    public void run() {
+        if (span != null) {
+            try {
+                span.activate();
+            } catch (Throwable throwable) {
+                try {
+                    logger.warn("Failed to activate span");
+                } catch (Throwable t) {
+                    // do nothing, just never fail
+                }
+            }
+        }
+
+        try {
+            //noinspection ConstantConditions
+            delegate.run();
+        } finally {
+            try {
+                if (span != null) {
+                    span.deactivate();
+                }
+                tracer.recycle(this);
+            } catch (Throwable throwable) {
+                try {
+                    logger.warn("Failed to deactivate span or recycle");
+                } catch (Throwable t) {
+                    // do nothing, just never fail
+                }
+            }
+        }
+    }
+
+    @Override
+    public void resetState() {
+        delegate = null;
+        span = null;
+    }
+}

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java

@@ -23,6 +23,7 @@
 import co.elastic.apm.agent.bci.HelperClassManager;
 import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.impl.transaction.Transaction;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.NamedElement;
 import net.bytebuddy.description.method.MethodDescription;
@@ -52,83 +53,124 @@
  * {@code javax.servlet}, as these are inlined into the matching methods.
  * The agent itself does not have access to the Servlet API classes, as they are loaded by a child class loader.
  * See https://github.com/raphw/byte-buddy/issues/465 for more information.
- * However, the helper class {@link StartAsyncAdviceHelper} has access to the Servlet API,
+ * However, the helper class {@link AsyncContextAdviceHelper} has access to the Servlet API,
  * as it is loaded by the child classloader of {@link AsyncContext}
- * (see {@link StartAsyncAdvice#onExitStartAsync(HttpServletRequest, AsyncContext)}).
+ * (see {@link StartAsyncInstrumentation.StartAsyncAdvice#onExitStartAsync(HttpServletRequest, AsyncContext)}
+ * and {@link AsyncContextInstrumentation.AsyncContextStartAdvice#onEnterAsyncContextStart(Runnable)}).
  */
-public class AsyncInstrumentation extends ElasticApmInstrumentation {
+public abstract class AsyncInstrumentation extends ElasticApmInstrumentation {
 
     private static final String SERVLET_API_ASYNC_GROUP_NAME = "servlet-api-async";
     @Nullable
     @VisibleForAdvice
     // referring to AsyncContext is legal because of type erasure
-    public static HelperClassManager<StartAsyncAdviceHelper<AsyncContext>> asyncHelper;
+    public static HelperClassManager<AsyncContextAdviceHelper<AsyncContext>> asyncHelper;
 
     @Override
     public void init(ElasticApmTracer tracer) {
         asyncHelper = HelperClassManager.ForSingleClassLoader.of(tracer,
-            "co.elastic.apm.agent.servlet.helper.StartAsyncAdviceHelperImpl",
+            "co.elastic.apm.agent.servlet.helper.AsyncContextAdviceHelperImpl",
             "co.elastic.apm.agent.servlet.helper.ApmAsyncListener");
     }
 
     @Override
-    public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
-        return nameContains("Request");
+    public Collection<String> getInstrumentationGroupNames() {
+        return Arrays.asList(ServletInstrumentation.SERVLET_API, SERVLET_API_ASYNC_GROUP_NAME);
     }
 
-    @Override
-    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
-        return not(isInterface())
-            .and(hasSuperType(named("javax.servlet.http.HttpServletRequest")));
+    public interface AsyncContextAdviceHelper<T> {
+        void onExitStartAsync(T asyncContext);
     }
 
-    /**
-     * Matches
-     * <ul>
-     * <li>{@link HttpServletRequest#startAsync()}</li>
-     * <li>{@link HttpServletRequest#startAsync(ServletRequest, ServletResponse)}</li>
-     * </ul>
-     *
-     * @return
-     */
-    @Override
-    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
-        return named("startAsync")
-            .and(returns(named("javax.servlet.AsyncContext")))
-            .and(takesArguments(0)
-                .or(
-                    takesArgument(0, named("javax.servlet.ServletRequest"))
-                        .and(takesArgument(1, named("javax.servlet.ServletResponse")))
-                )
-            )
-            .and(isPublic());
-    }
+    public static class StartAsyncInstrumentation extends AsyncInstrumentation {
+        @Override
+        public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
+            return nameContains("Request");
+        }
 
-    @Override
-    public Collection<String> getInstrumentationGroupNames() {
-        return Arrays.asList(ServletInstrumentation.SERVLET_API, SERVLET_API_ASYNC_GROUP_NAME);
-    }
+        @Override
+        public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+            return not(isInterface())
+                .and(hasSuperType(named("javax.servlet.http.HttpServletRequest")));
+        }
 
-    @Override
-    public Class<?> getAdviceClass() {
-        return StartAsyncAdvice.class;
-    }
+        /**
+         * Matches
+         * <ul>
+         * <li>{@link HttpServletRequest#startAsync()}</li>
+         * <li>{@link HttpServletRequest#startAsync(ServletRequest, ServletResponse)}</li>
+         * </ul>
+         *
+         * @return
+         */
+        @Override
+        public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+            return isPublic()
+                .and(named("startAsync"))
+                .and(returns(named("javax.servlet.AsyncContext")))
+                .and(takesArguments(0)
+                    .or(
+                        takesArgument(0, named("javax.servlet.ServletRequest"))
+                            .and(takesArgument(1, named("javax.servlet.ServletResponse")))
+                    )
+                );
+        }
 
-    public interface StartAsyncAdviceHelper<T> {
-        void onExitStartAsync(T asyncContext);
-    }
+        @Override
+        public Class<?> getAdviceClass() {
+            return StartAsyncAdvice.class;
+        }
 
-    @VisibleForAdvice
-    public static class StartAsyncAdvice {
+        @VisibleForAdvice
+        public static class StartAsyncAdvice {
 
-        @Advice.OnMethodExit(suppress = Throwable.class)
-        private static void onExitStartAsync(@Advice.This HttpServletRequest request, @Advice.Return AsyncContext asyncContext) {
-            if (tracer != null) {
-                if (asyncHelper != null) {
-                    asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
+            @Advice.OnMethodExit(suppress = Throwable.class)
+            private static void onExitStartAsync(@Advice.This HttpServletRequest request, @Advice.Return AsyncContext asyncContext) {
+                if (tracer != null) {
+                    if (asyncHelper != null) {
+                        asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
+                    }
                 }
             }
         }
     }
 
+    public static class AsyncContextInstrumentation extends AsyncInstrumentation {
+        @Override
+        public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
+            return nameContains("AsyncContext");
+        }
+
+        @Override
+        public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+            return not(isInterface())
+                .and(hasSuperType(named("javax.servlet.AsyncContext")));
+        }
+
+        @Override
+        public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+            return isPublic()
+                .and(named("start"))
+                .and(takesArguments(Runnable.class));
+        }
+
+        @Override
+        public Class<?> getAdviceClass() {
+            return AsyncContextStartAdvice.class;
+        }
+
+        @VisibleForAdvice
+        public static class AsyncContextStartAdvice {
+
+            @Advice.OnMethodEnter(suppress = Throwable.class)
+            private static void onEnterAsyncContextStart(@Advice.Argument(value = 0, readOnly = false) @Nullable Runnable runnable) {
+                if (tracer != null) {
+                    final Transaction transaction = tracer.currentTransaction();
+                    if (transaction != null && runnable != null && asyncHelper != null) {
+                        runnable = tracer.wrapRunnable(runnable, transaction);
+                    }
+                }
+            }
+        }
+    }
 }

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/StartAsyncAdviceHelperImpl.java → apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java

@@ -31,14 +31,14 @@
 import static co.elastic.apm.agent.servlet.ServletTransactionHelper.ASYNC_ATTRIBUTE;
 import static co.elastic.apm.agent.servlet.ServletTransactionHelper.TRANSACTION_ATTRIBUTE;
 
-public class StartAsyncAdviceHelperImpl implements AsyncInstrumentation.StartAsyncAdviceHelper<AsyncContext> {
+public class AsyncContextAdviceHelperImpl implements AsyncInstrumentation.AsyncContextAdviceHelper<AsyncContext> {
 
     private static final String ASYNC_LISTENER_ADDED = ServletApiAdvice.class.getName() + ".asyncListenerAdded";
 
     private final ServletTransactionHelper servletTransactionHelper;
     private final ElasticApmTracer tracer;
 
-    public StartAsyncAdviceHelperImpl(ElasticApmTracer tracer) {
+    public AsyncContextAdviceHelperImpl(ElasticApmTracer tracer) {
         this.tracer = tracer;
         servletTransactionHelper = new ServletTransactionHelper(tracer);
     }

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/package-info.java

@@ -48,13 +48,13 @@
  * </p>
  *
  * <pre>
- *          System/Bootstrap CL (Agent)     provides interface StartAsyncAdviceHelper
+ *          System/Bootstrap CL (Agent)     provides interface AsyncContextAdviceHelper
  *                     |
  *                     v
- *             App CL (Servlet API)         uses StartAsyncAdviceHelperImpl from Helper CL
+ *             App CL (Servlet API)         uses AsyncContextAdviceHelperImpl from Helper CL
  *               /              \
  *              v               v
- * WebApp CL (user code/libs)   Helper CL   implements StartAsyncAdviceHelperImpl
+ * WebApp CL (user code/libs)   Helper CL   implements AsyncContextAdviceHelperImpl
  * </pre>
  */
 @NonnullApi

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

@@ -1,3 +1,4 @@
 co.elastic.apm.agent.servlet.ServletInstrumentation
 co.elastic.apm.agent.servlet.FilterChainInstrumentation
-co.elastic.apm.agent.servlet.AsyncInstrumentation
+co.elastic.apm.agent.servlet.AsyncInstrumentation$StartAsyncInstrumentation
+co.elastic.apm.agent.servlet.AsyncInstrumentation$AsyncContextInstrumentation