extract javassist util method doTryFinallyForMethod

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

@@ -5,10 +5,8 @@ import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransform
 import javassist.*;
 
 import java.io.IOException;
-import java.lang.reflect.Modifier;
 
-import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.getCtClass;
-import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.signatureOfMethod;
+import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.*;
 
 /**
  * TTL {@link JavassistTransformlet} for {@link java.util.concurrent.ForkJoinTask}.
@@ -45,27 +43,15 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
         logger.info("add new field " + capturedFieldName + " to class " + className);
 
         final CtMethod doExecMethod = clazz.getDeclaredMethod("doExec", new CtClass[0]);
-        final CtMethod new_doExecMethod = CtNewMethod.copy(doExecMethod, clazz, null);
+        final String original_doExec_method_rename = renamedMethodNameByTtl(doExecMethod);
 
-        // rename original doExec method, and set to private method(avoid reflect out renamed method unexpectedly)
-        final String original_doExec_method_rename = "original$" + doExecMethod.getName()+ "$method$renamed$by$ttl";
-        doExecMethod.setName(original_doExec_method_rename);
-        doExecMethod.setModifiers(doExecMethod.getModifiers() & ~Modifier.PUBLIC /* remove public */ | Modifier.PRIVATE /* add private */);
-
-        // set new doExec method implementation
-        final String code = "{\n" +
-                // do nothing/directly return, if is TTL ForkJoinTask instance
-                "if (this instanceof " + TTL_RECURSIVE_ACTION_CLASS_NAME + " || this instanceof " + TTL_RECURSIVE_TASK_CLASS_NAME + ") {\n" +
-                "    return " + original_doExec_method_rename + "($$);\n" +
+        final String beforeCode = "if (this instanceof " + TTL_RECURSIVE_ACTION_CLASS_NAME + " || this instanceof " + TTL_RECURSIVE_TASK_CLASS_NAME + ") {\n" +
+                "    return " + original_doExec_method_rename + "($$);\n" + // do nothing/directly return, if is TTL ForkJoinTask instance
                 "}\n" +
-                "Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");\n" +
-                "try {\n" +
-                "    return " + original_doExec_method_rename + "($$);\n" +
-                "} finally {\n" +
-                "    com.alibaba.ttl.TransmittableThreadLocal.Transmitter.restore(backup);\n" +
-                "}\n" + "}";
-        new_doExecMethod.setBody(code);
-        clazz.addMethod(new_doExecMethod);
-        logger.info("insert code around method " + signatureOfMethod(doExecMethod) + " of class " + className + ": " + code);
+                "Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");";
+
+        final String finallyCode = "com.alibaba.ttl.TransmittableThreadLocal.Transmitter.restore(backup);";
+
+        doTryFinallyForMethod(doExecMethod, original_doExec_method_rename, beforeCode, finallyCode);
     }
 }

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

@@ -5,10 +5,9 @@ import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransform
 import javassist.*;
 
 import java.io.IOException;
-import java.lang.reflect.Modifier;
 
+import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.doTryFinallyForMethod;
 import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.getCtClass;
-import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.signatureOfMethod;
 
 /**
  * TTL {@link JavassistTransformlet} for {@link java.util.TimerTask}.
@@ -58,23 +57,10 @@ public class TtlTimerTaskTransformlet implements JavassistTransformlet {
         logger.info("add new field " + capturedFieldName + " to class " + className);
 
         final CtMethod runMethod = clazz.getDeclaredMethod(RUN_METHOD_NAME, new CtClass[0]);
-        final CtMethod new_runMethod = CtNewMethod.copy(runMethod, clazz, null);
 
-        // rename original run method, and set to private method(avoid reflect out renamed method unexpectedly)
-        final String original_run_method_rename = "original$" + runMethod.getName() + "$method$renamed$by$ttl";
-        runMethod.setName(original_run_method_rename);
-        runMethod.setModifiers(runMethod.getModifiers() & ~Modifier.PUBLIC /* remove public */ | Modifier.PRIVATE /* add private */);
+        final String beforeCode = "Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");";
+        final String finallyCode = "com.alibaba.ttl.TransmittableThreadLocal.Transmitter.restore(backup);";
 
-        // set new run method implementation
-        final String code = "{\n" +
-                "Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");\n" +
-                "try {\n" +
-                "    return " + original_run_method_rename + "($$);\n" +
-                "} finally {\n" +
-                "    com.alibaba.ttl.TransmittableThreadLocal.Transmitter.restore(backup);\n" +
-                "}\n" + "}";
-        new_runMethod.setBody(code);
-        clazz.addMethod(new_runMethod);
-        logger.info("insert code around method " + signatureOfMethod(runMethod) + " of class " + className + ": " + code);
+        doTryFinallyForMethod(runMethod, beforeCode, finallyCode);
     }
 }

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.threadpool.agent.internal.logging.Logger;
 import javassist.*;
 
 import java.io.ByteArrayInputStream;
@@ -11,6 +12,8 @@ import java.lang.reflect.Modifier;
  * @since 2.6.0
  */
 class Utils {
+    private static final Logger logger = Logger.getLogger(Utils.class);
+
     /**
      * String like {@code ScheduledFuture scheduleAtFixedRate(Runnable, long, long, TimeUnit)}
      * for {@link  java.util.concurrent.ScheduledThreadPoolExecutor#scheduleAtFixedRate}.
@@ -50,4 +53,37 @@ class Utils {
         clazz.defrost();
         return clazz;
     }
+
+
+    static String renamedMethodNameByTtl(CtMethod method) {
+        return "original$" + method.getName() + "$method$renamed$by$ttl";
+    }
+
+    static void doTryFinallyForMethod(CtMethod method, String beforeCode, String finallyCode) throws CannotCompileException, NotFoundException {
+        doTryFinallyForMethod(method, renamedMethodNameByTtl(method), beforeCode, finallyCode);
+    }
+
+    static void doTryFinallyForMethod(CtMethod method, String originalMethodRename, String beforeCode, String finallyCode) throws CannotCompileException, NotFoundException {
+        final CtClass clazz = method.getDeclaringClass();
+        final CtMethod new_method = CtNewMethod.copy(method, clazz, null);
+
+        // rename original method, and set to private method(avoid reflect out renamed method unexpectedly)
+        method.setName(originalMethodRename);
+        method.setModifiers(method.getModifiers()
+                & ~Modifier.PUBLIC /* remove public */
+                & ~Modifier.PROTECTED /* remove protected */
+                | Modifier.PRIVATE /* add private */);
+
+        // set new method implementation
+        final String code = "{\n" +
+                beforeCode + "\n" +
+                "try {\n" +
+                "    return " + originalMethodRename + "($$);\n" +
+                "} finally {\n" +
+                "    " + finallyCode + "\n" +
+                "} }";
+        new_method.setBody(code);
+        clazz.addMethod(new_method);
+        logger.info("insert code around method " + signatureOfMethod(method) + " of class " + clazz.getName() + ": " + code);
+    }
 }