refactor doTransform implementation

src/main/java/com/alibaba/ttl/threadpool/agent/TtlTransformer.java

 package com.alibaba.ttl.threadpool.agent;
 
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.ClassInfo;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 
 import javax.annotation.Nonnull;
@@ -41,9 +42,12 @@ public class TtlTransformer implements ClassFileTransformer {
             if (classFile == null) return EMPTY_BYTE_ARRAY;
 
             final String className = toClassName(classFile);
+
+            ClassInfo classInfo = new ClassInfo(className, classFileBuffer, loader);
+
             for (JavassistTransformlet transformlet : transformletList) {
-                final byte[] bytes = transformlet.doTransform(className, classFileBuffer, loader);
-                if (bytes != null) return bytes;
+                transformlet.doTransform(classInfo);
+                if (classInfo.isModified()) return classInfo.getCtClass().toBytecode();
             }
         } catch (Throwable t) {
             String msg = "Fail to transform class " + classFile + ", cause: " + t.toString();

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

+package com.alibaba.ttl.threadpool.agent.internal.transformlet;
+
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.LoaderClassPath;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+
+public class ClassInfo {
+    private final String className;
+    private final byte[] classFileBuffer;
+    private final ClassLoader loader;
+
+    public ClassInfo(String className, byte[] classFileBuffer, ClassLoader loader) {
+        this.className = className;
+        this.classFileBuffer = classFileBuffer;
+        this.loader = loader;
+    }
+
+    public String getClassName() {
+        return className;
+    }
+
+    private CtClass ctClass;
+
+    public CtClass getCtClass() throws IOException {
+        if (ctClass != null) return ctClass;
+
+        final ClassPool classPool = new ClassPool(true);
+        if (loader == null) {
+            classPool.appendClassPath(new LoaderClassPath(ClassLoader.getSystemClassLoader()));
+        } else {
+            classPool.appendClassPath(new LoaderClassPath(loader));
+        }
+
+        final CtClass clazz = classPool.makeClass(new ByteArrayInputStream(classFileBuffer), false);
+        clazz.defrost();
+
+        this.ctClass = clazz;
+        return clazz;
+    }
+
+    private boolean modified = false;
+
+    public boolean isModified() {
+        return modified;
+    }
+
+    public void setModified() {
+        this.modified = true;
+    }
+}

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

@@ -12,5 +12,5 @@ import java.io.IOException;
  * @since 2.5.1
  */
 public interface JavassistTransformlet {
-    byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException;
+    void doTransform(ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException;
 }

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

@@ -2,6 +2,7 @@ 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.ClassInfo;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 import javassist.*;
 
@@ -13,7 +14,6 @@ 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;
 
 /**
@@ -54,9 +54,9 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
     }
 
     @Override
-    public byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException {
-        if (EXECUTOR_CLASS_NAMES.contains(className)) {
-            final CtClass clazz = getCtClass(classFileBuffer, loader);
+    public void doTransform(final ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException {
+        if (EXECUTOR_CLASS_NAMES.contains(classInfo.getClassName())) {
+            final CtClass clazz = classInfo.getCtClass();
 
             for (CtMethod method : clazz.getDeclaredMethods()) {
                 updateSubmitMethodsOfExecutorClass_decorateToTtlWrapperAndSetAutoWrapper(method);
@@ -64,20 +64,19 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
 
             if (disableInheritable) updateConstructorDisableInheritable(clazz);
 
-            return clazz.toBytecode();
+            classInfo.setModified();
         } else {
-            final CtClass clazz = getCtClass(classFileBuffer, loader);
+            final CtClass clazz = classInfo.getCtClass();
 
             if (clazz.isPrimitive() || clazz.isArray() || clazz.isInterface() || clazz.isAnnotation()) {
-                return null;
+                return;
             }
-            if (!clazz.subclassOf(clazz.getClassPool().get(THREAD_POOL_EXECUTOR_CLASS_NAME))) return null;
+            if (!clazz.subclassOf(clazz.getClassPool().get(THREAD_POOL_EXECUTOR_CLASS_NAME))) return;
 
-            logger.info("Transforming class " + className);
+            logger.info("Transforming class " + classInfo.getClassName());
 
-            final boolean updated = updateBeforeAndAfterExecuteMethodOfExecutorSubclass(clazz);
-            if (updated) return clazz.toBytecode();
-            else return null;
+            final boolean modified = updateBeforeAndAfterExecuteMethodOfExecutorSubclass(clazz);
+            if (modified) classInfo.setModified();
         }
     }
 
@@ -134,7 +133,7 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
         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;
+        boolean modified = false;
 
         try {
             final CtMethod beforeExecute = clazz.getDeclaredMethod("beforeExecute", new CtClass[]{threadClass, runnableClass});
@@ -142,7 +141,7 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
             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;
+            modified = true;
         } catch (NotFoundException e) {
             // clazz does not override beforeExecute method, do nothing.
         }
@@ -153,11 +152,11 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
             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;
+            modified = true;
         } catch (NotFoundException e) {
             // clazz does not override afterExecute method, do nothing.
         }
 
-        return updated;
+        return modified;
     }
 }

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

@@ -2,6 +2,7 @@ package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;
 
 import com.alibaba.ttl.spi.TtlEnhanced;
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.ClassInfo;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 import javassist.*;
 
@@ -32,19 +33,14 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
     }
 
     @Override
-    public byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException {
-        if (FORK_JOIN_TASK_CLASS_NAME.equals(className)) {
-            final CtClass clazz = getCtClass(classFileBuffer, loader);
-
-            updateForkJoinTaskClass(clazz);
-            return clazz.toBytecode();
-        } else if (disableInheritable && FORK_JOIN_POOL_CLASS_NAME.equals(className)) {
-            final CtClass clazz = getCtClass(classFileBuffer, loader);
-            updateConstructorDisableInheritable(clazz);
-            return clazz.toBytecode();
+    public void doTransform(final ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException {
+        if (FORK_JOIN_TASK_CLASS_NAME.equals(classInfo.getClassName())) {
+            updateForkJoinTaskClass(classInfo.getCtClass());
+            classInfo.setModified();
+        } else if (disableInheritable && FORK_JOIN_POOL_CLASS_NAME.equals(classInfo.getClassName())) {
+            updateConstructorDisableInheritable(classInfo.getCtClass());
+            classInfo.setModified();
         }
-
-        return null;
     }
 
     /**

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

 package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;
 
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.ClassInfo;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 import javassist.*;
 
 import java.io.IOException;
 
-import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.*;
+import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.doTryFinallyForMethod;
 
 /**
  * TTL {@link JavassistTransformlet} for {@link java.util.TimerTask}.
@@ -24,27 +25,27 @@ public class TtlTimerTaskTransformlet implements JavassistTransformlet {
     private static final String RUN_METHOD_NAME = "run";
 
     @Override
-    public byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException {
-        if (TIMER_TASK_CLASS_NAME.equals(className)) return null; // No need transform TimerTask class
+    public void doTransform(final ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException {
+        if (TIMER_TASK_CLASS_NAME.equals(classInfo.getClassName())) return; // No need transform TimerTask class
 
-        final CtClass clazz = getCtClass(classFileBuffer, loader);
+        final CtClass clazz = classInfo.getCtClass();
 
         if (clazz.isPrimitive() || clazz.isArray() || clazz.isInterface() || clazz.isAnnotation()) {
-            return null;
+            return;
         }
         // class contains method `void run()` ?
         try {
             final CtMethod runMethod = clazz.getDeclaredMethod(RUN_METHOD_NAME, new CtClass[0]);
-            if (!CtClass.voidType.equals(runMethod.getReturnType())) return null;
+            if (!CtClass.voidType.equals(runMethod.getReturnType())) return;
         } catch (NotFoundException e) {
-            return null;
+            return;
         }
-        if (!clazz.subclassOf(clazz.getClassPool().get(TIMER_TASK_CLASS_NAME))) return null;
+        if (!clazz.subclassOf(clazz.getClassPool().get(TIMER_TASK_CLASS_NAME))) return;
 
-        logger.info("Transforming class " + className);
+        logger.info("Transforming class " + classInfo.getClassName());
 
         updateTimerTaskClass(clazz);
-        return clazz.toBytecode();
+        classInfo.setModified();
     }
 
     /**

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

@@ -46,20 +46,6 @@ public class Utils {
         return stringBuilder.toString();
     }
 
-    static CtClass getCtClass(final byte[] classFileBuffer, final ClassLoader classLoader) throws IOException {
-        final ClassPool classPool = new ClassPool(true);
-        if (classLoader == null) {
-            classPool.appendClassPath(new LoaderClassPath(ClassLoader.getSystemClassLoader()));
-        } else {
-            classPool.appendClassPath(new LoaderClassPath(classLoader));
-        }
-
-        final CtClass clazz = classPool.makeClass(new ByteArrayInputStream(classFileBuffer), false);
-        clazz.defrost();
-        return clazz;
-    }
-
-
     static String renamedMethodNameByTtl(CtMethod method) {
         return "original$" + method.getName() + "$method$renamed$by$ttl";
     }