Support JAX-RS @suspended and AsyncResponse

Closes #83

agent-parent/agent/src/main/java/org/glowroot/agent/model/OptionalThreadContextImpl.java

@@ -165,6 +165,14 @@ public void setTransactionAsync() {
         }
     }
 
+    @Override
+    public boolean isTransactionAsync() {
+        if (threadContext != null) {
+            return threadContext.isTransactionAsync();
+        }
+        return false;
+    }
+
     @Override
     public void completeAsyncTransaction() {
         if (threadContext != null) {

agent-parent/agent/src/main/java/org/glowroot/agent/model/ThreadContextImpl.java

@@ -808,6 +808,11 @@ public void setTransactionAsync() {
         transaction.setAsync();
     }
 
+    @Override
+    public boolean isTransactionAsync() {
+        return transaction.isAsync();
+    }
+
     @Override
     public void completeAsyncTransaction() {
         completeAsyncTransaction = true;

agent-parent/plugin-api/src/main/java/org/glowroot/agent/plugin/api/ThreadContext.java

@@ -87,6 +87,8 @@ AsyncTraceEntry startAsyncServiceCallEntry(String type, String text,
 
     void setTransactionAsync();
 
+    boolean isTransactionAsync();
+
     void completeAsyncTransaction();
 
     /**

agent-parent/plugins/executor-plugin/src/main/java/org/glowroot/agent/plugin/executor/ExecutorAspect.java

@@ -165,19 +165,28 @@ public static void onBefore(@BindReceiver Object futureTask,
     @Pointcut(className = "java.util.concurrent.Executor", methodName = "execute",
             methodParameterTypes = {"java.lang.Runnable"}, nestingGroup = "executor")
     public static class ExecuteAdvice {
+
         @IsEnabled
         public static boolean isEnabled(@BindParameter Object runnableCallable) {
-            // only capture execute if called on FutureTask
-            return runnableCallable instanceof FutureTaskMixin;
+            // this class may have been loaded before class file transformer was added to jvm
+            return runnableCallable instanceof RunnableCallableMixin;
         }
+
         @OnBefore
         public static void onBefore(ThreadContext context, @BindParameter Object runnableCallable) {
-            FutureTaskMixin futureTaskMixin = (FutureTaskMixin) runnableCallable;
-            AuxThreadContext auxContext = context.createAuxThreadContext();
-            RunnableCallableMixin innerRunnableCallable =
-                    futureTaskMixin.glowroot$getInnerRunnableCallable();
-            if (innerRunnableCallable != null) {
-                innerRunnableCallable.glowroot$setAuxContext(auxContext);
+            if (runnableCallable instanceof FutureTaskMixin) {
+                FutureTaskMixin futureTaskMixin = (FutureTaskMixin) runnableCallable;
+                AuxThreadContext auxContext = context.createAuxThreadContext();
+                RunnableCallableMixin innerRunnableCallable =
+                        futureTaskMixin.glowroot$getInnerRunnableCallable();
+                if (innerRunnableCallable != null) {
+                    innerRunnableCallable.glowroot$setAuxContext(auxContext);
+                }
+            } else if (context.isTransactionAsync()) {
+                RunnableCallableMixin runnableCallableMixin =
+                        (RunnableCallableMixin) runnableCallable;
+                AuxThreadContext asyncContext = context.createAuxThreadContext();
+                runnableCallableMixin.glowroot$setAuxContext(asyncContext);
             }
         }
     }

agent-parent/plugins/jaxrs-plugin/pom.xml

@@ -76,6 +76,13 @@
       <version>2.22.2</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <!-- executor plugin is included to test async servlets -->
+      <groupId>org.glowroot</groupId>
+      <artifactId>glowroot-agent-executor-plugin</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

agent-parent/plugins/jaxrs-plugin/src/main/java/org/glowroot/agent/plugin/jaxrs/ResourceAspect.java

@@ -106,6 +106,11 @@ public static class ResourceAdvice {
         @OnBefore
         public static TraceEntry onBefore(ThreadContext context,
                 @BindMethodMeta ResourceMethodMeta resourceMethodMeta) {
+
+            if (resourceMethodMeta.isAsync()) {
+                context.setTransactionAsync();
+            }
+
             if (useAltTransactionNaming.value()) {
                 context.setTransactionName(resourceMethodMeta.getAltTransactionName(),
                         Priority.CORE_PLUGIN);
@@ -146,6 +151,26 @@ private static String getTransactionName(@Nullable String method,
         }
     }
 
+    @Pointcut(className = "javax.ws.rs.container.AsyncResponse", methodName = "resume",
+            methodParameterTypes = {".."}, timerName = "jaxrs async response")
+    public static class AsyncResponseAdvice {
+        private static final TimerName timerName = Agent.getTimerName(AsyncResponseAdvice.class);
+        @OnBefore
+        public static TraceEntry onBefore(ThreadContext context) {
+            return context.startTraceEntry(MessageSupplier.from("jaxrs async response"), timerName);
+        }
+        @OnReturn
+        public static void onReturn(ThreadContext context, @BindTraveler TraceEntry traceEntry) {
+            traceEntry.end();
+            context.completeAsyncTransaction();
+        }
+        @OnThrow
+        public static void onThrow(@BindThrowable Throwable throwable,
+                @BindTraveler TraceEntry traceEntry) {
+            traceEntry.endWithError(throwable);
+        }
+    }
+
     private static class RequestInfo {
         private @Nullable String method;
         private @Nullable String servletPath;

agent-parent/plugins/jaxrs-plugin/src/main/java/org/glowroot/agent/plugin/jaxrs/ResourceMethodMeta.java

@@ -35,6 +35,8 @@ public class ResourceMethodMeta {
 
     private final String altTransactionName;
 
+    private final boolean async;
+
     public ResourceMethodMeta(Method method) {
         Class<?> resourceClass = method.getDeclaringClass();
         resourceClassName = resourceClass.getName();
@@ -43,6 +45,7 @@ public ResourceMethodMeta(Method method) {
         String methodPath = getPath(method);
         path = combine(classPath, methodPath);
         altTransactionName = resourceClass.getSimpleName() + "#" + methodName;
+        async = containsSuspended(method.getParameterAnnotations());
     }
 
     String getResourceClassName() {
@@ -61,6 +64,10 @@ String getAltTransactionName() {
         return altTransactionName;
     }
 
+    boolean isAsync() {
+        return async;
+    }
+
     private static @Nullable String getPath(AnnotatedElement annotatedElement) {
         try {
             for (Annotation annotation : annotatedElement.getDeclaredAnnotations()) {
@@ -107,4 +114,16 @@ private static String normalize(@Nullable String path) {
     private static String replacePathSegmentsWithAsterisk(String path) {
         return path.replaceAll("\\{[^}]*\\}", "*");
     }
+
+    private static boolean containsSuspended(Annotation[][] parameterAnnotations) {
+        for (Annotation[] annotationsForOneParameter : parameterAnnotations) {
+            for (Annotation annotation : annotationsForOneParameter) {
+                if (annotation.annotationType().getName()
+                        .equals("javax.ws.rs.container.Suspended")) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
 }

agent-parent/plugins/jaxrs-plugin/src/test/java/org/glowroot/agent/plugin/jaxrs/InvokeJaxrsResourceInTomcat.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2016 the original author or authors.
+ *
+ * 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.
+ */
+package org.glowroot.agent.plugin.jaxrs;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.ServerSocket;
+
+import com.ning.http.client.AsyncHttpClient;
+import org.apache.catalina.Context;
+import org.apache.catalina.loader.WebappLoader;
+import org.apache.catalina.startup.Tomcat;
+
+import org.glowroot.agent.it.harness.AppUnderTest;
+
+abstract class InvokeJaxrsResourceInTomcat implements AppUnderTest {
+
+    public void executeApp(String webapp, String url) throws Exception {
+        int port = getAvailablePort();
+        Tomcat tomcat = new Tomcat();
+        tomcat.setBaseDir("target/tomcat");
+        tomcat.setPort(port);
+        Context context = tomcat.addWebapp("",
+                new File("src/test/resources/" + webapp).getAbsolutePath());
+
+        WebappLoader webappLoader =
+                new WebappLoader(InvokeJaxrsResourceInTomcat.class.getClassLoader());
+        context.setLoader(webappLoader);
+
+        tomcat.start();
+        AsyncHttpClient asyncHttpClient = new AsyncHttpClient();
+        asyncHttpClient.prepareGet("http://localhost:" + port + url).execute().get();
+        asyncHttpClient.close();
+
+        tomcat.stop();
+        tomcat.destroy();
+    }
+
+    private static int getAvailablePort() throws IOException {
+        ServerSocket serverSocket = new ServerSocket(0);
+        int port = serverSocket.getLocalPort();
+        serverSocket.close();
+        return port;
+    }
+}

agent-parent/plugins/jaxrs-plugin/src/test/java/org/glowroot/agent/plugin/jaxrs/ResourceIT.java

@@ -15,26 +15,18 @@
  */
 package org.glowroot.agent.plugin.jaxrs;
 
-import java.io.File;
-import java.io.IOException;
-import java.net.ServerSocket;
 import java.util.List;
 
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.core.Response;
 
-import com.ning.http.client.AsyncHttpClient;
-import org.apache.catalina.Context;
-import org.apache.catalina.loader.WebappLoader;
-import org.apache.catalina.startup.Tomcat;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
-import org.glowroot.agent.it.harness.AppUnderTest;
 import org.glowroot.agent.it.harness.Container;
 import org.glowroot.agent.it.harness.Containers;
 import org.glowroot.wire.api.model.TraceOuterClass.Trace;
@@ -205,36 +197,6 @@ public void executeApp() throws Exception {
         }
     }
 
-    private static abstract class InvokeJaxrsResourceInTomcat implements AppUnderTest {
-
-        public void executeApp(String webapp, String url) throws Exception {
-            int port = getAvailablePort();
-            Tomcat tomcat = new Tomcat();
-            tomcat.setBaseDir("target/tomcat");
-            tomcat.setPort(port);
-            Context context = tomcat.addWebapp("",
-                    new File("src/test/resources/" + webapp).getAbsolutePath());
-
-            WebappLoader webappLoader =
-                    new WebappLoader(InvokeJaxrsResourceInTomcat.class.getClassLoader());
-            context.setLoader(webappLoader);
-
-            tomcat.start();
-            AsyncHttpClient asyncHttpClient = new AsyncHttpClient();
-            asyncHttpClient.prepareGet("http://localhost:" + port + url).execute().get();
-            asyncHttpClient.close();
-            tomcat.stop();
-            tomcat.destroy();
-        }
-
-        private static int getAvailablePort() throws IOException {
-            ServerSocket serverSocket = new ServerSocket(0);
-            int port = serverSocket.getLocalPort();
-            serverSocket.close();
-            return port;
-        }
-    }
-
     @Path("hello")
     public static class HelloResource {
         @Path("{param}")

agent-parent/plugins/jaxrs-plugin/src/test/java/org/glowroot/agent/plugin/jaxrs/ResourceMethodMetaTest.java

@@ -15,6 +15,9 @@
  */
 package org.glowroot.agent.plugin.jaxrs;
 
+import javax.ws.rs.container.AsyncResponse;
+import javax.ws.rs.container.Suspended;
+
 import org.junit.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -38,4 +41,14 @@ public void should() {
         assertThat(ResourceMethodMeta.combine("/abc", "")).isEqualTo("/abc");
         assertThat(ResourceMethodMeta.combine("", "/xyz")).isEqualTo("/xyz");
     }
+
+    @Test
+    public void shouldReturnAsync() throws NoSuchMethodException {
+        ResourceMethodMeta meta = new ResourceMethodMeta(ResourceMethodMetaTest.class
+                .getMethod("methodContainingSuspended", AsyncResponse.class));
+        assertThat(meta.isAsync()).isTrue();
+    }
+
+    public void methodContainingSuspended(
+            @SuppressWarnings("unused") @Suspended AsyncResponse response) {}
 }