Performance improvement in annotation lookup

src/main/java/org/testng/internal/annotations/JDK15AnnotationFinder.java

@@ -3,6 +3,7 @@
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -46,10 +47,20 @@
  * @author <a href="mailto:[email protected]">Cedric Beust</a>
  */
 public class JDK15AnnotationFinder implements IAnnotationFinder {
+
+  private static final IAnnotation NULL_MARKER = new IAnnotation() {
+    @Override
+    protected Object clone() throws CloneNotSupportedException {
+      return super.clone();
+    }
+  };
+
   private JDK15TagFactory m_tagFactory = new JDK15TagFactory();
-  private Map<Class<? extends IAnnotation>, Class<? extends Annotation>> m_annotationMap =
-      new ConcurrentHashMap<>();
-  private Map<Pair<Annotation, ?>, IAnnotation> m_annotations = new ConcurrentHashMap<>();
+  // readonly map. does not need ConcurrentHashMap
+  private final Map<Class<? extends IAnnotation>, Class<? extends Annotation>> m_annotationMap =
+      new HashMap<>();
+
+  private Map<Pair<AnnotationWrapper, ?>, IAnnotation> m_annotations = new ConcurrentHashMap<>();
 
   private IAnnotationTransformer m_transformer = null;
 
@@ -106,7 +117,7 @@ public <A extends IAnnotation> A findAnnotation(Method m, Class<A> annotationCla
     }
     Annotation annotation = m.getAnnotation(a);
     return findAnnotation(m.getDeclaringClass(), annotation, annotationClass, null, null, m,
-        new Pair<>(annotation, m));
+        new Pair<>(new AnnotationWrapper(annotation), m));
   }
 
   @Override
@@ -128,7 +139,7 @@ public <A extends IAnnotation> A findAnnotation(ITestNGMethod tm, Class<A> annot
       annotation = testClass.getAnnotation(a);
     }
     return findAnnotation(testClass, annotation, annotationClass, null, null, m,
-        new Pair<>(annotation, m));
+        new Pair<>(new AnnotationWrapper(annotation), m));
   }
 
   private void transform(IAnnotation a, Class<?> testClass,
@@ -176,7 +187,7 @@ public <A extends IAnnotation> A findAnnotation(Class<?> cls, Class<A> annotatio
     }
     Annotation annotation = findAnnotationInSuperClasses(cls, a);
     return findAnnotation(cls, annotation, annotationClass, cls, null, null,
-        new Pair<>(annotation, annotationClass));
+        new Pair<>(new AnnotationWrapper(annotation), annotationClass));
   }
 
   @Override
@@ -188,21 +199,23 @@ public <A extends IAnnotation> A findAnnotation(Constructor<?> cons, Class<A> an
     }
     Annotation annotation = cons.getAnnotation(a);
     return findAnnotation(cons.getDeclaringClass(), annotation, annotationClass, null, cons, null,
-        new Pair<>(annotation, cons));
+        new Pair<>(new AnnotationWrapper(annotation), cons));
   }
 
   private <A extends IAnnotation> A findAnnotation(Class cls, Annotation a,
       Class<A> annotationClass, Class<?> testClass,
-      Constructor<?> testConstructor, Method testMethod, Pair<Annotation, ?> p) {
+      Constructor<?> testConstructor, Method testMethod, Pair<AnnotationWrapper, ?> p) {
     if (a == null) {
       return null;
     }
 
     IAnnotation result = m_annotations.get(p);
     if (result == null) {
       result = m_tagFactory.createTag(cls, a, annotationClass, m_transformer);
-      m_annotations.put(p, result);
+      m_annotations.put(p, result != null ? result : NULL_MARKER);
       transform(result, testClass, testConstructor, testMethod);
+    } else if (result == NULL_MARKER) {
+      result = null;
     }
     //noinspection unchecked
     return (A) result;
@@ -244,4 +257,26 @@ private String[] optionalValues(Annotation[][] annotations) {
     }
     return result;
   }
+
+  /**
+   * Wrapper class to override {@link Annotation#hashCode()} and {@link Annotation#equals(Object)}
+   * for performance improvement.
+   */
+  private final static class AnnotationWrapper {
+    private final Annotation annotation;
+
+    public AnnotationWrapper(Annotation annotation) {
+      this.annotation = annotation;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      return obj instanceof AnnotationWrapper && annotation == ((AnnotationWrapper) obj).annotation;
+    }
+
+    @Override
+    public int hashCode() {
+      return System.identityHashCode(annotation);
+    }
+  }
 }