Some minor improvements to KeyedLocker:

(i) Change the semantics of KeyedLocker.AutoUnlocker#close such that it can be called at most once per AutoUnlocker instance.
(ii) Change KeyedLocker.AutoUnlocker#close to throw a IllegalUnlockException (RuntimeException) on error, rather than leave the behavior intentionally underspecified.
(iii) explicitly mention in AutoLocker#lock that a thread can call lock(k) multiple times before unlocking. Combined with (i), this implies that KeyedLocker#lock implementations will want to return fresh AutoUnlocker instances.

These semantics are bit nicer to use anyway, but I also want them because I will soon be introducing KeyedLocker#lockBatch, and it's much easier to specify that given the above.

--
MOS_MIGRATED_REVID=90259645

2 files changed, 75 insertions, 21 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/KeyedLocker.java b/src/main/java/com/google/devtools/build/lib/concurrent/KeyedLocker.java
index 8228f5b236..720db4d716 100644
--- a/src/main/java/com/google/devtools/build/lib/concurrent/KeyedLocker.java
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/KeyedLocker.java
@@ -13,7 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.concurrent;
 
-import javax.annotation.concurrent.ThreadSafe;
+import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 
 /** A keyed store of locks. */
 @ThreadSafe
@@ -21,12 +21,21 @@ public interface KeyedLocker<K> {
   /** Used to yield access to the implicit lock granted by {@link #lock}. */
   @ThreadSafe
   interface AutoUnlocker extends AutoCloseable {
+    /** Exception used to indicate illegal use of {@link AutoUnlocker#close}. */
+    public static class IllegalUnlockException extends RuntimeException {
+      public IllegalUnlockException(String msg) {
+        super(msg);
+      }
+    }
+
     /**
-     * If this was returned by {@code lock(k)}, yields exclusive access to {@code k}.
+     * Closes the {@link AutoUnlocker} instance. If this instance was the last unclosed one
+     * returned by {@code lock(k)} owned by then the current thread, then exclusive access to
+     * {@code k} is yielded.
      *
-     * <p>This method should be called at most once, and may only be called by the same thread that
-     * acquired the {@link AutoUnlocker} via {@link #lock}. Implementations are free to do anything
-     * if this is violated.
+     * <p>This method may only be called at most once per {@link AutoUnlocker} instance and must
+     * be called by the same thread that acquired the {@link AutoUnlocker} via {@link #lock}.
+     * Otherwise, an {@link IllegalUnlockException} is thrown.
      */
     @Override
     void close();
@@ -34,8 +43,13 @@ public interface KeyedLocker<K> {
 
   /**
    * Blocks the current thread until it has exclusive access to do things with {@code k} and
-   * returns a {@link AutoUnlocker} instance for yielding the implicit lock. The intended usage
-   * is:
+   * returns a {@link AutoUnlocker} instance for yielding the implicit lock.
+   *
+   * <p>Notably, this means that a thread is allowed to call {@code lock(k)} again before calling
+   * {@link AutoUnlocker#close} for the first call to {@code lock(k)}. Each call to {@link #lock}
+   * will return a different {@link AutoUnlocker} instance.
+   *
+   * <p>The intended usage is:
    *
    * <pre>
    * {@code
@@ -44,6 +58,28 @@ public interface KeyedLocker<K> {
    * }
    * }
    * </pre>
+   *
+   * <p>Note that the usual caveats about mutexes apply here, e.g. the following may deadlock:
+   *
+   * <pre>
+   * {@code
+   * // Thread A
+   * try (AutoUnlocker unlocker = locker.lock(k1)) {
+   *   // This will deadlock if Thread A already acquired a lock for k2.
+   *   try (AutoUnlocker unlocker = locker.lock(k2)) {
+   *   }
+   * }
+   * // end Thread A
+   *
+   * // Thread B
+   * try (AutoUnlocker unlocker = locker.lock(k2)) {
+   *   // This will deadlock if Thread A already acquired a lock for k1.
+   *   try (AutoUnlocker unlocker = locker.lock(k1)) {
+   *   }
+   * }
+   * // end Thread B
+   * }
+   * </pre>
    */
   AutoUnlocker lock(K key);
 }
\ No newline at end of file
diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java b/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java
index 4efa90f567..0656285fae 100644
--- a/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java
@@ -16,9 +16,11 @@ package com.google.devtools.build.lib.concurrent;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ConcurrentHashMultiset;
 import com.google.common.util.concurrent.Striped;
+import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -26,6 +28,7 @@ import java.util.concurrent.locks.ReentrantLock;
  * An implementation of {@link KeyedLocker} that uses ref counting to efficiently only store locks
  * that are live.
  */
+@ThreadSafe
 public class RefCountedMultisetKeyedLocker<K> implements KeyedLocker<K> {
   // Multiset of keys that have threads waiting on a lock or using a lock.
   private final ConcurrentHashMultiset<K> waiters = ConcurrentHashMultiset.<K>create();
@@ -36,19 +39,30 @@ public class RefCountedMultisetKeyedLocker<K> implements KeyedLocker<K> {
   private final Striped<Lock> waitersLocks = Striped.lazyWeakLock(NUM_STRIPES);
 
   // Map of key to current lock, for keys that have at least one waiting or live thread.
-  private final ConcurrentMap<K, RefCountedLockImpl> locks = new ConcurrentHashMap<>();
+  private final ConcurrentMap<K, ReentrantLock> locks = new ConcurrentHashMap<>();
 
-  private class RefCountedLockImpl extends ReentrantLock implements AutoUnlocker {
+  private class RefCountedAutoUnlocker implements AutoUnlocker {
     private final K key;
+    private final ReentrantLock lock;
+    private final AtomicBoolean closeCalled = new AtomicBoolean(false);
 
-    private RefCountedLockImpl(K key) {
+    private RefCountedAutoUnlocker(K key, ReentrantLock lock) {
       this.key = key;
+      this.lock = lock;
     }
 
     @Override
     public void close() {
-      Preconditions.checkState(isHeldByCurrentThread(), "For key %s, 'close' can be called at most "
-          + "once and the calling thread must be the one that acquired the AutoUnlocker", key);
+      if (closeCalled.getAndSet(true)) {
+        String msg = String.format("For key %s, 'close' can be called at most once per "
+            + "AutoUnlocker instance", key);
+        throw new IllegalUnlockException(msg);
+      }
+      if (!lock.isHeldByCurrentThread()) {
+        String msg = String.format("For key %s, the calling thread to 'close' must be the one "
+            + "that acquired the AutoUnlocker", key);
+        throw new IllegalUnlockException(msg);
+      }
       try {
         Lock waitersLock = waitersLocks.get(key);
         try {
@@ -57,42 +71,46 @@ public class RefCountedMultisetKeyedLocker<K> implements KeyedLocker<K> {
           // count is 0.
           waiters.remove(key);
           if (waiters.count(key) == 0) {
-            // No other thread is waiting to access this key, so we garbage collect the lock.
-            Preconditions.checkState(locks.remove(key, this), key);
+            // No other thread is waiting to access this key, nor does the current thread have
+            // another AutoUnlocker instance, so we garbage collect the lock.
+            Preconditions.checkState(locks.remove(key, lock), key);
           }
         } finally {
           waitersLock.unlock();
         }
       } finally {
-        unlock();
+        lock.unlock();
       }
     }
   }
 
   @Override
   public AutoUnlocker lock(K key) {
-    RefCountedLockImpl newLock = new RefCountedLockImpl(key);
+    ReentrantLock newLock = new ReentrantLock();
     // Pre-lock our fresh lock, in case we win the race to get access to 'key'.
     newLock.lock();
     Lock waitersLock = waitersLocks.get(key);
     try {
       waitersLock.lock();
       // Add us to the set of waiters, so that in case we lose the race to access 'key', the winner
-      // will know that we are waiting.
+      // will know that we are waiting. If we already have access to 'key', this simply bumps up
+      // the ref count.
       waiters.add(key);
     } finally {
       waitersLock.unlock();
     }
-    RefCountedLockImpl lock;
+    ReentrantLock lock;
     lock = locks.putIfAbsent(key, newLock);
     if (lock != null) {
-      // Another thread won the race to get access to 'key', so we wait for our turn.
       Preconditions.checkState(lock != newLock);
       newLock.unlock();
+      // Either another thread won the race to get access to 'key', or we already have exclusive
+      // access to 'key'. Either way, we lock; in the former case we wait for our turn and in the
+      // latter case the lock's implicit counter is incremented.
       lock.lock();
-      return lock;
+      return new RefCountedAutoUnlocker(key, lock);
     }
     // We won the race, so the current lock for 'key' is the one we locked and inserted.
-    return newLock;
+    return new RefCountedAutoUnlocker(key, newLock);
   }
 }