no more EnableTestable annotation in agent

testable-agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java

@@ -12,6 +12,5 @@ public class ConstPool {
     public static final String TEST_POSTFIX = "Test";
     public static final String TESTABLE_INJECT_REF = "_testableInternalRef";
 
-    public static final String ENABLE_TESTABLE = "com.alibaba.testable.core.annotation.EnableTestable";
     public static final String TESTABLE_INJECT = "com.alibaba.testable.core.annotation.TestableInject";
 }

testable-agent/src/main/java/com/alibaba/testable/agent/handler/TestClassHandler.java

@@ -41,7 +41,7 @@ public class TestClassHandler extends BaseClassHandler {
 
     private void transformMethod(ClassNode cn, MethodNode mn) {
         handleAnnotation(cn, mn);
-        handleInstruction(cn, mn);
+        handleInstruction(mn);
     }
 
     private void handleAnnotation(ClassNode cn, MethodNode mn) {
@@ -61,7 +61,7 @@ public class TestClassHandler extends BaseClassHandler {
         }
     }
 
-    private void handleInstruction(ClassNode cn, MethodNode mn) {
+    private void handleInstruction(MethodNode mn) {
         AbstractInsnNode[] instructions = mn.instructions.toArray();
         int i = 0;
         do {

testable-agent/src/main/java/com/alibaba/testable/agent/tool/ComparableWeakRef.java

@@ -0,0 +1,33 @@
+package com.alibaba.testable.agent.tool;
+
+import java.lang.ref.WeakReference;
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
+
+/**
+ * @author flin
+ */
+public class ComparableWeakRef<T> extends WeakReference<T> {
+
+    public ComparableWeakRef(T referent) {
+        super(referent);
+    }
+
+    static public <T> Set<ComparableWeakRef<T>> getWeekHashSet() {
+        return Collections.newSetFromMap(new WeakHashMap<ComparableWeakRef<T>, Boolean>());
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj instanceof WeakReference) {
+            return this.get().equals(((WeakReference)obj).get());
+        }
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        return this.get().hashCode();
+    }
+}

testable-agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java

@@ -5,6 +5,7 @@ import com.alibaba.testable.agent.handler.SourceClassHandler;
 import com.alibaba.testable.agent.handler.TestClassHandler;
 import com.alibaba.testable.agent.model.ImmutablePair;
 import com.alibaba.testable.agent.model.MethodInfo;
+import com.alibaba.testable.agent.tool.ComparableWeakRef;
 import com.alibaba.testable.agent.util.ClassUtil;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.Type;
@@ -17,7 +18,6 @@ import java.lang.instrument.ClassFileTransformer;
 import java.net.URLClassLoader;
 import java.security.ProtectionDomain;
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 
@@ -29,28 +29,26 @@ import static com.alibaba.testable.agent.util.ClassUtil.toSlashSeparatedName;
  */
 public class TestableClassTransformer implements ClassFileTransformer {
 
-    private final Set<String> loadedClassNames = new HashSet<String>();
+    private final Set<ComparableWeakRef<String>> loadedClassNames = ComparableWeakRef.getWeekHashSet();
     private static final String TARGET_CLASS = "targetClass";
     private static final String TARGET_METHOD = "targetMethod";
 
+    @Override
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                             ProtectionDomain protectionDomain, byte[] classFileBuffer) {
-        if (isSystemClass(loader, className) || loadedClassNames.contains(className)) {
+        if (isSystemClass(loader, className) || loadedClassNames.contains(new ComparableWeakRef<String>(className))) {
             // Ignore system class and reloaded class
             return null;
         }
-
-        List<String> annotations = ClassUtil.getAnnotations(className);
-        List<String> testAnnotations = ClassUtil.getAnnotations(ClassUtil.getTestClassName(className));
         try {
-            if (testAnnotations.contains(ConstPool.ENABLE_TESTABLE)) {
+            if (shouldTransformAsSourceClass(className)) {
                 // it's a source class with testable enabled
-                loadedClassNames.add(className);
+                loadedClassNames.add(new ComparableWeakRef<String>(className));
                 List<MethodInfo> injectMethods = getTestableInjectMethods(ClassUtil.getTestClassName(className));
                 return new SourceClassHandler(injectMethods).getBytes(classFileBuffer);
-            } else if (annotations.contains(ConstPool.ENABLE_TESTABLE)) {
+            } else if (shouldTransformAsTestClass(className)) {
                 // it's a test class with testable enabled
-                loadedClassNames.add(className);
+                loadedClassNames.add(new ComparableWeakRef<String>(className));
                 return new TestClassHandler().getBytes(classFileBuffer);
             }
         } catch (IOException e) {
@@ -59,6 +57,15 @@ public class TestableClassTransformer implements ClassFileTransformer {
         return null;
     }
 
+    private boolean shouldTransformAsSourceClass(String className) {
+        return ClassUtil.anyMethodHasAnnotation(ClassUtil.getTestClassName(className), ConstPool.TESTABLE_INJECT);
+    }
+
+    private boolean shouldTransformAsTestClass(String className) {
+        return className.endsWith(ConstPool.TEST_POSTFIX) &&
+            ClassUtil.anyMethodHasAnnotation(className, ConstPool.TESTABLE_INJECT);
+    }
+
     private boolean isSystemClass(ClassLoader loader, String className) {
         return !(loader instanceof URLClassLoader) || null == className || className.startsWith("jdk/");
     }

testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java

@@ -1,14 +1,13 @@
 package com.alibaba.testable.agent.util;
 
 import com.alibaba.testable.agent.constant.ConstPool;
+import com.alibaba.testable.agent.tool.ComparableWeakRef;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.MethodNode;
 
-import java.util.ArrayList;
+import java.util.*;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
 
 /**
  * @author flin
@@ -29,6 +28,8 @@ public class ClassUtil {
     private static final char TYPE_ARRAY = '[';
 
     private static final Map<Character, String> TYPE_MAPPING = new HashMap<Character, String>();
+    private static final Map<ComparableWeakRef<String>, Boolean> loadedClass =
+        new WeakHashMap<ComparableWeakRef<String>, Boolean>();
 
     static {
         TYPE_MAPPING.put(TYPE_BYTE, "java/lang/Byte");
@@ -42,21 +43,33 @@ public class ClassUtil {
     }
 
     /**
-     * Get annotation on class definition
+     * Check whether any method in specified class has specified annotation
      * @param className class that need to explore
+     * @param annotationName annotation to look for
      */
-    public static List<String> getAnnotations(String className) {
+    public static boolean anyMethodHasAnnotation(String className, String annotationName) {
+        Boolean found = loadedClass.get(new ComparableWeakRef<String>(className));
+        if (found != null) {
+            return found;
+        }
         try {
-            List<String> annotations = new ArrayList<String>();
             ClassNode cn = new ClassNode();
             new ClassReader(className).accept(cn, 0);
-            for (AnnotationNode an : cn.visibleAnnotations) {
+            for (MethodNode mn : cn.methods) {
-                annotations.add(toDotSeparateFullClassName(an.desc));
+                if (mn.visibleAnnotations != null) {
+                    for (AnnotationNode an : mn.visibleAnnotations) {
+                        if (toDotSeparateFullClassName(an.desc).equals(annotationName)) {
+                            loadedClass.put(new ComparableWeakRef<String>(className), true);
+                            return true;
+                        }
+                    }
+                }
             }
-            return annotations;
         } catch (Exception e) {
-            return new ArrayList<String>();
+            // ignore
         }
+        loadedClass.put(new ComparableWeakRef<String>(className), false);
+        return false;
     }
 
     /**

testable-agent/src/test/java/com/alibaba/testable/agent/tool/ComparableWeakRefTest.java

@@ -0,0 +1,38 @@
+package com.alibaba.testable.agent.tool;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.*;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class ComparableWeakRefTest {
+
+    @Test
+    void should_works_with_weak_hash_map() throws Exception {
+        Map<ComparableWeakRef<String>, String> map = new WeakHashMap<ComparableWeakRef<String>, String>();
+        map.put(new ComparableWeakRef<String>("a"), "bc");
+        map.put(new ComparableWeakRef<String>("ab"), "ac");
+        map.put(new ComparableWeakRef<String>("abc"), "bc");
+        assertEquals(3, map.size());
+        assertEquals("bc", map.get(new ComparableWeakRef<String>("abc")));
+        System.gc();
+        Thread.sleep(100);
+        assertEquals(0, map.size());
+        assertNull(map.get(new ComparableWeakRef<String>("abc")));
+    }
+
+    @Test
+    void should_works_with_hash_set() throws Exception {
+        Set<ComparableWeakRef<String>> set = ComparableWeakRef.getWeekHashSet();
+        set.add(new ComparableWeakRef<String>("a"));
+        set.add(new ComparableWeakRef<String>("ab"));
+        set.add(new ComparableWeakRef<String>("abc"));
+        assertEquals(3, set.size());
+        assertTrue(set.contains(new ComparableWeakRef<String>("ab")));
+        System.gc();
+        Thread.sleep(100);
+        assertEquals(0, set.size());
+        assertFalse(set.contains(new ComparableWeakRef<String>("ab")));
+    }
+}

testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java

@@ -2,15 +2,19 @@ package com.alibaba.testable.agent.util;
 
 import org.junit.jupiter.api.Test;
 
+import java.lang.ref.WeakReference;
+import java.util.Map;
+import java.util.WeakHashMap;
+
 import static org.junit.jupiter.api.Assertions.*;
 
 class ClassUtilTest {
 
     @Test
     void should_able_to_get_annotation() {
-        assertEquals(0, ClassUtil.getAnnotations("class.not.exist").size());
+        assertEquals(false, ClassUtil.anyMethodHasAnnotation("class.not.exist", ""));
-        assertEquals(0, ClassUtil.getAnnotations("com.alibaba.testable.agent.util.ClassUtil").size());
+        assertEquals(false, ClassUtil.anyMethodHasAnnotation("org.junit.jupiter.api.Assertions", "annotation.not.exist"));
-        assertEquals("org.apiguardian.api.API", ClassUtil.getAnnotations("org.junit.jupiter.api.Assertions").get(0));
+        assertEquals(true, ClassUtil.anyMethodHasAnnotation("org.junit.jupiter.api.Assertions", "org.apiguardian.api.API"));
     }
 
     @Test