Introduce KeyedLocker#lockBatch, which does what it sounds like it does.

--
MOS_MIGRATED_REVID=90282858

2 files changed, 127 insertions, 11 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/BatchedKeyedLocker.java b/src/main/java/com/google/devtools/build/lib/concurrent/BatchedKeyedLocker.java
new file mode 100644
index 0000000000..d65b3c165d
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/BatchedKeyedLocker.java
@@ -0,0 +1,46 @@
+// Copyright 2015 Google Inc. 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.devtools.build.lib.concurrent.KeyedLocker.AutoUnlocker;
+
+import java.util.Comparator;
+import java.util.Set;
+
+/** A {@link KeyedLocker} that additionally offers batched locking functionality. */
+public interface BatchedKeyedLocker<K> extends KeyedLocker<K> {
+  /** Factory for {@link BatchedKeyedLocker} instances. */
+  interface Factory<K> {
+    /**
+     * Returns a fresh {@link BatchedKeyedLocker} instance.
+     *
+     * <p>The given {@link Comparator} instance is used to get consistent ordering for
+     * {@link BatchedKeyedLocker#lockBatch}.
+     */
+    BatchedKeyedLocker<K> create(Comparator<K> comparator);
+  }
+
+  /**
+   * Similar to {@link #lock}, blocks the current thread until it has exclusive access to do
+   * things with all the keys in {@code keys} and returns a single {@link AutoUnlocker} instance
+   * for yielding the implicit locks on all the given keys.
+   *
+   * <p>If a thread has an unclosed {@link AutoUnlocker} instance returned by a call to
+   * {@code lockBatch(keys)}, this is equivalent to having separate, unclosed {@link AutoUnlocker}
+   * instances for each {@code k} in {@code keys}.
+   *
+   * <p>Note that use of this method avoids the concerns described in {@link #lock}.
+   */
+  AutoUnlocker lockBatch(Set<K> keys);
+}
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 0656285fae..5cab307e4e 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
@@ -15,21 +15,26 @@ package com.google.devtools.build.lib.concurrent;
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ConcurrentHashMultiset;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.util.concurrent.Striped;
-import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Set;
 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;
 
+import javax.annotation.Nullable;
+
 /**
  * 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> {
+public class RefCountedMultisetKeyedLocker<K> implements BatchedKeyedLocker<K> {
   // Multiset of keys that have threads waiting on a lock or using a lock.
   private final ConcurrentHashMultiset<K> waiters = ConcurrentHashMultiset.<K>create();
 
@@ -41,10 +46,44 @@ public class RefCountedMultisetKeyedLocker<K> implements KeyedLocker<K> {
   // 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<>();
 
-  private class RefCountedAutoUnlocker implements AutoUnlocker {
+  // Used to enforce a consistent ordering in lockBatch.
+  @Nullable
+  private final Comparator<K> comparator;
+
+  private RefCountedMultisetKeyedLocker(Comparator<K> comparator) {
+    this.comparator = comparator;
+  }
+
+  /** Factory for {@link RefCountedMultisetKeyedLocker} instances. */ 
+  public static class Factory<K> implements BatchedKeyedLocker.Factory<K> {
+    @Override
+    public BatchedKeyedLocker<K> create(Comparator<K> comparator) {
+      return new RefCountedMultisetKeyedLocker<>(comparator);
+    }
+
+    public KeyedLocker<K> create() {
+      return new RefCountedMultisetKeyedLocker<>(/*comparator=*/null);
+    }
+  }
+
+  private abstract static class AtMostOnceAutoUnlockerBase<K> implements AutoUnlocker {
+    private final AtomicBoolean closeCalled = new AtomicBoolean(false);
+
+    @Override
+    public final void close() {
+      if (closeCalled.getAndSet(true)) {
+        String msg = "'close' can be called at most once per AutoUnlocker instance";
+        throw new IllegalUnlockException(msg);
+      }
+      doClose();
+    }
+
+    protected abstract void doClose();
+  }
+
+  private class RefCountedAutoUnlocker extends AtMostOnceAutoUnlockerBase<K> {
     private final K key;
     private final ReentrantLock lock;
-    private final AtomicBoolean closeCalled = new AtomicBoolean(false);
 
     private RefCountedAutoUnlocker(K key, ReentrantLock lock) {
       this.key = key;
@@ -52,12 +91,7 @@ public class RefCountedMultisetKeyedLocker<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);
-        throw new IllegalUnlockException(msg);
-      }
+    protected void doClose() {
       if (!lock.isHeldByCurrentThread()) {
         String msg = String.format("For key %s, the calling thread to 'close' must be the one "
             + "that acquired the AutoUnlocker", key);
@@ -113,4 +147,40 @@ public class RefCountedMultisetKeyedLocker<K> implements KeyedLocker<K> {
     // We won the race, so the current lock for 'key' is the one we locked and inserted.
     return new RefCountedAutoUnlocker(key, newLock);
   }
+
+  private static void unlockAll(Iterable<KeyedLocker.AutoUnlocker> unlockers) {
+    // Note that order doesn't matter here because we always acquire locks in a consistent order.
+    for (KeyedLocker.AutoUnlocker unlocker : unlockers) {
+      unlocker.close();
+    }
+  }
+
+  @Override
+  public AutoUnlocker lockBatch(Set<K> keys) {
+    // This indicates the client did some unsafe casting - not our problem.
+    Preconditions.checkNotNull(comparator);
+    // We acquire locks in a consistent order. This prevents a deadlock that would otherwise occur
+    // on two concurrent calls to lockBatch(keys(k1, k2)) if the callers acquired the locks in a
+    // different order.
+    ImmutableSortedSet<K> sortedKeys = ImmutableSortedSet.copyOf(comparator, keys);
+    final List<KeyedLocker.AutoUnlocker> unlockers = new ArrayList<>(sortedKeys.size());
+    boolean success = false;
+    try {
+      for (K key : sortedKeys) {
+        unlockers.add(lock(key));
+      }
+      success = true;
+      return new AtMostOnceAutoUnlockerBase<K>() {
+        @Override
+        public void doClose() {
+          unlockAll(unlockers);
+        }
+      };
+    } finally {
+      // Just in case we encounter a crash, e.g. if there is a bug in #lock.
+      if (!success) {
+        unlockAll(unlockers);
+      }
+    }
+  }
 }