ttl agent support TimerTask #7

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

@@ -2,13 +2,14 @@ package com.alibaba.ttl.threadpool.agent;
 
 
 import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlExecutorTransformlet;
 import com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlForkJoinTransformlet;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlTimerTaskTransformlet;
 
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.*;
 import java.util.logging.Level;
 
 /**
@@ -50,8 +51,12 @@ public final class TtlAgent {
 
     /**
      * Entrance method of TTL Java Agent.
-     * <p>
-     * The log of TTL Java Agent is config by agent argument, using key {@code ttl.agent.logger}.
+     *
+     * <h2>TTL Agent configuration</h2>
+     * Configure TTL agent via agent arguments, format is {@code k1:v1,k2:v2}. Below is available configuration keys.
+     *
+     * <h3>The log configuration</h3>
+     * The log of TTL Java Agent is config by key {@code ttl.agent.logger}.
      *
      * <ul>
      * <li>{@code ttl.agent.logger : STDERR}<br>
@@ -67,6 +72,14 @@ public final class TtlAgent {
      * <li>{@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.logger:STDOUT}</li>
      * </ul>
      *
+     * <h3>Enable TimerTask class decoration</h3>
+     * Enable TimerTask class decoration is config by key {@code ttl.agent.enable.timer.task}. Do not enable when no configuration for this key.
+     * <p>
+     * configuration example: {@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.enable.timer.task:true}
+     *
+     * <h3>Multi key configuration example</h3>
+     * {@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.logger:STDOUT,ttl.agent.enable.timer.task:true}
+     *
      * @see <a href="https://docs.oracle.com/javase/10/docs/api/java/lang/instrument/package-summary.html">The mechanism for instrumentation</a>
      * @see Logger
      * @see Logger#TTL_AGENT_LOGGER_KEY
@@ -74,14 +87,22 @@ public final class TtlAgent {
      * @see Logger#STDOUT
      */
     public static void premain(String agentArgs, Instrumentation inst) {
-        Logger.setLoggerImplType(getLogImplTypeFromAgentArgs(agentArgs));
+        final Map<String, String> kvs = splitCommaColonStringToKV(agentArgs);
+
+        Logger.setLoggerImplType(getLogImplTypeFromAgentArgs(kvs));
         final Logger logger = Logger.getLogger(TtlAgent.class);
 
         try {
             logger.info("[TtlAgent.premain] begin, agentArgs: " + agentArgs + ", Instrumentation: " + inst);
 
-            @SuppressWarnings("unchecked")
-            ClassFileTransformer transformer = new TtlTransformer(TtlExecutorTransformlet.class, TtlForkJoinTransformlet.class);
+            final List<Class<? extends JavassistTransformlet>> transformletList = new ArrayList<Class<? extends JavassistTransformlet>>();
+            transformletList.add(TtlExecutorTransformlet.class);
+            transformletList.add(TtlForkJoinTransformlet.class);
+            if (enableTimerTask(kvs)) {
+                transformletList.add(TtlTimerTaskTransformlet.class);
+            }
+
+            final ClassFileTransformer transformer = new TtlTransformer(transformletList);
             inst.addTransformer(transformer, true);
             logger.info("[TtlAgent.premain] addTransformer " + transformer.getClass() + " success");
 
@@ -93,9 +114,17 @@ public final class TtlAgent {
         }
     }
 
-    private static String getLogImplTypeFromAgentArgs(String agentArgs) {
-        final Map<String, String> kv = splitCommaColonStringToKV(agentArgs);
-        return kv.get(Logger.TTL_AGENT_LOGGER_KEY);
+    private static String getLogImplTypeFromAgentArgs(final Map<String, String> kvs) {
+        return kvs.get(Logger.TTL_AGENT_LOGGER_KEY);
+    }
+
+    private static final String TTL_AGENT_ENABLE_TIMER_TASK_KEY = "ttl.agent.enable.timer.task";
+
+    private static boolean enableTimerTask(final Map<String, String> kvs) {
+        final boolean hasEnableKey = kvs.containsKey(TTL_AGENT_ENABLE_TIMER_TASK_KEY);
+        if (!hasEnableKey) return false;
+
+        return !"false".equalsIgnoreCase(kvs.get(TTL_AGENT_ENABLE_TIMER_TASK_KEY));
     }
 
     /**

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.JavassistTransformlet;
-import javassist.ClassPool;
-import javassist.CtClass;
-import javassist.LoaderClassPath;
 
-import java.io.ByteArrayInputStream;
-import java.io.IOException;
 import java.lang.instrument.ClassFileTransformer;
 import java.security.ProtectionDomain;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.logging.Level;
 
-import com.alibaba.ttl.threadpool.agent.internal.logging.Logger;
-
 /**
  * TTL {@link ClassFileTransformer} of Java Agent
  *
@@ -30,8 +24,7 @@ public class TtlTransformer implements ClassFileTransformer {
 
     private final List<JavassistTransformlet> transformletList = new ArrayList<JavassistTransformlet>();
 
-    @SuppressWarnings("unchecked")
-    TtlTransformer(Class<? extends JavassistTransformlet>... transformletClasses) throws Exception {
+    TtlTransformer(List<Class<? extends JavassistTransformlet>> transformletClasses) throws Exception {
         for (Class<? extends JavassistTransformlet> transformletClass : transformletClasses) {
             final JavassistTransformlet transformlet = transformletClass.getConstructor().newInstance();
             transformletList.add(transformlet);

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

@@ -40,7 +40,7 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
         final String className = clazz.getName();
 
         final String capturedFieldName = "captured$field$add$by$ttl";
-        final CtField capturedField = CtField.make("private final java.lang.Object " + capturedFieldName + ";", clazz);
+        final CtField capturedField = CtField.make("private final Object " + capturedFieldName + ";", clazz);
         clazz.addField(capturedField, "com.alibaba.ttl.TransmittableThreadLocal.Transmitter.capture();");
         logger.info("add new field " + capturedFieldName + " to class " + className);
 
@@ -59,7 +59,7 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
                 "if (this instanceof " + TTL_RECURSIVE_ACTION_CLASS_NAME + " || this instanceof " + TTL_RECURSIVE_TASK_CLASS_NAME + ") {\n" +
                 "    return " + original_doExec_method_rename + "($$);\n" +
                 "}\n" +
-                "java.lang.Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");\n" +
+                "Object backup = com.alibaba.ttl.TransmittableThreadLocal.Transmitter.replay(" + capturedFieldName + ");\n" +
                 "try {\n" +
                 "    return " + original_doExec_method_rename + "($$);\n" +
                 "} finally {\n" +

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.JavassistTransformlet;
+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;
+
+/**
+ * TTL {@link JavassistTransformlet} for {@link java.util.TimerTask}.
+ *
+ * @author Jerry Lee (oldratlee at gmail dot com)
+ * @author wuwen5 (wuwen.55 at aliyun dot com)
+ * @see java.util.TimerTask
+ * @see java.util.Timer
+ * @since 2.7.0
+ */
+public class TtlTimerTaskTransformlet implements JavassistTransformlet {
+    private static final Logger logger = Logger.getLogger(TtlTimerTaskTransformlet.class);
+
+    private static final String TIMER_TASK_CLASS_NAME = "java.util.TimerTask";
+    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
+
+        final CtClass clazz = getCtClass(classFileBuffer, loader);
+
+        if (clazz.isPrimitive() || clazz.isArray() || clazz.isInterface() || clazz.isAnnotation()) {
+            return null;
+        }
+        // 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;
+            }
+        } catch (NotFoundException e) {
+            return null;
+        }
+        if (!clazz.subclassOf(clazz.getClassPool().get(TIMER_TASK_CLASS_NAME))) {
+            return null;
+        }
+
+        logger.info("Transforming class " + className);
+        updateTimerTaskClass(clazz);
+        return clazz.toBytecode();
+    }
+
+    private void updateTimerTaskClass(final CtClass clazz) throws CannotCompileException, NotFoundException {
+        // add new field
+        final String className = clazz.getName();
+
+        final String capturedFieldName = "captured$field$add$by$ttl";
+        final CtField capturedField = CtField.make("private final Object " + capturedFieldName + ";", clazz);
+        clazz.addField(capturedField, "com.alibaba.ttl.TransmittableThreadLocal.Transmitter.capture();");
+        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, RUN_METHOD_NAME, clazz, null);
+
+        // rename original run method, and set to private method(avoid reflect out renamed method unexpectedly)
+        final String original_run_method_rename = "original$run$method$renamed$by$ttl";
+        runMethod.setName(original_run_method_rename);
+        runMethod.setModifiers(runMethod.getModifiers() & ~Modifier.PUBLIC /* remove public */ | Modifier.PRIVATE /* add private */);
+
+        // 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);
+    }
+}

src/test/java/com/alibaba/ttl/threadpool/agent/TtlAgentTest.kt

@@ -12,11 +12,21 @@ class TtlAgentTest {
         assertEquals(emptyMap<String, String>(), splitCommaColonStringToKV(""))
         assertEquals(emptyMap<String, String>(), splitCommaColonStringToKV("   "))
 
+
+        assertEquals(mapOf("k1" to "", "k2" to ""), splitCommaColonStringToKV(
+                "k1,k2"))
+        assertEquals(mapOf("k1" to "", "k2" to ""), splitCommaColonStringToKV(
+                "   k1,   k2 "))
+
         assertEquals(mapOf("ttl.agent.logger" to "STDOUT"), splitCommaColonStringToKV(
                 "ttl.agent.logger:STDOUT"))
         assertEquals(mapOf("k1" to "v1", "ttl.agent.logger" to "STDOUT"), splitCommaColonStringToKV(
                 "k1:v1,ttl.agent.logger:STDOUT"))
+
         assertEquals(mapOf("k1" to "v1", "ttl.agent.logger" to "STDOUT"), splitCommaColonStringToKV(
                 "     k1     :v1  , ttl.agent.logger    :STDOUT   "))
+
+        assertEquals(mapOf("k1" to "v1", "ttl.agent.logger" to "STDOUT", "k3" to ""), splitCommaColonStringToKV(
+                "     k1     :v1  , ttl.agent.logger    :STDOUT   ,k3"))
     }
 }