fix agent order problem, WIP

alibaba · Jan 18, 2021 · 4886b7e · 4886b7e
1 parent d522e7d
commit 4886b7e
Show file tree

Hide file tree

Showing 8 changed files with 66 additions and 2 deletions.
diff --git a/pom.xml b/pom.xml
@@ -267,7 +267,7 @@
 							-->
 							<Premain-Class>com.alibaba.ttl.threadpool.agent.TtlAgent</Premain-Class>
 							<Boot-Class-Path>${project.artifactId}-${project.version}.jar</Boot-Class-Path>
-							<Can-Redefine-Classes>false</Can-Redefine-Classes>
+							<Can-Redefine-Classes>true</Can-Redefine-Classes>
 							<Can-Retransform-Classes>true</Can-Retransform-Classes>
 							<Can-Set-Native-Method-Prefix>false</Can-Set-Native-Method-Prefix>
 						</manifestEntries>

diff --git a/pom4ide.xml b/pom4ide.xml
@@ -267,7 +267,7 @@
 							-->
 							<Premain-Class>com.alibaba.ttl.threadpool.agent.TtlAgent</Premain-Class>
 							<Boot-Class-Path>${project.artifactId}-${project.version}.jar</Boot-Class-Path>
-							<Can-Redefine-Classes>false</Can-Redefine-Classes>
+							<Can-Redefine-Classes>true</Can-Redefine-Classes>
 							<Can-Retransform-Classes>true</Can-Retransform-Classes>
 							<Can-Set-Native-Method-Prefix>false</Can-Set-Native-Method-Prefix>
 						</manifestEntries>

diff --git a/src/main/java/com/alibaba/ttl/threadpool/agent/TtlAgent.java b/src/main/java/com/alibaba/ttl/threadpool/agent/TtlAgent.java
@@ -1,5 +1,6 @@
 package com.alibaba.ttl.threadpool.agent;
 
+import com.alibaba.ttl.spi.TtlEnhanced;
 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;
@@ -143,6 +144,17 @@ public static void premain(final String agentArgs, @NonNull final Instrumentatio
 
             logger.info("[TtlAgent.premain] end");
 
+            if (isRetransformLoadedClasses() && inst.isRetransformClassesSupported()) {
+                for (Class<?> clazz : inst.getAllLoadedClasses()) {
+                    if (!inst.isModifiableClass(clazz)) continue;
+                    if (clazz.isAssignableFrom(TtlEnhanced.class)) continue;
+                    if (clazz.isPrimitive() || clazz.isArray() || clazz.isInterface() || clazz.isAnnotation()) continue;
+
+                    logger.info("retransform class: " + clazz);
+                    inst.retransformClasses(clazz);
+                }
+            }
+
             ttlAgentLoaded = true;
         } catch (Exception e) {
             String msg = "Fail to load TtlAgent , cause: " + e.toString();
@@ -168,10 +180,21 @@ public static boolean isTtlAgentLoaded() {
         return ttlAgentLoaded;
     }
 
+    private static final String TTL_AGENT_RETRANSFORM_LOADED_CLASSES_KEY = "ttl.agent.retransform.loaded.classes";
+
     private static final String TTL_AGENT_ENABLE_TIMER_TASK_KEY = "ttl.agent.enable.timer.task";
 
     private static final String TTL_AGENT_DISABLE_INHERITABLE_FOR_THREAD_POOL = "ttl.agent.disable.inheritable.for.thread.pool";
 
+    /**
+     * Whether timer task is enhanced by ttl agent, check {@link #isTtlAgentLoaded()} first.
+     *
+     * @since 2.13.0
+     */
+    public static boolean isRetransformLoadedClasses() {
+        return isBooleanOptionSet(kvs, TTL_AGENT_RETRANSFORM_LOADED_CLASSES_KEY, true);
+    }
+
     /**
      * Whether disable inheritable for thread pool is enhanced by ttl agent, check {@link #isTtlAgentLoaded()} first.
      *

diff --git a/src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/TtlAgentEnhanced.java b/src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/TtlAgentEnhanced.java
@@ -0,0 +1,15 @@
+package com.alibaba.ttl.threadpool.agent.internal.transformlet;
+
+/**
+ * a Ttl marker/tag interface, for ttl enhanced class by {@link com.alibaba.ttl.threadpool.agent.TtlAgent}
+ *
+ * @author Jerry Lee (oldratlee at gmail dot com)
+ * @see com.alibaba.ttl.threadpool.agent.TtlAgent
+ * @see com.alibaba.ttl.threadpool.agent.internal.transformlet.JavassistTransformlet
+ * @see com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlExecutorTransformlet
+ * @see com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlForkJoinTransformlet
+ * @see com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.TtlTimerTaskTransformlet
+ * @since 2.13.0
+ */
+public interface TtlAgentEnhanced {
+}
diff --git a/.../com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlExecutorTransformlet.java b/.../com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlExecutorTransformlet.java
@@ -16,6 +16,7 @@
 import java.util.Set;
 import java.util.concurrent.Callable;
 
+import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.addTtlAgentEnhancedInterfaceForClass;
 import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.signatureOfMethod;
 
 /**
@@ -59,6 +60,8 @@ public TtlExecutorTransformlet(boolean disableInheritableForThreadPool) {
     public void doTransform(@NonNull final ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException {
         final CtClass clazz = classInfo.getCtClass();
         if (EXECUTOR_CLASS_NAMES.contains(classInfo.getClassName())) {
+            if (!addTtlAgentEnhancedInterfaceForClass(classInfo)) return;
+
             for (CtMethod method : clazz.getDeclaredMethods()) {
                 updateSubmitMethodsOfExecutorClass_decorateToTtlWrapperAndSetAutoWrapperAttachment(method);
             }
@@ -72,6 +75,8 @@ public void doTransform(@NonNull final ClassInfo classInfo) throws IOException,
             }
             if (!clazz.subclassOf(clazz.getClassPool().get(THREAD_POOL_EXECUTOR_CLASS_NAME))) return;
 
+            if (!addTtlAgentEnhancedInterfaceForClass(classInfo)) return;
+
             logger.info("Transforming class " + classInfo.getClassName());
 
             final boolean modified = updateBeforeAndAfterExecuteMethodOfExecutorSubclass(clazz);

diff --git a/.../com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlForkJoinTransformlet.java b/.../com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlForkJoinTransformlet.java
@@ -36,9 +36,13 @@ public TtlForkJoinTransformlet(boolean disableInheritableForThreadPool) {
     @Override
     public void doTransform(@NonNull final ClassInfo classInfo) throws IOException, NotFoundException, CannotCompileException {
         if (FORK_JOIN_TASK_CLASS_NAME.equals(classInfo.getClassName())) {
+            if (!addTtlAgentEnhancedInterfaceForClass(classInfo)) return;
+
             updateForkJoinTaskClass(classInfo.getCtClass());
             classInfo.setModified();
         } else if (disableInheritableForThreadPool && FORK_JOIN_POOL_CLASS_NAME.equals(classInfo.getClassName())) {
+            if (!addTtlAgentEnhancedInterfaceForClass(classInfo)) return;
+
             updateConstructorDisableInheritable(classInfo.getCtClass());
             classInfo.setModified();
         }

diff --git a/...com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlTimerTaskTransformlet.java b/...com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/TtlTimerTaskTransformlet.java
@@ -8,6 +8,7 @@
 
 import java.io.IOException;
 
+import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.addTtlAgentEnhancedInterfaceForClass;
 import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.doTryFinallyForMethod;
 
 /**
@@ -43,6 +44,8 @@ public void doTransform(@NonNull final ClassInfo classInfo) throws IOException,
         }
         if (!clazz.subclassOf(clazz.getClassPool().get(TIMER_TASK_CLASS_NAME))) return;
 
+        if (!addTtlAgentEnhancedInterfaceForClass(classInfo)) return;
+
         logger.info("Transforming class " + classInfo.getClassName());
 
         updateTimerTaskClass(clazz);

diff --git a/src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/Utils.java b/src/main/java/com/alibaba/ttl/threadpool/agent/internal/transformlet/impl/Utils.java
@@ -1,13 +1,16 @@
 package com.alibaba.ttl.threadpool.agent.internal.transformlet.impl;
 
 import com.alibaba.ttl.TtlRunnable;
+import com.alibaba.ttl.threadpool.agent.internal.transformlet.TtlAgentEnhanced;
 import com.alibaba.ttl.spi.TtlAttachments;
 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 edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import javassist.*;
 
+import java.io.IOException;
 import java.lang.reflect.Modifier;
 
 import static com.alibaba.ttl.TransmittableThreadLocal.Transmitter.capture;
@@ -116,4 +119,15 @@ private static boolean isAutoWrapper(@Nullable Runnable ttlAttachments) {
 
         return value;
     }
+
+    public static boolean addTtlAgentEnhancedInterfaceForClass(ClassInfo classInfo) throws NotFoundException, IOException {
+        final CtClass clazz = classInfo.getCtClass();
+        final CtClass ttlAgentEnhancedClass = clazz.getClassPool().get(TtlAgentEnhanced.class.getName());
+
+        if (clazz.subclassOf(ttlAgentEnhancedClass)) return false;
+
+        clazz.addInterface(ttlAgentEnhancedClass);
+        classInfo.setModified();
+        return true;
+    }
 }