Optimize the runtime library for try-with-resources, by reducing the

granularity of locks. Now it uses a customized concurrent weak identity hash
map.

RELNOTES: n/a
PiperOrigin-RevId: 155688279
GitOrigin-RevId: 3bf15e757a801ff813370aaa01ebc9143a8834d4
Change-Id: I0e52abcd7979e59f22be76f37379b06cc470f343

2 files changed, 126 insertions, 34 deletions

diff --git a/java/com/google/devtools/build/android/desugar/TryWithResourcesRewriter.java b/java/com/google/devtools/build/android/desugar/TryWithResourcesRewriter.java
index 38255a1..4093e12 100644
--- a/java/com/google/devtools/build/android/desugar/TryWithResourcesRewriter.java
+++ b/java/com/google/devtools/build/android/desugar/TryWithResourcesRewriter.java
@@ -52,6 +52,8 @@ public class TryWithResourcesRewriter extends ClassVisitor {
       ImmutableSet.of(
           THROWABLE_EXTENSION_INTERNAL_NAME,
           THROWABLE_EXTENSION_INTERNAL_NAME + "$AbstractDesugaringStrategy",
+          THROWABLE_EXTENSION_INTERNAL_NAME + "$ConcurrentWeakIdentityHashMap",
+          THROWABLE_EXTENSION_INTERNAL_NAME + "$ConcurrentWeakIdentityHashMap$WeakKey",
           THROWABLE_EXTENSION_INTERNAL_NAME + "$MimicDesugaringStrategy",
           THROWABLE_EXTENSION_INTERNAL_NAME + "$NullDesugaringStrategy",
           THROWABLE_EXTENSION_INTERNAL_NAME + "$ReuseDesugaringStrategy");
diff --git a/java/com/google/devtools/build/android/desugar/runtime/ThrowableExtension.java b/java/com/google/devtools/build/android/desugar/runtime/ThrowableExtension.java
index 3581fe8..365884b 100644
--- a/java/com/google/devtools/build/android/desugar/runtime/ThrowableExtension.java
+++ b/java/com/google/devtools/build/android/desugar/runtime/ThrowableExtension.java
@@ -15,10 +15,13 @@ package com.google.devtools.build.android.desugar.runtime;
 
 import java.io.PrintStream;
 import java.io.PrintWriter;
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
-import java.util.ArrayList;
 import java.util.List;
-import java.util.WeakHashMap;
+import java.util.Vector;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * This is an extension class for java.lang.Throwable. It emulates the methods
@@ -27,7 +30,7 @@ import java.util.WeakHashMap;
  *
  * <p>Note that the Desugar should avoid desugaring this class.
  */
-public class ThrowableExtension {
+public final class ThrowableExtension {
 
   static final AbstractDesugaringStrategy STRATEGY;
   /**
@@ -142,7 +145,7 @@ public class ThrowableExtension {
   }
 
   /** This strategy just delegates all the method calls to java.lang.Throwable. */
-  static class ReuseDesugaringStrategy extends AbstractDesugaringStrategy {
+  static final class ReuseDesugaringStrategy extends AbstractDesugaringStrategy {
 
     @Override
     public void addSuppressed(Throwable receiver, Throwable suppressed) {
@@ -171,17 +174,14 @@ public class ThrowableExtension {
   }
 
   /** This strategy mimics the behavior of suppressed exceptions with a map. */
-  static class MimicDesugaringStrategy extends AbstractDesugaringStrategy {
+  static final class MimicDesugaringStrategy extends AbstractDesugaringStrategy {
 
-    public static final String SUPPRESSED_PREFIX = "Suppressed: ";
-    private final WeakHashMap<Throwable, List<Throwable>> map = new WeakHashMap<>();
+    static final String SUPPRESSED_PREFIX = "Suppressed: ";
+    private final ConcurrentWeakIdentityHashMap map = new ConcurrentWeakIdentityHashMap();
 
     /**
      * Suppress an exception. If the exception to be suppressed is {@receiver} or {@null}, an
      * exception will be thrown.
-     *
-     * @param receiver
-     * @param suppressed
      */
     @Override
     public void addSuppressed(Throwable receiver, Throwable suppressed) {
@@ -191,23 +191,17 @@ public class ThrowableExtension {
       if (suppressed == null) {
         throw new NullPointerException("The suppressed exception cannot be null.");
       }
-      synchronized (this) {
-        List<Throwable> list = map.get(receiver);
-        if (list == null) {
-          list = new ArrayList<>(1);
-          map.put(receiver, list);
-        }
-        list.add(suppressed);
-      }
+      // The returned list is a synchrnozed list.
+      map.get(receiver, /*createOnAbsence=*/true).add(suppressed);
     }
 
     @Override
-    public synchronized Throwable[] getSuppressed(Throwable receiver) {
-      List<Throwable> list = map.get(receiver);
+    public Throwable[] getSuppressed(Throwable receiver) {
+      List<Throwable> list = map.get(receiver, /*createOnAbsence=*/false);
       if (list == null || list.isEmpty()) {
         return EMPTY_THROWABLE_ARRAY;
       }
-      return list.toArray(new Throwable[0]);
+      return list.toArray(EMPTY_THROWABLE_ARRAY);
     }
 
     /**
@@ -218,35 +212,131 @@ public class ThrowableExtension {
      * {@code receiver} and its suppressed exceptions are printed in two different streams.
      */
     @Override
-    public synchronized void printStackTrace(Throwable receiver) {
+    public void printStackTrace(Throwable receiver) {
       receiver.printStackTrace();
-      for (Throwable suppressed : getSuppressed(receiver)) {
-        System.err.print(SUPPRESSED_PREFIX);
-        suppressed.printStackTrace();
+      List<Throwable> suppressedList = map.get(receiver, /*createOnAbsence=*/false);
+      if (suppressedList == null) {
+        return;
+      }
+      synchronized (suppressedList) {
+        for (Throwable suppressed : suppressedList) {
+          System.err.print(SUPPRESSED_PREFIX);
+          suppressed.printStackTrace();
+        }
       }
     }
 
     @Override
-    public synchronized void printStackTrace(Throwable receiver, PrintStream stream) {
+    public void printStackTrace(Throwable receiver, PrintStream stream) {
       receiver.printStackTrace(stream);
-      for (Throwable suppressed : getSuppressed(receiver)) {
-        stream.print(SUPPRESSED_PREFIX);
-        suppressed.printStackTrace(stream);
+      List<Throwable> suppressedList = map.get(receiver, /*createOnAbsence=*/false);
+      if (suppressedList == null) {
+        return;
+      }
+      synchronized (suppressedList) {
+        for (Throwable suppressed : suppressedList) {
+          stream.print(SUPPRESSED_PREFIX);
+          suppressed.printStackTrace(stream);
+        }
       }
     }
 
     @Override
-    public synchronized void printStackTrace(Throwable receiver, PrintWriter writer) {
+    public void printStackTrace(Throwable receiver, PrintWriter writer) {
       receiver.printStackTrace(writer);
-      for (Throwable suppressed : getSuppressed(receiver)) {
-        writer.print(SUPPRESSED_PREFIX);
-        suppressed.printStackTrace(writer);
+      List<Throwable> suppressedList = map.get(receiver, /*createOnAbsence=*/false);
+      if (suppressedList == null) {
+        return;
+      }
+      synchronized (suppressedList) {
+        for (Throwable suppressed : suppressedList) {
+          writer.print(SUPPRESSED_PREFIX);
+          suppressed.printStackTrace(writer);
+        }
+      }
+    }
+  }
+
+  /** A hash map, that is concurrent, weak-key, and identity-hashing. */
+  static final class ConcurrentWeakIdentityHashMap {
+
+    private final ConcurrentHashMap<WeakKey, List<Throwable>> map =
+        new ConcurrentHashMap<>(16, 0.75f, 10);
+    private final ReferenceQueue<Throwable> referenceQueue = new ReferenceQueue<>();
+
+    /**
+     * @param throwable, the key to retrieve or create associated list.
+     * @param createOnAbsence {@code true} to create a new list if there is no value for the key.
+     * @return the associated value with the given {@code throwable}. If {@code createOnAbsence} is
+     *     {@code true}, the returned value will be non-null. Otherwise, it can be {@code null}
+     */
+    public List<Throwable> get(Throwable throwable, boolean createOnAbsence) {
+      deleteEmptyKeys();
+      WeakKey keyForQuery = new WeakKey(throwable, null);
+      List<Throwable> list = map.get(keyForQuery);
+      if (!createOnAbsence) {
+        return list;
+      }
+      if (list != null) {
+        return list;
+      }
+      List<Throwable> newValue = new Vector<>(2);
+      list = map.putIfAbsent(new WeakKey(throwable, referenceQueue), newValue);
+      return list == null ? newValue : list;
+    }
+
+    /** For testing-purpose */
+    int size() {
+      return map.size();
+    }
+
+    void deleteEmptyKeys() {
+      // The ReferenceQueue.poll() is thread-safe.
+      for (Reference<?> key = referenceQueue.poll(); key != null; key = referenceQueue.poll()) {
+        map.remove(key);
+      }
+    }
+
+    private static final class WeakKey extends WeakReference<Throwable> {
+
+      /**
+       * The hash code is used later to retrieve the entry, of which the key is the current weak
+       * key. If the referent is marked for garbage collection and is set to null, we are still able
+       * to locate the entry.
+       */
+      private final int hash;
+
+      public WeakKey(Throwable referent, ReferenceQueue<Throwable> q) {
+        super(referent, q);
+        if (referent == null) {
+          throw new NullPointerException("The referent cannot be null");
+        }
+        hash = System.identityHashCode(referent);
+      }
+
+      @Override
+      public int hashCode() {
+        return hash;
+      }
+
+      @Override
+      public boolean equals(Object obj) {
+        if (obj == null || obj.getClass() != getClass()) {
+          return false;
+        }
+        if (this == obj) {
+          return true;
+        }
+        WeakKey other = (WeakKey) obj;
+        // Note that, after the referent is garbage collected, then the referent will be null.
+        // And the equality test still holds.
+        return this.hash == other.hash && this.get() == other.get();
       }
     }
   }
 
   /** This strategy ignores all suppressed exceptions, which is how retrolambda does. */
-  static class NullDesugaringStrategy extends AbstractDesugaringStrategy {
+  static final class NullDesugaringStrategy extends AbstractDesugaringStrategy {
 
     @Override
     public void addSuppressed(Throwable receiver, Throwable suppressed) {