Support JAX-RS @Suspended and AsyncResponse

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,23 +165,24 @@ 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 asyncContext = context.createAuxThreadContext();
-            RunnableCallableMixin innerRunnableCallable =
-                    futureTaskMixin.glowroot$getInnerRunnableCallable();
-            if (innerRunnableCallable != null) {
-                innerRunnableCallable.glowroot$setAuxAsyncContext(asyncContext);
+            if (context.isTransactionAsync()) {
+                RunnableCallableMixin runnableCallableMixin = (RunnableCallableMixin) runnableCallable;
+                AuxThreadContext asyncContext = context.createAuxThreadContext();
+                runnableCallableMixin.glowroot$setAuxAsyncContext(asyncContext);
             }
         }
     }
 
+
     // this method uses submit() and returns Future, but none of the callers use/wait on the Future
     @Pointcut(className = "net.sf.ehcache.store.disk.DiskStorageFactory", methodName = "schedule",
             methodParameterTypes = {"java.util.concurrent.Callable"}, nestingGroup = "executor")

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

@@ -17,23 +17,10 @@
 
 import javax.annotation.Nullable;
 
-import org.glowroot.agent.plugin.api.Agent;
-import org.glowroot.agent.plugin.api.MessageSupplier;
-import org.glowroot.agent.plugin.api.ThreadContext;
+import org.glowroot.agent.plugin.api.*;
 import org.glowroot.agent.plugin.api.ThreadContext.Priority;
-import org.glowroot.agent.plugin.api.TimerName;
-import org.glowroot.agent.plugin.api.TraceEntry;
 import org.glowroot.agent.plugin.api.config.BooleanProperty;
-import org.glowroot.agent.plugin.api.weaving.BindMethodMeta;
-import org.glowroot.agent.plugin.api.weaving.BindParameter;
-import org.glowroot.agent.plugin.api.weaving.BindThrowable;
-import org.glowroot.agent.plugin.api.weaving.BindTraveler;
-import org.glowroot.agent.plugin.api.weaving.OnAfter;
-import org.glowroot.agent.plugin.api.weaving.OnBefore;
-import org.glowroot.agent.plugin.api.weaving.OnReturn;
-import org.glowroot.agent.plugin.api.weaving.OnThrow;
-import org.glowroot.agent.plugin.api.weaving.Pointcut;
-import org.glowroot.agent.plugin.api.weaving.Shim;
+import org.glowroot.agent.plugin.api.weaving.*;
 
 // TODO optimize away servletPath thread local, e.g. store servlet path in thread context via
 // servlet plugin and retrieve here
@@ -106,6 +93,11 @@ public static class ResourceAdvice {
         @OnBefore
         public static TraceEntry onBefore(ThreadContext context,
                 @BindMethodMeta ResourceMethodMeta resourceMethodMeta) {
+
+            if (resourceMethodMeta.isAsynchronous()) {
+                context.setTransactionAsync();
+            }
+
             if (useAltTransactionNaming.value()) {
                 context.setTransactionName(resourceMethodMeta.getAltTransactionName(),
                         Priority.CORE_PLUGIN);
@@ -120,7 +112,7 @@ public static TraceEntry onBefore(ThreadContext context,
                     timerName);
         }
         @OnReturn
-        public static void onReturn(@BindTraveler TraceEntry traceEntry) {
+        public static void onReturn(ThreadContext context, @BindTraveler TraceEntry traceEntry) {
             traceEntry.end();
         }
         @OnThrow
@@ -146,8 +138,128 @@ 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);
+        }
+
+        @OnAfter
+        public static void onAfter(ThreadContext context, @BindTraveler TraceEntry traceEntry) {
+            context.completeAsyncTransaction();
+            traceEntry.end();
+        }
+
+        @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;
     }
+
+
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    //
+    // Put this here for test purposes. Should be removed, and there should be a test that covers two plugins.
+    // the method names are verbose to avoid conflict since they will become methods in all classes
+    // that extend java.lang.Runnable and/or java.util.concurrent.Callable
+    public interface RunnableCallableMixin {
+
+        @Nullable
+        AuxThreadContext glowroot$getAuxThreadContext();
+
+        void glowroot$setAuxAsyncContext(@Nullable AuxThreadContext auxThreadContext);
+    }
+
+    // the field and method names are verbose to avoid conflict since they will become fields
+    // and methods in all classes that extend java.lang.Runnable and/or
+    // java.util.concurrent.Callable
+    @Mixin({"java.lang.Runnable", "java.util.concurrent.Callable"})
+    public abstract static class RunnableImpl implements RunnableCallableMixin {
+
+        private volatile @Nullable AuxThreadContext glowroot$auxThreadContext;
+
+        @Override
+        public @Nullable AuxThreadContext glowroot$getAuxThreadContext() {
+            return glowroot$auxThreadContext;
+        }
+
+        @Override
+        public void glowroot$setAuxAsyncContext(@Nullable AuxThreadContext auxThreadContext) {
+            this.glowroot$auxThreadContext = auxThreadContext;
+        }
+    }
+
+
+    @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) {
+            // 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) {
+            if (context.isTransactionAsync()) {
+                RunnableCallableMixin runnableCallableMixin = (RunnableCallableMixin) runnableCallable;
+                AuxThreadContext asyncContext = context.createAuxThreadContext();
+                runnableCallableMixin.glowroot$setAuxAsyncContext(asyncContext);
+            }
+        }
+    }
+
+    @Pointcut(className = "java.lang.Runnable", methodName = "run", methodParameterTypes = {})
+    public static class RunnableAdvice {
+        @OnBefore
+        public static @Nullable TraceEntry onBefore(@BindReceiver Runnable runnable) {
+            if (!(runnable instanceof RunnableCallableMixin)) {
+                // this class was loaded before class file transformer was added to jvm
+                return null;
+            }
+            RunnableCallableMixin runnableMixin = (RunnableCallableMixin) runnable;
+            AuxThreadContext asyncContext = runnableMixin.glowroot$getAuxThreadContext();
+            if (asyncContext == null) {
+                return null;
+            }
+            return asyncContext.start();
+        }
+        @OnReturn
+        public static void onReturn(@BindTraveler @Nullable TraceEntry traceEntry) {
+            if (traceEntry != null) {
+                traceEntry.end();
+            }
+        }
+        @OnThrow
+        public static void onThrow(@BindThrowable Throwable t,
+                                   @BindTraveler @Nullable TraceEntry traceEntry) {
+            if (traceEntry != null) {
+                traceEntry.endWithError(t);
+            }
+        }
+        @OnAfter
+        public static void onAfter(@BindReceiver Runnable runnable) {
+            if (!(runnable instanceof RunnableCallableMixin)) {
+                // this class was loaded before class file transformer was added to jvm
+                return;
+            }
+            RunnableCallableMixin runnableMixin = (RunnableCallableMixin) runnable;
+            runnableMixin.glowroot$setAuxAsyncContext(null);
+        }
+    }
+
+
+
 }

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 asynchronous;
+
     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;
+        asynchronous = containsSuspended(method.getParameterAnnotations());
     }
 
     String getResourceClassName() {
@@ -61,7 +64,13 @@ String getAltTransactionName() {
         return altTransactionName;
     }
 
-    private static @Nullable String getPath(AnnotatedElement annotatedElement) {
+    boolean isAsynchronous() {
+        return asynchronous;
+    }
+
+    private static
+    @Nullable
+    String getPath(AnnotatedElement annotatedElement) {
         try {
             for (Annotation annotation : annotatedElement.getDeclaredAnnotations()) {
                 Class<?> annotationClass = annotation.annotationType();
@@ -75,8 +84,10 @@ String getAltTransactionName() {
         return null;
     }
 
-    private static @Nullable String getPathAttribute(Class<?> pathClass, Object path,
-            String attributeName) throws Exception {
+    private static
+    @Nullable
+    String getPathAttribute(Class<?> pathClass, Object path,
+                            String attributeName) throws Exception {
         Method method = pathClass.getMethod(attributeName);
         return (String) method.invoke(path);
     }
@@ -107,4 +118,15 @@ private static String normalize(@Nullable String path) {
     private static String replacePathSegmentsWithAsterisk(String path) {
         return path.replaceAll("\\{[^}]*\\}", "*");
     }
+
+    private static boolean containsSuspended(Annotation[][] parameterAnnotations) {
+        for (int i = 0; i < parameterAnnotations.length; i++) {
+            for (int j = 0; j < parameterAnnotations[i].length; j++) {
+                if (parameterAnnotations[i][j].annotationType().getName().equals("javax.ws.rs.container.Suspended")) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
 }