Distinguish between read and write locks for KeyedLocker.

--
MOS_MIGRATED_REVID=109835697

3 files changed, 50 insertions, 136 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 1d75f38b25..907237db19 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
@@ -18,7 +18,9 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 /** A keyed store of locks. */
 @ThreadSafe
 public interface KeyedLocker<K> {
-  /** Used to yield access to the implicit lock granted by {@link #lock}. */
+  /**
+   * Used to yield access to the implicit locks granted by {@link #writeLock} or {@link #readLock}.
+   */
   @ThreadSafe
   interface AutoUnlocker extends AutoCloseable {
     /** Exception used to indicate illegal use of {@link AutoUnlocker#close}. */
@@ -30,12 +32,14 @@ public interface KeyedLocker<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.
+     * returned by {@link #writeLock} with argument {@code k} owned by the current
+     * thread, then exclusive access to {@code k} is yielded. If this instance was the last unclosed
+     * one returned by {@link #readLock} with argument {@code k}, then a thread can request
+     * exclusive write access using {@link #writeLock} with argument {@code k}.
      *
      * <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.
+     * be called by the same thread that acquired the {@link AutoUnlocker} via {@link #writeLock}
+     * or {@link #readLock}. Otherwise, an {@link IllegalUnlockException} is thrown.
      */
     @Override
     void close();
@@ -45,15 +49,15 @@ 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.
    *
-   * <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>Notably, this means that a thread is allowed to call {@code writeLock(k)} again before
+   * calling {@link AutoUnlocker#close} for the first call to {@code writeLock(k)}. Each call to
+   * {@code #writeLock} will return a different {@link AutoUnlocker} instance.
    *
    * <p>The intended usage is:
    *
    * <pre>
    * {@code
-   * try (AutoUnlocker unlocker = locker.lock(k)) {
+   * try (AutoUnlocker unlocker = locker.writeLock(k)) {
    *   // Your code here.
    * }
    * }
@@ -64,22 +68,32 @@ public interface KeyedLocker<K> {
    * <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)) {
+   * try (AutoUnlocker unlocker = locker.writeLock(k1)) {
+   *   // This will deadlock if Thread B already acquired a writeLock for k2.
+   *   try (AutoUnlocker unlocker = locker.writeLock(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)) {
+   * try (AutoUnlocker unlocker = locker.writeLock(k2)) {
+   *   // This will deadlock if Thread A already acquired a writeLock for k1.
+   *   try (AutoUnlocker unlocker = locker.writeLock(k1)) {
    *   }
    * }
    * // end Thread B
    * }
    * </pre>
    */
-  AutoUnlocker lock(K key);
-}
\ No newline at end of file
+  AutoUnlocker writeLock(K key);
+
+  /**
+   * Blocks the current thread until it has access to read things that have to do with {@code k}.
+   * Multiple threads may acquire simultaneous read locks, so long as there is no thread with a
+   * write lock.
+   *
+   * <p>As with {@link #writeLock}, the same thread can call {@code readLock(k)} multiple times for
+   * the same k before closing the lock.
+   */
+  AutoUnlocker readLock(K key);
+}
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
deleted file mode 100644
index ac83f03b28..0000000000
--- a/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java
+++ /dev/null
@@ -1,112 +0,0 @@
-// Copyright 2015 The Bazel Authors. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//    http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-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 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;
-
-/**
- * An implementation of {@link KeyedLocker} that uses ref counting to efficiently only store locks
- * that are live.
- */
-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();
-
-  private static final int NUM_STRIPES = 2048;
-  // For each key, gives the striped lock to use for atomically managing the waiters on that key
-  // internally.
-  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, ReentrantLock> locks =
-      new ConcurrentHashMap<>(1024, .75f, NUM_STRIPES);
-
-  public RefCountedMultisetKeyedLocker() {
-  }
-
-  private class RefCountedAutoUnlocker implements AutoUnlocker {
-    private final K key;
-    private final ReentrantLock lock;
-    private final AtomicBoolean closeCalled = new AtomicBoolean(false);
-
-    private RefCountedAutoUnlocker(K key, ReentrantLock lock) {
-      this.key = key;
-      this.lock = lock;
-    }
-
-    @Override
-    public void close() {
-      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);
-      }
-      try {
-        Lock waitersLock = waitersLocks.get(key);
-        try {
-          waitersLock.lock();
-          // Note that ConcurrentHashMultiset automatically removes removes entries for keys whose
-          // count is 0.
-          if (waiters.remove(key, 1) == 1) {
-            // 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 {
-        lock.unlock();
-      }
-    }
-  }
-
-  @Override
-  public AutoUnlocker lock(K 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. If we already have access to 'key', this simply bumps up
-      // the ref count.
-      waiters.add(key);
-    } finally {
-      waitersLock.unlock();
-    }
-    ReentrantLock lock;
-    lock = locks.putIfAbsent(key, newLock);
-    if (lock != null) {
-      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 new RefCountedAutoUnlocker(key, lock);
-    }
-    // We won the race, so the current lock for 'key' is the one we locked and inserted.
-    return new RefCountedAutoUnlocker(key, newLock);
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/StripedKeyedLocker.java b/src/main/java/com/google/devtools/build/lib/concurrent/StripedKeyedLocker.java
index 2947913d5e..891b8eb60f 100644
--- a/src/main/java/com/google/devtools/build/lib/concurrent/StripedKeyedLocker.java
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/StripedKeyedLocker.java
@@ -17,21 +17,21 @@ import com.google.common.util.concurrent.Striped;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
 
 /**
  * An implementation of {@link KeyedLocker} backed by a {@link Striped}.
  */
 public class StripedKeyedLocker<K> implements KeyedLocker<K> {
 
-  private final Striped<Lock> locks;
+  private final Striped<ReadWriteLock> locks;
 
   public StripedKeyedLocker(int stripes) {
-    locks = Striped.lock(stripes);
+    locks = Striped.readWriteLock(stripes);
   }
 
-  @Override
-  public AutoUnlocker lock(final K key) {
-    final Lock lock = locks.get(key);
+  private static AutoUnlocker lockAndMakeAutoUnlocker(
+      final Lock lock, final Object keyForDebugging) {
     lock.lock();
     return new AutoUnlocker() {
       private final AtomicBoolean closeCalled = new AtomicBoolean(false);
@@ -39,12 +39,24 @@ public class StripedKeyedLocker<K> implements KeyedLocker<K> {
       @Override
       public void close() {
         if (closeCalled.getAndSet(true)) {
-          String msg = String.format("For key %s, 'close' can be called at most once per "
-              + "AutoUnlocker instance", key);
+          String msg =
+              String.format(
+                  "For key %s, 'close' can be called at most once per AutoUnlocker instance",
+                  keyForDebugging);
           throw new IllegalUnlockException(msg);
         }
         lock.unlock();
       }
     };
   }
+
+  @Override
+  public AutoUnlocker writeLock(K key) {
+    return lockAndMakeAutoUnlocker(locks.get(key).writeLock(), key);
+  }
+
+  @Override
+  public AutoUnlocker readLock(K key) {
+    return lockAndMakeAutoUnlocker(locks.get(key).readLock(), key);
+  }
 }