ClassLoader Code Update To New Branch

src/main/java/com/alibaba/ttl/classloader/TtlAgentJarUtil.java

+package com.alibaba.ttl.classloader;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.net.URLDecoder;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import java.util.regex.Pattern;
+
+import com.alibaba.ttl.threadpool.agent.TtlAgent;
+
+public class TtlAgentJarUtil {
+
+    private static final Logger logger = Logger.getLogger(TtlAgentJarUtil.class.getName());
+
+    private static final Pattern CLASS_NAME = Pattern.compile(TtlAgent.class.getName().replace('.', '/') + ".class");
+
+    private static URL cachedURL = null;
+
+    @SuppressWarnings({ "rawtypes" })
+    public static URL getAgentJarFilePath() {
+
+        if (cachedURL != null) {
+            return cachedURL;
+        }
+        ClassLoader sysClassLoader = ClassLoader.getSystemClassLoader();
+        if ((sysClassLoader instanceof URLClassLoader)) {
+            URL[] jarUrls = ((URLClassLoader) sysClassLoader).getURLs();
+            for (URL jarFullPath : jarUrls) {
+                if ((jarFullPath.getFile().endsWith(".jar"))
+                        && (jarFullPath.getFile().indexOf("transmittable-thread-local") != -1)) {
+                    Collection<String> jarClassCollection = getClassCollectionsByJarUrl(jarFullPath, CLASS_NAME);
+                    if (!jarClassCollection.isEmpty()) {
+                        cachedURL = jarFullPath;
+                        return cachedURL;
+                    }
+                }
+            }
+        }
+        // JDK 10 and later needs this code
+        Class loadClasssysClassLoader = sysClassLoader.getClass();
+        if (loadClasssysClassLoader.getName().equals("jdk.internal.loader.BuiltinClassLoader")
+                || loadClasssysClassLoader.getSuperclass().getName().equals("jdk.internal.loader.BuiltinClassLoader")) {
+            try {
+
+                //ucp field has class jar file system location information that we need
+                Field classPathField = loadClasssysClassLoader.getDeclaredField("ucp");
+                Object classLoaderModule = java.lang.Class.class.getDeclaredMethod("getModule")
+                        .invoke(loadClasssysClassLoader);
+                //get system classloader's module
+                Field moduleField = java.lang.Class.class.getDeclaredField("module");
+                moduleField.setAccessible(true);
+
+                // set agent class module the same as classloader's module
+                moduleField.set(TtlAgentJarUtil.class, classLoaderModule);
+
+                // in jdk10 after caller( in this case TtlAgentJarUtil.class) has the same module with target class (systemClassLoad of jdk 10)
+                // then we can call setAccessible(true) to access private field
+                classPathField.setAccessible(true);
+                Object urlclasspath = classPathField.get(sysClassLoader);
+
+                // jdk.internal.loader.URLClassPath is JDK 10 special class, we can only reflect invoke
+                Method getUrlMethod = Class.forName("jdk.internal.loader.URLClassPath").getDeclaredMethod("getURLs");
+                URL[] jarUrls = (URL[]) getUrlMethod.invoke(urlclasspath);
+                for (URL jarFullPath : jarUrls) {
+                    if ((jarFullPath.getFile().endsWith(".jar"))
+                            && (jarFullPath.getFile().indexOf("transmittable-thread-local") != -1)) {
+                        Collection<String> jarClassCollection = getClassCollectionsByJarUrl(jarFullPath, CLASS_NAME);
+                        if (!jarClassCollection.isEmpty()) {
+                            cachedURL = jarFullPath;
+                            return cachedURL;
+                        }
+                    }
+                }
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }
+        cachedURL = TtlAgent.class.getProtectionDomain().getCodeSource().getLocation();
+        return cachedURL;
+    }
+
+    private static Collection<String> getClassCollectionsByJarUrl(URL paramURL, Pattern classNamePattern) {
+        JarFile localJarFile = null;
+        try {
+            localJarFile = getJarFileByPath(paramURL);
+            List<String> destinationClassList = new ArrayList<String>();
+            Enumeration<JarEntry> jarEntry = localJarFile.entries();
+            while (jarEntry.hasMoreElements()) {
+                JarEntry jarEntryObject = jarEntry.nextElement();
+                if (classNamePattern.matcher(jarEntryObject.getName()).matches()) {
+                    destinationClassList.add(jarEntryObject.getName());
+                }
+            }
+            return destinationClassList;
+        } catch (Exception t) {
+            if (logger.isLoggable(Level.WARNING)) {
+                logger.log(Level.WARNING, "Unable to search the agent jar for " + classNamePattern.pattern(), t);
+            }
+            return Collections.emptyList();
+        } finally {
+            if (localJarFile != null) {
+                try {
+                    localJarFile.close();
+                } catch (IOException localIOException3) {
+                }
+            }
+        }
+    }
+
+    public static Collection<String> getClassCollectionsByClassPattern(Pattern classPattern) {
+        URL jarUrl = getAgentJarFilePath();
+        return getClassCollectionsByJarUrl(jarUrl, classPattern);
+    }
+
+    public static File getJarFile() {
+        URL argentUrl = getAgentJarFilePath();
+        if (argentUrl != null) {
+            File jarFile = new File(replaceJarFullPath(argentUrl));
+            if (jarFile.exists()) {
+                return jarFile.getParentFile();
+            }
+        }
+        return null;
+    }
+
+    public static JarFile getJarFileByPath(URL jarFullPath) {
+        if (jarFullPath != null) {
+            try {
+                return new JarFile(replaceJarFullPath(jarFullPath));
+            } catch (IOException localIOException) {
+            }
+        }
+        return null;
+    }
+
+    private static String replaceJarFullPath(URL jarFullPath) {
+        if (jarFullPath == null) {
+            return null;
+        }
+        try {
+            return URLDecoder.decode(jarFullPath.getFile().replace("+", "%2B"), "UTF-8");
+        } catch (IOException localIOException) {
+        }
+        return null;
+    }
+
+}

src/main/java/com/alibaba/ttl/classloader/TtlClassCache.java

+package com.alibaba.ttl.classloader;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+//c
+public class TtlClassCache<T> {
+
+    private Map<String, Map<ClassLoader, T>> classCacheMap = new ConcurrentHashMap<String, Map<ClassLoader, T>>();
+
+    public T loadClassByName(ClassLoader paramClassLoader, String className) {
+        Map<ClassLoader, T> classLoaderKeyAndClassValueMap = (Map<ClassLoader, T>) this.classCacheMap.get(className);
+        if (classLoaderKeyAndClassValueMap != null) {
+            if (classLoaderKeyAndClassValueMap.containsKey(paramClassLoader)) {
+                return (T) classLoaderKeyAndClassValueMap.get(paramClassLoader);
+            }
+            ClassLoader parentClassLoader = paramClassLoader == null ? null : paramClassLoader.getParent();
+            while ((parentClassLoader != null) && (!parentClassLoader.getClass().equals(ClassLoader.class))) {
+                if (classLoaderKeyAndClassValueMap.containsKey(parentClassLoader)) {
+                    return (T) classLoaderKeyAndClassValueMap.get(parentClassLoader);
+                }
+                if (parentClassLoader.equals(paramClassLoader.getParent())) {
+                    break;
+                }
+                parentClassLoader = paramClassLoader.getParent();
+            }
+        }
+        return null;
+    }
+
+    public void setClazzLoaderMap(ClassLoader paramClassLoader, String className, T clazz) {
+        if (!this.classCacheMap.containsKey(className)) {
+            this.classCacheMap.put(className, new HashMap<ClassLoader, T>());
+        }
+        ((Map<ClassLoader, T>) this.classCacheMap.get(className)).put(paramClassLoader, clazz);
+    }
+
+}

src/main/java/com/alibaba/ttl/classloader/TtlExtendLoader.java

+package com.alibaba.ttl.classloader;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.URL;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import java.util.jar.JarInputStream;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public class TtlExtendLoader {
+
+    private static final TtlClassCache<Class<?>> LOAD_CLASS_TREE = new TtlClassCache<Class<?>>();
+
+    private static final Map<String, byte[]> CLASS_BYTE_MAP = new HashMap<String, byte[]>(30);
+
+    private static final Logger logger = Logger.getLogger(TtlExtendLoader.class.getName());
+
+    public static boolean isExtendLoadClass(String className) {
+        if (CLASS_BYTE_MAP.isEmpty()) {
+            init();
+        }
+        return CLASS_BYTE_MAP.containsKey(className);
+    }
+
+    public static boolean isAgentLoadClass(String className) {
+        return className.startsWith("com.alibaba.ttl.");
+    }
+
+    public static Class<?> getClazz(ClassLoader paramClassLoader, String className) {
+        return (Class<?>) LOAD_CLASS_TREE.loadClassByName(paramClassLoader, className);
+    }
+
+    public static void setClazzLoaderMap(ClassLoader paramClassLoader, String className, Class<?> paramClass) {
+        LOAD_CLASS_TREE.setClazzLoaderMap(paramClassLoader, className, paramClass);
+    }
+
+    public static byte[] getClazzBytes(String paramString) {
+        if (CLASS_BYTE_MAP.isEmpty()) {
+            init();
+        }
+        return (byte[]) CLASS_BYTE_MAP.get(paramString);
+    }
+
+    private static final int EOF = -1;
+
+    public static void init() {
+        synchronized (TtlExtendLoader.class) {
+            if (CLASS_BYTE_MAP.isEmpty()) {
+                JarInputStream jarInputStream = null;
+                try {
+                    URL agentJarUrl = TtlAgentJarUtil.getAgentJarFilePath();
+                    JarFile jarfile = TtlAgentJarUtil.getJarFileByPath(agentJarUrl);
+                    jarInputStream = new JarInputStream(agentJarUrl.openStream());
+                    JarEntry agentJarEntry = null;
+                    while ((agentJarEntry = jarInputStream.getNextJarEntry()) != null) {
+                        if (agentJarEntry.getName().endsWith(".class")) {
+                            ByteArrayOutputStream output = new ByteArrayOutputStream();
+                            copy(jarfile.getInputStream(agentJarEntry), output);
+                            byte[] classBytes = output.toByteArray();
+                            final String className = agentJarEntry.getName()
+                                    .substring(0, agentJarEntry.getName().length() - 6).replace('/', '.');
+                            CLASS_BYTE_MAP.put(className, classBytes);
+                        }
+                    }
+                } catch (Exception t) {
+                    if (logger.isLoggable(Level.WARNING)) {
+                        logger.log(Level.WARNING, "TTL exception when load class, cause: " + t.toString(), t);
+                    }
+                } finally {
+                    try {
+                        if (jarInputStream != null) {
+                            jarInputStream.close();
+                        }
+                    } catch (IOException e) {
+                    }
+                }
+            }
+        }
+    }
+
+    public static int copy(InputStream input, OutputStream output) throws IOException {
+        long count = copyLarge(input, output);
+        if (count > Integer.MAX_VALUE) {
+            return -1;
+        }
+        return (int) count;
+    }
+
+    private static final int DEFAULT_BUFFER_SIZE = 1024 * 4;
+
+    public static long copyLarge(InputStream input, OutputStream output) throws IOException {
+        return copyLarge(input, output, new byte[DEFAULT_BUFFER_SIZE]);
+    }
+
+    public static long copyLarge(InputStream input, OutputStream output, byte[] buffer) throws IOException {
+        long count = 0;
+        int n = 0;
+        while (EOF != (n = input.read(buffer))) {
+            output.write(buffer, 0, n);
+            count += n;
+        }
+        return count;
+    }
+
+}

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

+package com.alibaba.ttl.threadpool.agent;
+
+import java.io.IOException;
+
+import javassist.CannotCompileException;
+import javassist.CtClass;
+import javassist.NotFoundException;
+
+public interface JavassistDynamicMutipleClassTransformlet extends Transformlet {
+
+    boolean needTransform(CtClass clazz, String className);
+
+    void doTransform(CtClass clazz) throws NotFoundException, CannotCompileException, IOException;
+}

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

@@ -12,7 +12,8 @@ import java.io.IOException;
  * @author Jerry Lee (oldratlee at gmail dot com)
  * @since 2.5.1
  */
-public interface JavassistTransformlet {
+public interface JavassistTransformlet extends Transformlet {
+
     boolean needTransform(String className);
 
     void doTransform(CtClass clazz) throws NotFoundException, CannotCompileException, IOException;

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

+package com.alibaba.ttl.threadpool.agent;
+
+public interface Transformlet {
+
+}

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

 package com.alibaba.ttl.threadpool.agent;
 
-
-import com.alibaba.ttl.threadpool.agent.transformlet.TtlExecutorTransformlet;
-import com.alibaba.ttl.threadpool.agent.transformlet.TtlForkJoinTransformlet;
-
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.util.jar.JarFile;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import com.alibaba.ttl.classloader.TtlAgentJarUtil;
+import com.alibaba.ttl.threadpool.agent.transformlet.TtlClassloaderTransformlet;
+import com.alibaba.ttl.threadpool.agent.transformlet.TtlExecutorTransformlet;
+import com.alibaba.ttl.threadpool.agent.transformlet.TtlForkJoinTransformlet;
+
+public class TtlAgent {
 
-/**
- * TTL Java Agent.
- *
- * @author Jerry Lee (oldratlee at gmail dot com)
- * @see <a href="https://docs.oracle.com/javase/8/docs/api/java/lang/instrument/package-summary.html">The mechanism for instrumentation</a>
- * @since 0.9.0
- */
-public final class TtlAgent {
     private static final Logger logger = Logger.getLogger(TtlAgent.class.getName());
 
     private TtlAgent() {
         throw new InstantiationError("Must not instantiate this class");
     }
 
-    public static void premain(String agentArgs, Instrumentation inst) {
+    public static void premain(String agentArgs, Instrumentation paramInstrumentation) {
         try {
-
-            logger.info("[TtlAgent.premain] begin, agentArgs: " + agentArgs + ", Instrumentation: " + inst);
-
-            @SuppressWarnings("unchecked") ClassFileTransformer transformer = new TtlTransformer(TtlExecutorTransformlet.class, TtlForkJoinTransformlet.class);
-            inst.addTransformer(transformer, true);
+            //even api can append agent jar to xbootclasspath,we still need enhance subclassloaders,
+            //because in tomcat's WebAppBaseClassLoader, our business class such as filter serlvet still can't load TransmittableThreadLocal
+            JarFile localJarFile = TtlAgentJarUtil.getJarFileByPath(TtlAgentJarUtil.getAgentJarFilePath());
+            paramInstrumentation.appendToBootstrapClassLoaderSearch(localJarFile);
+            paramInstrumentation.appendToSystemClassLoaderSearch(localJarFile);
+            logger.info("[TtlAgent.premain] begin, agentArgs: " + agentArgs + ", Instrumentation: " + paramInstrumentation);
+
+            @SuppressWarnings("unchecked")
+            ClassFileTransformer transformer = new TtlTransformer(
+                    TtlClassloaderTransformlet.class, 
+                    TtlExecutorTransformlet.class,
+                    TtlForkJoinTransformlet.class);
+
+            paramInstrumentation.addTransformer(transformer, true);
             logger.info("[TtlAgent.premain] addTransformer " + transformer.getClass() + " success");
-
             logger.info("[TtlAgent.premain] end");
 
         } catch (Exception e) {
@@ -41,6 +44,8 @@ public final class TtlAgent {
                 logger.log(Level.SEVERE, msg, e);
             }
             throw new IllegalStateException(msg, e);
+
         }
+
     }
 }

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

@@ -25,15 +25,24 @@ public class TtlTransformer implements ClassFileTransformer {
 
     private static final byte[] EMPTY_BYTE_ARRAY = {};
 
-    @SuppressWarnings("unchecked")
-    final List<JavassistTransformlet> transformletList = new ArrayList();
+    final List<JavassistTransformlet> transformletList = new ArrayList<JavassistTransformlet>();
 
-    @SuppressWarnings("unchecked")
-    public TtlTransformer(Class<? extends JavassistTransformlet>... transformletClasses) throws Exception {
-        for (Class<? extends JavassistTransformlet> transformletClass : transformletClasses) {
-            final JavassistTransformlet transformlet = transformletClass.getConstructor().newInstance();
-            transformletList.add(transformlet);
+    final List<JavassistDynamicMutipleClassTransformlet> dynamicTransformletList 
+    	= new ArrayList<JavassistDynamicMutipleClassTransformlet>();
 
+    public TtlTransformer(Class<? extends Transformlet>... transformletClasses ) throws Exception {
+        for (Class<? extends Transformlet> transformletClass : transformletClasses) {
+            
+            final Transformlet transformlet = transformletClass.getConstructor().newInstance();
+            
+            if(transformlet instanceof JavassistTransformlet){
+                JavassistTransformlet jtransformlet =(JavassistTransformlet) transformlet;
+                transformletList.add(jtransformlet);
+            }else if(transformlet instanceof JavassistDynamicMutipleClassTransformlet){
+                JavassistDynamicMutipleClassTransformlet dynamictransformlet =(JavassistDynamicMutipleClassTransformlet) transformlet;
+                dynamicTransformletList.add(dynamictransformlet);
+            }
+            
             logger.info("[TtlTransformer] add Transformlet " + transformletClass + " success");
         }
     }
@@ -48,6 +57,16 @@ public class TtlTransformer implements ClassFileTransformer {
             }
 
             final String className = toClassName(classFile);
+            
+            for (JavassistDynamicMutipleClassTransformlet dynamicTransformlet : dynamicTransformletList) {
+            	final CtClass clazz = getCtClass(classFileBuffer, loader);
+            	if (dynamicTransformlet.needTransform(clazz, className )) {
+                    logger.info("Dynamic Transforming class " + className);
+                    dynamicTransformlet.doTransform(clazz);
+                    return clazz.toBytecode();
+                }
+            }
+            
             for (JavassistTransformlet transformlet : transformletList) {
                 if (transformlet.needTransform(className)) {
                     logger.info("Transforming class " + className);
@@ -56,6 +75,7 @@ public class TtlTransformer implements ClassFileTransformer {
                     return clazz.toBytecode();
                 }
             }
+            
         } catch (Throwable t) {
             String msg = "Fail to transform class " + classFile + ", cause: " + t.toString();
             if (logger.isLoggable(Level.SEVERE)) {

src/main/java/com/alibaba/ttl/threadpool/agent/transformlet/TtlClassloaderTransformlet.java

+package com.alibaba.ttl.threadpool.agent.transformlet;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.logging.Logger;
+
+import com.alibaba.ttl.threadpool.agent.JavassistDynamicMutipleClassTransformlet;
+
+import javassist.CannotCompileException;
+import javassist.CtClass;
+import javassist.CtMethod;
+import javassist.CtNewMethod;
+import javassist.NotFoundException;
+
+public class TtlClassloaderTransformlet implements JavassistDynamicMutipleClassTransformlet {
+
+    private static final Logger logger = Logger.getLogger(TtlClassloaderTransformlet.class.getName());
+
+    @Override
+    public boolean needTransform(CtClass clazz, String className) {
+        List<CtClass> superClassList = new ArrayList<CtClass>();
+        addAllSuperClass(clazz, superClassList);
+
+        //All the sub class of java.lang.ClassLoader will be enhanced
+        //Here we enhance ClassLoader's sub classes，
+        //But in fact, only 3rd part sub classes such as Tomcat WebAppBaseClassLoader can be enhanced
+        //In jdk, bootstrap ext app, this 3 classloaders can't be enhanced
+        //I found this hint by system.out.prinln all the class in premain
+        //This hints that jdk's own 3 classloader can't be enhanced and it is already loader，
+        //JDK's own classloaders is loaded before agent started
+        for (Iterator<CtClass> iterator = superClassList.iterator(); iterator.hasNext();) {
+            CtClass ctClass = (CtClass) iterator.next();
+            if ("java.lang.ClassLoader".equals(ctClass.getName())) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public void doTransform(CtClass clazz) throws NotFoundException, CannotCompileException, IOException {
+        final String className = clazz.getName();
+        // enhance two loadClass method
+        // if subclass doesn't have any load class method override, then ignore it
+        // as method does not search the superclasses, so we only modify subclass
+        String loadClass_methodName = "loadClass";
+        final CtMethod[] loadClassMethods = clazz.getDeclaredMethods(loadClass_methodName);
+        if (loadClassMethods != null) {
+            for (CtMethod loadClassMethod : loadClassMethods) {
+                final CtMethod new_loadClassMethod = CtNewMethod.copy(loadClassMethod, loadClass_methodName, clazz,
+                        null);
+                final String original_load_method_rename = "original$loadClass$method$renamed$by$ttl";
+                loadClassMethod.setName(original_load_method_rename);// rename loadClass
+
+                CtClass[] parameterTypes = loadClassMethod.getParameterTypes();
+                String codeWeaveParameter = ((parameterTypes.length == 1) ? ("$1") : ("$1,$2"));
+
+                final String code = "{\n" + 
+                        "boolean isAgentClass = com.alibaba.ttl.classloader.TtlExtendLoader.isAgentLoadClass($1);\n"+ 
+                        "if(isAgentClass){\n"+ 
+                        "    boolean isExtendloadedClass = com.alibaba.ttl.classloader.TtlExtendLoader.isExtendLoadClass($1);\n"+ 
+                        "    if(isExtendloadedClass){\n  " +
+                        "        Class clazzTtlInner = com.alibaba.ttl.classloader.TtlExtendLoader.getClazz(this, $1);\n "+ 
+                        "        if(clazzTtlInner==null){\n "+ 
+                        "            byte[] clazzBytes = com.alibaba.ttl.classloader.TtlExtendLoader.getClazzBytes($1);\n"+ 
+                        "            if(clazzBytes==null){ return " + original_load_method_rename + "(" + codeWeaveParameter + ");  }\n         "+
+                        "                Class resultClass = this.defineClass($1, clazzBytes, 0, clazzBytes.length);\n   "+
+                        "                com.alibaba.ttl.classloader.TtlExtendLoader.setClazzLoaderMap(this, $1, resultClass );  \n   "+
+                        "                return resultClass;\n "+ 
+                        "            }else{\n"+
+                        "                return  clazzTtlInner;\n"+ 
+                        "            }\n "+
+                        "        }else{\n"+
+                        "            return " + original_load_method_rename + "(" + codeWeaveParameter + "); \n "+
+                        "        }\n"+
+                        "    }else{  return " + original_load_method_rename + "(" + codeWeaveParameter + "); }\n" + 
+                        "}\n";
+                
+                new_loadClassMethod.setBody(code);
+                clazz.addMethod(new_loadClassMethod);
+                logger.info(
+                        "insert code around method " + new_loadClassMethod + " of class " + className + ": " + code);
+            }
+
+        }
+    }
+
+    private void addAllSuperClass(CtClass clazz, List<CtClass> superClassList) {
+        CtClass supserClazz = null;
+        try {
+            supserClazz = clazz.getSuperclass();
+        } catch (NotFoundException e) {
+        }
+
+        if (supserClazz != null) {
+            superClassList.add(supserClazz);
+            addAllSuperClass(supserClazz, superClassList);
+        }
+    }
+}