fix #141

unwrap ttl runnable for before/afterExecute methods
of executor subclass when decorate by ttl agent

- setAutoWrapper for TtlRunnable/TtlCallable
  in TtlExecutorTransformlet(Agent)
- update before/afterExecute methods of executor subclasses,
  unwrap runnable if isAutoWrapper

src/main/java/com/alibaba/ttl/spi/TtlAttachmentsDelegate.java

 package com.alibaba.ttl.spi;
 
-import javax.annotation.Nonnull;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
 public class TtlAttachmentsDelegate implements TtlAttachments {
-    private volatile ConcurrentMap<String, Object> attachment = new ConcurrentHashMap<String, Object>();
+    private volatile ConcurrentMap<String, Object> attachments = new ConcurrentHashMap<String, Object>();
 
     @Override
     public void setTtlAttachment(String key, Object value) {
-        attachment.put(key, value);
+        attachments.put(key, value);
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public <T> T getTtlAttachment(String key) {
-        return (T) attachment.get(key);
-    }
-
-    public static void setAutoWrapper(Object ttlAttachment) {
-        if (!(ttlAttachment instanceof TtlAttachments)) return;
-        ((TtlAttachments) ttlAttachment).setTtlAttachment(TtlAttachments.KEY_IS_AUTO_WRAPPER, true);
-    }
-
-    public static boolean isAutoWrapper(@Nonnull TtlAttachments ttlAttachments) {
-        return ttlAttachments.getTtlAttachment(TtlAttachments.KEY_IS_AUTO_WRAPPER);
+        return (T) attachments.get(key);
     }
 }

src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlExecutorTransformlet.java

 package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;
 
+import com.alibaba.ttl.threadpool.TtlExecutors;
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 import javassist.*;
@@ -10,6 +11,7 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.Callable;
 
 import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.getCtClass;
 import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.signatureOfMethod;
@@ -32,11 +34,14 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
     private static Set<String> EXECUTOR_CLASS_NAMES = new HashSet<String>();
     private static final Map<String, String> PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS = new HashMap<String, String>();
 
+    private static final String THREAD_POOL_EXECUTOR_CLASS_NAME = "java.util.concurrent.ThreadPoolExecutor";
+    private static final String RUNNABLE_CLASS_NAME = "java.lang.Runnable";
+
     static {
-        EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ThreadPoolExecutor");
+        EXECUTOR_CLASS_NAMES.add(THREAD_POOL_EXECUTOR_CLASS_NAME);
         EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ScheduledThreadPoolExecutor");
 
-        PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put("java.lang.Runnable", "com.alibaba.ttl.TtlRunnable");
+        PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put(RUNNABLE_CLASS_NAME, "com.alibaba.ttl.TtlRunnable");
         PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put("java.util.concurrent.Callable", "com.alibaba.ttl.TtlCallable");
     }
 
@@ -54,17 +59,34 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
             final CtClass clazz = getCtClass(classFileBuffer, loader);
 
             for (CtMethod method : clazz.getDeclaredMethods()) {
-                updateMethodOfExecutorClass(method);
+                updateSubmitMethodsOfExecutorClass_decorateToTtlWrapperAndSetAutoWrapper(method);
             }
 
             if (disableInheritable) updateConstructorDisableInheritable(clazz);
 
             return clazz.toBytecode();
+        } else {
+            final CtClass clazz = getCtClass(classFileBuffer, loader);
+
+            if (clazz.isPrimitive() || clazz.isArray() || clazz.isInterface() || clazz.isAnnotation()) {
+                return null;
+            }
+            if (!clazz.subclassOf(clazz.getClassPool().get(THREAD_POOL_EXECUTOR_CLASS_NAME))) return null;
+
+            logger.info("Transforming class " + className);
+
+            final boolean updated = updateBeforeAndAfterExecuteMethodOfExecutorSubclass(clazz);
+            if (updated) return clazz.toBytecode();
+            else return null;
         }
-        return null;
     }
 
-    private void updateMethodOfExecutorClass(final CtMethod method) throws NotFoundException, CannotCompileException {
+    /**
+     * @see com.alibaba.ttl.TtlRunnable#get(Runnable, boolean, boolean)
+     * @see com.alibaba.ttl.TtlCallable#get(Callable, boolean, boolean)
+     * @see com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils#setAutoWrapper(Object)
+     */
+    private void updateSubmitMethodsOfExecutorClass_decorateToTtlWrapperAndSetAutoWrapper(final CtMethod method) throws NotFoundException, CannotCompileException {
         final int modifiers = method.getModifiers();
         if (!Modifier.isPublic(modifiers) || Modifier.isStatic(modifiers)) return;
 
@@ -73,7 +95,12 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
         for (int i = 0; i < parameterTypes.length; i++) {
             final String paramTypeName = parameterTypes[i].getName();
             if (PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.containsKey(paramTypeName)) {
-                String code = String.format("$%d = %s.get($%d, false, true);", i + 1, PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.get(paramTypeName), i + 1);
+                String code = String.format(
+                        // decorate to TTL wrapper,
+                        // and then set AutoWrapper attachment/Tag
+                        "$%d = %s.get($%d, false, true);"
+                                + "\ncom.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.setAutoWrapper($%<d);",
+                        i + 1, PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.get(paramTypeName), i + 1);
                 logger.info("insert code before method " + signatureOfMethod(method) + " of class " + method.getDeclaringClass().getName() + ": " + code);
                 insertCode.append(code);
             }
@@ -81,6 +108,9 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
         if (insertCode.length() > 0) method.insertBefore(insertCode.toString());
     }
 
+    /**
+     * @see TtlExecutors#getDisableInheritableThreadFactory(java.util.concurrent.ThreadFactory)
+     */
     private void updateConstructorDisableInheritable(final CtClass clazz) throws NotFoundException, CannotCompileException {
         for (CtConstructor constructor : clazz.getDeclaredConstructors()) {
             final CtClass[] parameterTypes = constructor.getParameterTypes();
@@ -96,4 +126,38 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
             if (insertCode.length() > 0) constructor.insertBefore(insertCode.toString());
         }
     }
+
+    /**
+     * @see Utils#unwrapIfIsAutoWrapper(Runnable)
+     */
+    private boolean updateBeforeAndAfterExecuteMethodOfExecutorSubclass(final CtClass clazz) throws NotFoundException, CannotCompileException {
+        final CtClass runnableClass = clazz.getClassPool().get(RUNNABLE_CLASS_NAME);
+        final CtClass threadClass = clazz.getClassPool().get("java.lang.Thread");
+        final CtClass throwableClass = clazz.getClassPool().get("java.lang.Throwable");
+        boolean updated = false;
+
+        try {
+            final CtMethod beforeExecute = clazz.getDeclaredMethod("beforeExecute", new CtClass[]{threadClass, runnableClass});
+            // unwrap runnable if IsAutoWrapper
+            String code = "$2 = com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.unwrapIfIsAutoWrapper($2);";
+            logger.info("insert code before method " + signatureOfMethod(beforeExecute) + " of class " + beforeExecute.getDeclaringClass().getName() + ": " + code);
+            beforeExecute.insertBefore(code);
+            updated = true;
+        } catch (NotFoundException e) {
+            // clazz does not override beforeExecute method, do nothing.
+        }
+
+        try {
+            final CtMethod afterExecute = clazz.getDeclaredMethod("afterExecute", new CtClass[]{runnableClass, throwableClass});
+            // unwrap runnable if IsAutoWrapper
+            String code = "$1 = com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.unwrapIfIsAutoWrapper($1);";
+            logger.info("insert code before method " + signatureOfMethod(afterExecute) + " of class " + afterExecute.getDeclaringClass().getName() + ": " + code);
+            afterExecute.insertBefore(code);
+            updated = true;
+        } catch (NotFoundException e) {
+            // clazz does not override afterExecute method, do nothing.
+        }
+
+        return updated;
+    }
 }

src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlForkJoinTransformlet.java

@@ -47,6 +47,9 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
         return null;
     }
 
+    /**
+     * @see Utils#doCaptureWhenNotTtlEnhanced(java.lang.Object)
+     */
     private void updateForkJoinTaskClass(final CtClass clazz) throws CannotCompileException, NotFoundException {
         final String className = clazz.getName();
 

src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlTimerTaskTransformlet.java

@@ -47,6 +47,9 @@ public class TtlTimerTaskTransformlet implements JavassistTransformlet {
         return clazz.toBytecode();
     }
 
+    /**
+     * @see Utils#doCaptureWhenNotTtlEnhanced(java.lang.Object)
+     */
     private void updateTimerTaskClass(final CtClass clazz) throws CannotCompileException, NotFoundException {
         final String className = clazz.getName();
 

src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/Utils.java

 package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;
 
 import com.alibaba.ttl.TransmittableThreadLocal;
+import com.alibaba.ttl.TtlRunnable;
+import com.alibaba.ttl.spi.TtlAttachments;
 import com.alibaba.ttl.spi.TtlEnhanced;
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
 import javassist.*;
@@ -90,10 +92,27 @@ public class Utils {
         logger.info("insert code around method " + signatureOfMethod(method) + " of class " + clazz.getName() + ": " + code);
     }
 
-    @SuppressWarnings("unused")
     public static Object doCaptureWhenNotTtlEnhanced(Object obj) {
         if (obj instanceof TtlEnhanced) return null;
         else return TransmittableThreadLocal.Transmitter.capture();
     }
 
+    public static void setAutoWrapper(Object ttlAttachment) {
+        if (notTtlAttachments(ttlAttachment)) return;
+        ((TtlAttachments) ttlAttachment).setTtlAttachment(TtlAttachments.KEY_IS_AUTO_WRAPPER, true);
+    }
+
+    public static Runnable unwrapIfIsAutoWrapper(Runnable runnable) {
+        if (notTtlAttachments(runnable)) return runnable;
+        else if (isAutoWrapper(runnable)) return TtlRunnable.unwrap(runnable);
+        else return runnable;
+    }
+
+    private static boolean notTtlAttachments(Object ttlAttachment) {
+        return !(ttlAttachment instanceof TtlAttachments);
+    }
+
+    private static boolean isAutoWrapper(Runnable ttlAttachments) {
+        return ((TtlAttachments) ttlAttachments).getTtlAttachment(TtlAttachments.KEY_IS_AUTO_WRAPPER);
+    }
 }