Async Servlet state management changes (#385)

Author: eyalkoren

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

@@ -1,3 +1,22 @@
+/*-
+ * #%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.api;
 
 import javax.annotation.Nonnull;

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

@@ -23,7 +23,6 @@
 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;
@@ -38,7 +37,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 
-import static co.elastic.apm.agent.servlet.ServletApiAdvice.TRANSACTION_ATTRIBUTE;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -126,12 +124,8 @@ public static class StartAsyncAdvice {
         @Advice.OnMethodExit(suppress = Throwable.class)
         private static void onExitStartAsync(@Advice.This HttpServletRequest request, @Advice.Return AsyncContext asyncContext) {
             if (tracer != null) {
-                final Transaction transaction = tracer.currentTransaction();
-                if (transaction != null) {
-                    request.setAttribute(TRANSACTION_ATTRIBUTE, transaction);
-                    if (asyncHelper != null) {
-                        asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
-                    }
+                if (asyncHelper != null) {
+                    asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
                 }
             }
         }

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

@@ -41,6 +41,8 @@
 import java.util.Enumeration;
 import java.util.Map;
 
+import static co.elastic.apm.agent.servlet.ServletTransactionHelper.TRANSACTION_ATTRIBUTE;
+
 /**
  * Only the methods annotated with {@link Advice.OnMethodEnter} and {@link Advice.OnMethodExit} may contain references to
  * {@code javax.servlet}, as these are inlined into the matching methods.
@@ -49,8 +51,6 @@
  */
 public class ServletApiAdvice {
 
-    @VisibleForAdvice
-    public static final String TRANSACTION_ATTRIBUTE = ServletApiAdvice.class.getName() + ".transaction";
     @Nullable
     @VisibleForAdvice
     public static ServletTransactionHelper servletTransactionHelper;
@@ -149,8 +149,9 @@ public static void onExitServletService(@Advice.Argument(0) ServletRequest servl
             servletResponse instanceof HttpServletResponse) {
 
             final HttpServletRequest request = (HttpServletRequest) servletRequest;
-            if (request.isAsyncStarted()) {
-                // the response is not ready yet; the request is handled asynchronously
+            if (request.getAttribute(ServletTransactionHelper.ASYNC_ATTRIBUTE) != null) {
+                // HttpServletRequest.startAsync was invoked on this request.
+                // The transaction should be handled from now on by the other thread committing the response
                 transaction.deactivate();
             } else {
                 // this is not an async request, so we can end the transaction immediately
@@ -170,7 +171,6 @@ public static void onExitServletService(@Advice.Argument(0) ServletRequest servl
                 } else {
                     parameterMap = null;
                 }
-                request.removeAttribute(TRANSACTION_ATTRIBUTE);
                 servletTransactionHelper.onAfter(transaction, t, response.isCommitted(), response.getStatus(), request.getMethod(),
                     parameterMap, request.getServletPath(), request.getPathInfo(), contentTypeHeader);
             }

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

@@ -51,6 +51,11 @@
 @VisibleForAdvice
 public class ServletTransactionHelper {
 
+    @VisibleForAdvice
+    public static final String TRANSACTION_ATTRIBUTE = ServletApiAdvice.class.getName() + ".transaction";
+    @VisibleForAdvice
+    public static final String ASYNC_ATTRIBUTE = ServletApiAdvice.class.getName() + ".async";
+
     private final Logger logger = LoggerFactory.getLogger(ServletTransactionHelper.class);
 
     private final Set<String> METHODS_WITH_BODY = new HashSet<>(Arrays.asList("POST", "PUT", "PATCH", "DELETE"));

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

@@ -29,14 +29,20 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
+import static co.elastic.apm.agent.servlet.ServletTransactionHelper.TRANSACTION_ATTRIBUTE;
 
 /**
  * Based on brave.servlet.ServletRuntime$TracingAsyncListener (under Apache license 2.0)
  */
 public class ApmAsyncListener implements AsyncListener {
+    private static final AtomicIntegerFieldUpdater<ApmAsyncListener> EVENT_COUNTER_UPDATER =
+        AtomicIntegerFieldUpdater.newUpdater(ApmAsyncListener.class, "endEventCounter");
+
     private final ServletTransactionHelper servletTransactionHelper;
     private final Transaction transaction;
-    private volatile boolean completed = false;
+    private volatile int endEventCounter = 0;
 
     ApmAsyncListener(ServletTransactionHelper servletTransactionHelper, Transaction transaction) {
         this.servletTransactionHelper = servletTransactionHelper;
@@ -45,26 +51,17 @@ public class ApmAsyncListener implements AsyncListener {
 
     @Override
     public void onComplete(AsyncEvent event) {
-        if (!completed) {
-            endTransaction(event);
-            completed = true;
-        }
+        endTransaction(event);
     }
 
     @Override
     public void onTimeout(AsyncEvent event) {
-        if (!completed) {
-            endTransaction(event);
-            completed = true;
-        }
+        endTransaction(event);
     }
 
     @Override
     public void onError(AsyncEvent event) {
-        if (!completed) {
-            endTransaction(event);
-            completed = true;
-        }
+        endTransaction(event);
     }
 
     /**
@@ -82,7 +79,14 @@ public void onStartAsync(AsyncEvent event) {
     // because only the onExitServletService method may contain references to the servlet API
     // (see class-level Javadoc)
     private void endTransaction(AsyncEvent event) {
+        // To ensure transaction is ended only by a single event
+        if (EVENT_COUNTER_UPDATER.getAndIncrement(this) > 0) {
+            return;
+        }
+
         HttpServletRequest request = (HttpServletRequest) event.getSuppliedRequest();
+        request.removeAttribute(TRANSACTION_ATTRIBUTE);
+
         HttpServletResponse response = (HttpServletResponse) event.getSuppliedResponse();
         final Response resp = transaction.getContext().getResponse();
         if (transaction.isSampled() && servletTransactionHelper.isCaptureHeaders()) {

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

@@ -28,6 +28,9 @@
 import javax.servlet.AsyncContext;
 import javax.servlet.ServletRequest;
 
+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> {
 
     private static final String ASYNC_LISTENER_ADDED = ServletApiAdvice.class.getName() + ".asyncListenerAdded";
@@ -47,9 +50,7 @@ public void onExitStartAsync(AsyncContext asyncContext) {
             return;
         }
         final Transaction transaction = tracer.currentTransaction();
-        if (transaction != null &&
-            transaction.isSampled() &&
-            request.getAttribute(ASYNC_LISTENER_ADDED) == null) {
+        if (transaction != null && transaction.isSampled() && request.getAttribute(ASYNC_LISTENER_ADDED) == null) {
             // makes sure that the listener is only added once, even if the request is wrapped
             // which leads to multiple invocations of startAsync for the same underlying request
             request.setAttribute(ASYNC_LISTENER_ADDED, Boolean.TRUE);
@@ -58,6 +59,9 @@ public void onExitStartAsync(AsyncContext asyncContext) {
             // however, only some application server like WebSphere actually implement it that way
             asyncContext.addListener(new ApmAsyncListener(servletTransactionHelper, transaction),
                 asyncContext.getRequest(), asyncContext.getResponse());
+
+            request.setAttribute(ASYNC_ATTRIBUTE, Boolean.TRUE);
+            request.setAttribute(TRANSACTION_ATTRIBUTE, transaction);
         }
     }
 }