fix #75

alibaba · Jun 24, 2017 · cd84835 · cd84835
1 parent e8b2f3d
commit cd84835
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 26 deletions.
diff --git a/src/main/java/com/alibaba/ttl/TtlCallable.java b/src/main/java/com/alibaba/ttl/TtlCallable.java
@@ -29,7 +29,6 @@ public final class TtlCallable<V> implements Callable<V> {
     private final Callable<V> callable;
     private final boolean releaseTtlValueReferenceAfterCall;
 
-
     private TtlCallable(Callable<V> callable, boolean releaseTtlValueReferenceAfterCall) {
         this.copiedRef = new AtomicReference<Map<TransmittableThreadLocal<?>, Object>>(TransmittableThreadLocal.copy());
         this.callable = callable;
@@ -58,6 +57,26 @@ public Callable<V> getCallable() {
         return callable;
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        TtlCallable<?> that = (TtlCallable<?>) o;
+
+        return callable.equals(that.callable);
+    }
+
+    @Override
+    public int hashCode() {
+        return callable.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return this.getClass().getName() + " - " + callable.toString();
+    }
+
     /**
      * Factory method, wrapper input {@link Callable} to {@link TtlCallable}.
      * <p>

diff --git a/src/main/java/com/alibaba/ttl/TtlRunnable.java b/src/main/java/com/alibaba/ttl/TtlRunnable.java
@@ -57,6 +57,26 @@ public Runnable getRunnable() {
         return runnable;
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        TtlRunnable that = (TtlRunnable) o;
+
+        return runnable.equals(that.runnable);
+    }
+
+    @Override
+    public int hashCode() {
+        return runnable.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return this.getClass().getName() + " - " + runnable.toString();
+    }
+
     /**
      * Factory method, wrapper input {@link Runnable} to {@link TtlRunnable}.
      *

diff --git a/src/main/java/com/alibaba/ttl/threadpool/agent/TtlTransformer.java b/src/main/java/com/alibaba/ttl/threadpool/agent/TtlTransformer.java
@@ -35,42 +35,44 @@ public class TtlTransformer implements ClassFileTransformer {
     private static final String TTL_RUNNABLE_CLASS_NAME = TtlRunnable.class.getName();
     private static final String TTL_CALLABLE_CLASS_NAME = TtlCallable.class.getName();
 
-    private static final String THREAD_POOL_CLASS_FILE = "java.util.concurrent.ThreadPoolExecutor".replace('.', '/');
-    private static final String SCHEDULER_CLASS_FILE = "java.util.concurrent.ScheduledThreadPoolExecutor".replace('.', '/');
 
-    private static final String TIMER_TASK_CLASS_FILE = "java.util.TimerTask".replace('.', '/');
-    private static final byte[] EMPTY_BYTE_ARRAY = {};
+    private static Set<String> EXECUTOR_CLASS_NAMES = new HashSet<String>();
 
-    private static String toClassName(String classFile) {
-        return classFile.replace('/', '.');
+    static {
+        EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ThreadPoolExecutor");
+        EXECUTOR_CLASS_NAMES.add("java.util.concurrent.ScheduledThreadPoolExecutor");
     }
 
+    private static final String TIMER_TASK_CLASS_NAME = "java.util.TimerTask";
+
+    private static final byte[] EMPTY_BYTE_ARRAY = {};
+
     @Override
     public byte[] transform(ClassLoader loader, String classFile, Class<?> classBeingRedefined,
                             ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
         try {
             // Lambda has no class file, no need to transform, just return.
-            if(classFile == null) {
+            if (classFile == null) {
                 return EMPTY_BYTE_ARRAY;
             }
 
             final String className = toClassName(classFile);
-            if (THREAD_POOL_CLASS_FILE.equals(classFile) || SCHEDULER_CLASS_FILE.equals(classFile)) {
+            if (EXECUTOR_CLASS_NAMES.contains(className)) {
                 logger.info("Transforming class " + className);
                 CtClass clazz = getCtClass(classFileBuffer, loader);
 
                 for (CtMethod method : clazz.getDeclaredMethods()) {
                     updateMethod(clazz, method);
                 }
                 return clazz.toBytecode();
-            } else if (TIMER_TASK_CLASS_FILE.equals(classFile)) {
+            } else if (TIMER_TASK_CLASS_NAME.equals(className)) {
                 CtClass clazz = getCtClass(classFileBuffer, loader);
                 while (true) {
                     String name = clazz.getSuperclass().getName();
                     if (Object.class.getName().equals(name)) {
                         break;
                     }
-                    if (TIMER_TASK_CLASS_FILE.equals(name)) {
+                    if (TIMER_TASK_CLASS_NAME.equals(name)) {
                         logger.info("Transforming class " + className);
                         // FIXME add code here
                         return EMPTY_BYTE_ARRAY;
@@ -88,7 +90,11 @@ public byte[] transform(ClassLoader loader, String classFile, Class<?> classBein
         return EMPTY_BYTE_ARRAY;
     }
 
-    private CtClass getCtClass(byte[] classFileBuffer, ClassLoader classLoader) throws IOException {
+    private static String toClassName(String classFile) {
+        return classFile.replace('/', '.');
+    }
+
+    private static CtClass getCtClass(byte[] classFileBuffer, ClassLoader classLoader) throws IOException {
         ClassPool classPool = new ClassPool(true);
         if (null != classLoader) {
             classPool.appendClassPath(new LoaderClassPath(classLoader));
@@ -99,20 +105,7 @@ private CtClass getCtClass(byte[] classFileBuffer, ClassLoader classLoader) thro
         return clazz;
     }
 
-    static final Set<String> updateMethodNames = new HashSet<String>();
-
-    static {
-        updateMethodNames.add("execute");
-        updateMethodNames.add("submit");
-        updateMethodNames.add("schedule");
-        updateMethodNames.add("scheduleAtFixedRate");
-        updateMethodNames.add("scheduleWithFixedDelay");
-    }
-
-    static void updateMethod(CtClass clazz, CtMethod method) throws NotFoundException, CannotCompileException {
-        if (!updateMethodNames.contains(method.getName())) {
-            return;
-        }
+    private static void updateMethod(CtClass clazz, CtMethod method) throws NotFoundException, CannotCompileException {
         if (method.getDeclaringClass() != clazz) {
             return;
         }