agent support for disable ttl inheritable for thread created by thread pool

- expose isDisableInheritableForThreadPool/isEnableTimerTask
- improve test for DisableInheritableForThreadPool

scripts/run-agent-test.sh

@@ -9,6 +9,13 @@ runCmd "${JAVA_CMD[@]}" -cp "$(getClasspathWithoutTtlJar)" \
     -Drun-ttl-test-under-agent=true \
     org.junit.runner.JUnitCore $(junit_test_case | grep -vE '\.TtlAgentTest$|\.JavassistTest$')
 
+# Run unit test under ttl agent, and turn on the disable inheritable for thread pool enhancement
+runCmd "${JAVA_CMD[@]}" -cp "$(getClasspathWithoutTtlJar)" \
+    "-javaagent:$(getTtlJarPath)=ttl.agent.logger:STDOUT,ttl.agent.disable.inheritable.for.thread.pool:true" \
+    -Drun-ttl-test-under-agent=true \
+    -Drun-ttl-test-under-agent-with-disable-inheritable=true \
+    org.junit.runner.JUnitCore $(junit_test_case | grep -vE '\.TtlAgentTest$|\.JavassistTest$')
+
 # Run agent check for Timer/TimerTask
 runCmd "${JAVA_CMD[@]}" -cp "$(getClasspathWithoutTtlJar)" \
     "-javaagent:$(getTtlJarPath)=ttl.agent.logger:STDOUT,ttl.agent.enable.timer.task:true" \

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

@@ -73,15 +73,36 @@ public final class TtlAgent {
      * <li>{@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.logger:STDOUT}</li>
      * </ul>
      *
+     * <h3>Disable inheritable for thread pool</h3>
+     * <p>
+     * Enable "disable inheritable" for thread pool, config by key {@code ttl.agent.disable.inheritable.for.thread.pool}.
+     * When no configuration for this key, default does <b>not</b> enabled.
+     *
+     * <ul>
+     * <li>rewrite the {@link java.util.concurrent.ThreadFactory} constructor parameter
+     * of {@link java.util.concurrent.ThreadPoolExecutor}
+     * to {@link com.alibaba.ttl.threadpool.DisableInheritableThreadFactory}
+     * by util method {@link com.alibaba.ttl.threadpool.TtlExecutors#getDisableInheritableThreadFactory(java.util.concurrent.ThreadFactory)}.
+     * </li>
+     * <li>rewrite the {@link java.util.concurrent.ForkJoinPool} constructor parameter
+     * of {@link java.util.concurrent.ForkJoinPool}
+     * to {@link com.alibaba.ttl.threadpool.DisableInheritableForkJoinWorkerThreadFactory}
+     * by util method {@link com.alibaba.ttl.threadpool.TtlForkJoinPoolHelper#getDisableInheritableForkJoinWorkerThreadFactory(java.util.concurrent.ForkJoinPool.ForkJoinWorkerThreadFactory)}.
+     * </li>
+     * </ul>
+     * <p>
+     * Configuration example:<br>
+     * {@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.disable.inheritable.for.thread.pool:true}
+     *
      * <h3>Enable TimerTask class decoration</h3>
      * Enable TimerTask class decoration is config by key {@code ttl.agent.enable.timer.task}.
-     * When no configuration for this key, default does not enabled.
+     * When no configuration for this key, default does <b>not</b> enabled.
      * <p>
      * Configuration example:<br>
      * {@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}
+     * {@code -javaagent:/path/to/transmittable-thread-local-2.x.x.jar=ttl.agent.logger:STDOUT,ttl.agent.disable.inheritable.for.thread.pool: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
@@ -90,18 +111,19 @@ public final class TtlAgent {
      * @see Logger#STDOUT
      */
     public static void premain(String agentArgs, @Nonnull Instrumentation inst) {
-        final Map<String, String> kvs = splitCommaColonStringToKV(agentArgs);
+        kvs = splitCommaColonStringToKV(agentArgs);
 
         Logger.setLoggerImplType(getLogImplTypeFromAgentArgs(kvs));
         final Logger logger = Logger.getLogger(TtlAgent.class);
 
         try {
             logger.info("[TtlAgent.premain] begin, agentArgs: " + agentArgs + ", Instrumentation: " + inst);
+            final boolean disableInheritable = isDisableInheritableForThreadPool();
 
-            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 List<JavassistTransformlet> transformletList = new ArrayList<JavassistTransformlet>();
+            transformletList.add(new TtlExecutorTransformlet(disableInheritable));
+            transformletList.add(new TtlForkJoinTransformlet(disableInheritable));
+            if (isEnableTimerTask()) transformletList.add(new TtlTimerTaskTransformlet());
 
             final ClassFileTransformer transformer = new TtlTransformer(transformletList);
             inst.addTransformer(transformer, true);
@@ -121,13 +143,54 @@ public final class TtlAgent {
         return kvs.get(Logger.TTL_AGENT_LOGGER_KEY);
     }
 
+    private static volatile Map<String, String> kvs;
+
+    private static volatile boolean ttlAgentLoaded = false;
+
+    /**
+     * Whether TTL agent is loaded.
+     *
+     * @since 2.9.0
+     */
+    public static boolean isTtlAgentLoaded() {
+        return ttlAgentLoaded;
+    }
+
     private static final String TTL_AGENT_ENABLE_TIMER_TASK_KEY = "ttl.agent.enable.timer.task";
 
-    private static boolean enableTimerTask(@Nonnull final Map<String, String> kvs) {
-        final boolean hasEnableKey = kvs.containsKey(TTL_AGENT_ENABLE_TIMER_TASK_KEY);
+    private static final String TTL_AGENT_DISABLE_INHERITABLE_FOR_THREAD_POOL = "ttl.agent.disable.inheritable.for.thread.pool";
+
+    /**
+     * Whether disable inheritable for thread pool is enhanced by ttl agent, check {@link #isTtlAgentLoaded()} first.
+     *
+     * @see com.alibaba.ttl.threadpool.DisableInheritableThreadFactory
+     * @see com.alibaba.ttl.threadpool.TtlExecutors#getDisableInheritableThreadFactory(java.util.concurrent.ThreadFactory)
+     * @see com.alibaba.ttl.threadpool.TtlExecutors#getDefaultDisableInheritableThreadFactory()
+     * @see com.alibaba.ttl.threadpool.DisableInheritableForkJoinWorkerThreadFactory
+     * @see com.alibaba.ttl.threadpool.TtlForkJoinPoolHelper#getDisableInheritableForkJoinWorkerThreadFactory(java.util.concurrent.ForkJoinPool.ForkJoinWorkerThreadFactory)
+     * @see com.alibaba.ttl.threadpool.TtlForkJoinPoolHelper#getDefaultDisableInheritableForkJoinWorkerThreadFactory()
+     * @since 2.10.1
+     */
+    public static boolean isDisableInheritableForThreadPool() {
+        return isOptionSet(kvs, TTL_AGENT_DISABLE_INHERITABLE_FOR_THREAD_POOL);
+    }
+
+    /**
+     * Whether timer task is enhanced by ttl agent, check {@link #isTtlAgentLoaded()} first.
+     *
+     * @since 2.10.1
+     */
+    public static boolean isEnableTimerTask() {
+        return isOptionSet(kvs, TTL_AGENT_ENABLE_TIMER_TASK_KEY);
+    }
+
+    private static boolean isOptionSet(@Nullable final Map<String, String> kvs, @Nonnull String key) {
+        if (null == kvs) return false;
+
+        final boolean hasEnableKey = kvs.containsKey(key);
         if (!hasEnableKey) return false;
 
-        return !"false".equalsIgnoreCase(kvs.get(TTL_AGENT_ENABLE_TIMER_TASK_KEY));
+        return !"false".equalsIgnoreCase(kvs.get(key));
     }
 
     /**
@@ -150,15 +213,6 @@ public final class TtlAgent {
         return ret;
     }
 
-    private static boolean ttlAgentLoaded = false;
-
-    /**
-     * whether TTL agent is loaded.
-     */
-    public static boolean isTtlAgentLoaded() {
-        return ttlAgentLoaded;
-    }
-
     private TtlAgent() {
         throw new InstantiationError("Must not instantiate this class");
     }

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

@@ -26,12 +26,10 @@ public class TtlTransformer implements ClassFileTransformer {
 
     private final List<JavassistTransformlet> transformletList = new ArrayList<JavassistTransformlet>();
 
-    TtlTransformer(List<Class<? extends JavassistTransformlet>> transformletClasses) throws Exception {
-        for (Class<? extends JavassistTransformlet> transformletClass : transformletClasses) {
-            final JavassistTransformlet transformlet = transformletClass.getConstructor().newInstance();
-            transformletList.add(transformlet);
-
-            logger.info("[TtlTransformer] add Transformlet " + transformletClass + " success");
+    TtlTransformer(List<? extends JavassistTransformlet> transformletList) {
+        for (JavassistTransformlet transformlet : transformletList) {
+            this.transformletList.add(transformlet);
+            logger.info("[TtlTransformer] add Transformlet " + transformlet.getClass() + " success");
         }
     }
 

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

@@ -2,10 +2,7 @@ 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.CannotCompileException;
-import javassist.CtClass;
-import javassist.CtMethod;
-import javassist.NotFoundException;
+import javassist.*;
 
 import java.io.IOException;
 import java.lang.reflect.Modifier;
@@ -14,7 +11,8 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
-import static com.alibaba.ttl.threadpool.agent.internal.transformlet.impl.Utils.*;
+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.concurrent.Executor}.
@@ -42,6 +40,14 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
         PARAM_TYPE_NAME_TO_DECORATE_METHOD_CLASS.put("java.util.concurrent.Callable", "com.alibaba.ttl.TtlCallable");
     }
 
+    private static final String THREAD_FACTORY_CLASS_NAME = "java.util.concurrent.ThreadFactory";
+
+    private final boolean disableInheritable;
+
+    public TtlExecutorTransformlet(boolean disableInheritable) {
+        this.disableInheritable = disableInheritable;
+    }
+
     @Override
     public byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException {
         if (EXECUTOR_CLASS_NAMES.contains(className)) {
@@ -50,6 +56,9 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
             for (CtMethod method : clazz.getDeclaredMethods()) {
                 updateMethodOfExecutorClass(clazz, method);
             }
+
+            if (disableInheritable) updateConstructorDisableInheritable(clazz);
+
             return clazz.toBytecode();
         }
         return null;
@@ -73,4 +82,20 @@ public class TtlExecutorTransformlet implements JavassistTransformlet {
         }
         if (insertCode.length() > 0) method.insertBefore(insertCode.toString());
     }
+
+    private void updateConstructorDisableInheritable(final CtClass clazz) throws NotFoundException, CannotCompileException {
+        for (CtConstructor constructor : clazz.getDeclaredConstructors()) {
+            final CtClass[] parameterTypes = constructor.getParameterTypes();
+            final StringBuilder insertCode = new StringBuilder();
+            for (int i = 0; i < parameterTypes.length; i++) {
+                final String paramTypeName = parameterTypes[i].getName();
+                if (THREAD_FACTORY_CLASS_NAME.equals(paramTypeName)) {
+                    String code = String.format("$%d = com.alibaba.ttl.threadpool.TtlExecutors.getDisableInheritableThreadFactory($%<d);", i + 1);
+                    logger.info("insert code before method " + signatureOfMethod(constructor) + " of class " + constructor.getDeclaringClass().getName() + ": " + code);
+                    insertCode.append(code);
+                }
+            }
+            if (insertCode.length() > 0) constructor.insertBefore(insertCode.toString());
+        }
+    }
 }

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

@@ -22,6 +22,14 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
     private static final Logger logger = Logger.getLogger(TtlForkJoinTransformlet.class);
 
     private static final String FORK_JOIN_TASK_CLASS_NAME = "java.util.concurrent.ForkJoinTask";
+    private static final String FORK_JOIN_POOL_CLASS_NAME = "java.util.concurrent.ForkJoinPool";
+    private static final String FORK_JOIN_WORKER_THREAD_FACTORY_CLASS_NAME = "java.util.concurrent.ForkJoinPool$ForkJoinWorkerThreadFactory";
+
+    private final boolean disableInheritable;
+
+    public TtlForkJoinTransformlet(boolean disableInheritable) {
+        this.disableInheritable = disableInheritable;
+    }
 
     @Override
     public byte[] doTransform(String className, byte[] classFileBuffer, ClassLoader loader) throws IOException, NotFoundException, CannotCompileException {
@@ -30,7 +38,12 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
 
             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();
         }
+
         return null;
     }
 
@@ -55,4 +68,20 @@ public class TtlForkJoinTransformlet implements JavassistTransformlet {
 
         doTryFinallyForMethod(doExecMethod, doExec_renamed_method_rename, beforeCode, finallyCode);
     }
+
+    private void updateConstructorDisableInheritable(final CtClass clazz) throws NotFoundException, CannotCompileException {
+        for (CtConstructor constructor : clazz.getDeclaredConstructors()) {
+            final CtClass[] parameterTypes = constructor.getParameterTypes();
+            final StringBuilder insertCode = new StringBuilder();
+            for (int i = 0; i < parameterTypes.length; i++) {
+                final String paramTypeName = parameterTypes[i].getName();
+                if (FORK_JOIN_WORKER_THREAD_FACTORY_CLASS_NAME.equals(paramTypeName)) {
+                    String code = String.format("$%d = com.alibaba.ttl.threadpool.TtlForkJoinPoolHelper.getDisableInheritableForkJoinWorkerThreadFactory($%<d);", i + 1);
+                    logger.info("insert code before method " + signatureOfMethod(constructor) + " of class " + constructor.getDeclaringClass().getName() + ": " + code);
+                    insertCode.append(code);
+                }
+            }
+            if (insertCode.length() > 0) constructor.insertBefore(insertCode.toString());
+        }
+    }
 }

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

@@ -23,14 +23,15 @@ public class Utils {
      * @param method method object
      * @return method signature string
      */
-    static String signatureOfMethod(final CtMethod method) throws NotFoundException {
+    static String signatureOfMethod(final CtBehavior method) throws NotFoundException {
         final StringBuilder stringBuilder = new StringBuilder();
 
-        final String returnType = method.getReturnType().getSimpleName();
-        final String modifier = Modifier.toString(method.getModifiers());
-        final String methodName = method.getName();
-        stringBuilder.append(modifier).append(" ").append(returnType).append(" ")
-                .append(methodName).append("(");
+        stringBuilder.append(Modifier.toString(method.getModifiers()));
+        if (method instanceof CtMethod) {
+            final String returnType = ((CtMethod) method).getReturnType().getSimpleName();
+            stringBuilder.append(" ").append(returnType);
+        }
+        stringBuilder.append(" ").append(method.getName()).append("(");
 
         final CtClass[] parameterTypes = method.getParameterTypes();
         for (int i = 0; i < parameterTypes.length; i++) {

src/test/java/com/alibaba/Utils.kt

@@ -4,8 +4,7 @@ package com.alibaba
 
 import com.alibaba.ttl.TransmittableThreadLocal
 import com.alibaba.ttl.threadpool.agent.TtlAgent
-import org.junit.Assert.assertEquals
-import org.junit.Assert.assertNull
+import org.junit.Assert.*
 import java.lang.Thread.sleep
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentMap
@@ -144,8 +143,22 @@ fun <T> assertParentTtlValues(values: TtlValues<T>) {
     )
 }
 
+
 fun noTtlAgentRun(): Boolean = TtlAgent.isTtlAgentLoaded().let {
     val isUnderAgentSet = System.getProperties().containsKey("run-ttl-test-under-agent")
     assertEquals(isUnderAgentSet, it)
     !it
 }
+
+fun hasTtlAgentRun(): Boolean = !noTtlAgentRun()
+
+
+fun noTtlAgentDisableInheritableForThreadPool(): Boolean = TtlAgent.isDisableInheritableForThreadPool().let {
+    val isUnderAgentSet = System.getProperties().containsKey("run-ttl-test-under-agent-with-disable-inheritable")
+    assertEquals(isUnderAgentSet, it)
+    if (it) assertTrue(hasTtlAgentRun())
+
+    !it
+}
+
+fun hasTtlTtlAgentRunWithDisableInheritableForThreadPool(): Boolean = !noTtlAgentDisableInheritableForThreadPool()

src/test/java/com/alibaba/support/junit/conditional/NoAgentRun.kt

+package com.alibaba.support.junit.conditional
+
+import com.alibaba.noTtlAgentRun
+import com.alibaba.support.junit.conditional.ConditionalIgnoreRule.IgnoreCondition
+
+/**
+ * @see [Getting Java version at runtime](https://stackoverflow.com/a/23706899/922688)
+ */
+class NoAgentRun : IgnoreCondition {
+    override fun isSatisfied(): Boolean = noTtlAgentRun()
+}

src/test/java/com/alibaba/ttl/InheritableTest.kt

 package com.alibaba.ttl
 
+import com.alibaba.hasTtlTtlAgentRunWithDisableInheritableForThreadPool
 import com.alibaba.support.junit.conditional.BelowJava7
 import com.alibaba.support.junit.conditional.ConditionalIgnoreRule
 import com.alibaba.support.junit.conditional.ConditionalIgnoreRule.ConditionalIgnore
@@ -15,6 +16,7 @@ import java.util.*
 import java.util.concurrent.Callable
 import java.util.concurrent.Executors
 import java.util.concurrent.ForkJoinPool
+import java.util.concurrent.ThreadPoolExecutor
 
 private const val hello = "hello"
 private val defaultValue = "${Date()} ${Math.random()}"
@@ -144,6 +146,17 @@ class InheritableTest {
         }
     }
 
+    @Test
+    fun disableInheritable_Executors_ByAgent() {
+        val threadPool = Executors.newCachedThreadPool() as ThreadPoolExecutor
+        try {
+            assertEquals(hasTtlTtlAgentRunWithDisableInheritableForThreadPool(),
+                    TtlExecutors.isDisableInheritableThreadFactory(threadPool.threadFactory))
+        } finally {
+            threadPool.shutdown()
+        }
+    }
+
     // ===================================================
     // ForkJoinPool
     // ===================================================
@@ -265,4 +278,16 @@ class InheritableTest {
         }
     }
 
+    @Test
+    @ConditionalIgnore(condition = BelowJava7::class)
+    fun disableInheritable_ForkJoinPool_ByAgent() {
+        val threadPool = ForkJoinPool(4)
+        try {
+            assertEquals(hasTtlTtlAgentRunWithDisableInheritableForThreadPool(),
+                    TtlForkJoinPoolHelper.isDisableInheritableForkJoinWorkerThreadFactory(threadPool.factory))
+        } finally {
+            threadPool.shutdown()
+        }
+    }
+
 }