Introduce KeyedLocker, a nice concurrency abstraction for managing lots of mutexes, and RefCountedMultisetKeyedLocker, an efficient implementation of this abstraction.

--
MOS_MIGRATED_REVID=88000985

2 files changed, 147 insertions, 0 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
new file mode 100644
index 0000000000..8228f5b236
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/KeyedLocker.java
@@ -0,0 +1,49 @@
+// 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 javax.annotation.concurrent.ThreadSafe;
+
+/** A keyed store of locks. */
+@ThreadSafe
+public interface KeyedLocker<K> {
+  /** Used to yield access to the implicit lock granted by {@link #lock}. */
+  @ThreadSafe
+  interface AutoUnlocker extends AutoCloseable {
+    /**
+     * If this was returned by {@code lock(k)}, yields exclusive access to {@code k}.
+     *
+     * <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.
+     */
+    @Override
+    void close();
+  }
+
+  /**
+   * 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:
+   *
+   * <pre>
+   * {@code
+   * try (AutoUnlocker unlocker = locker.lock(k)) {
+   *   // Your code here.
+   * }
+   * }
+   * </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
new file mode 100644
index 0000000000..4efa90f567
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/RefCountedMultisetKeyedLocker.java
@@ -0,0 +1,98 @@
+// 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.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.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 = 256;
+  // 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, RefCountedLockImpl> locks = new ConcurrentHashMap<>();
+
+  private class RefCountedLockImpl extends ReentrantLock implements AutoUnlocker {
+    private final K key;
+
+    private RefCountedLockImpl(K key) {
+      this.key = key;
+    }
+
+    @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);
+      try {
+        Lock waitersLock = waitersLocks.get(key);
+        try {
+          waitersLock.lock();
+          // Note that ConcurrentHashMultiset automatically removes removes entries for keys whose
+          // 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);
+          }
+        } finally {
+          waitersLock.unlock();
+        }
+      } finally {
+        unlock();
+      }
+    }
+  }
+
+  @Override
+  public AutoUnlocker lock(K key) {
+    RefCountedLockImpl newLock = new RefCountedLockImpl(key);
+    // 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.
+      waiters.add(key);
+    } finally {
+      waitersLock.unlock();
+    }
+    RefCountedLockImpl 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();
+      lock.lock();
+      return lock;
+    }
+    // We won the race, so the current lock for 'key' is the one we locked and inserted.
+    return newLock;
+  }
+}