Fix MultisetSemaphore.

We use a fixed version of the previous algorithm. See the comments for details.

Fancier algorithms exist. I thought of a cool one that makes use of BatchKeyedLocker (would give me an excuse to revive it, heh), but fancy algorithms would be overkill. As noted in the initial commit of NaiveMultisetSemaphore, performance isn't critical.

RELNOTES: None
PiperOrigin-RevId: 208560559

1 files changed, 34 insertions, 14 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/MultisetSemaphore.java b/src/main/java/com/google/devtools/build/lib/concurrent/MultisetSemaphore.java
index 8a0bb93bde..94611b5ef3 100644
--- a/src/main/java/com/google/devtools/build/lib/concurrent/MultisetSemaphore.java
+++ b/src/main/java/com/google/devtools/build/lib/concurrent/MultisetSemaphore.java
@@ -146,30 +146,50 @@ public abstract class MultisetSemaphore<T> {
 
     @Override
     public void acquireAll(Set<T> valuesToAcquire) throws InterruptedException {
-      int numUniqueValuesToAcquire = 0;
+      int oldNumNeededPermits;
       synchronized (lock) {
-        for (T value : valuesToAcquire) {
-          int oldCount = actualValues.add(value, 1);
-          if (oldCount == 0) {
-            numUniqueValuesToAcquire++;
+        oldNumNeededPermits = computeNumNeededPermitsLocked(valuesToAcquire);
+      }
+      while (true) {
+        semaphore.acquire(oldNumNeededPermits);
+        synchronized (lock) {
+          int newNumNeededPermits = computeNumNeededPermitsLocked(valuesToAcquire);
+          if (newNumNeededPermits != oldNumNeededPermits) {
+            // While we were doing 'acquire' above, another thread won the race to acquire the first
+            // usage of one of the values in 'valuesToAcquire' or release the last usage of one of
+            // the values. This means we either acquired too many or too few permits, respectively,
+            // above. Release the permits we did acquire, in order to restore the accuracy of the
+            // semaphore's current count, and then try again.
+            semaphore.release(oldNumNeededPermits);
+            oldNumNeededPermits = newNumNeededPermits;
+            continue;
+          } else {
+            // Our modification to the semaphore was correct, so it's sound to update the multiset.
+            valuesToAcquire.forEach(actualValues::add);
+            return;
           }
         }
       }
-      semaphore.acquire(numUniqueValuesToAcquire);
+    }
+
+    private int computeNumNeededPermitsLocked(Set<T> valuesToAcquire) {
+      // We need a permit for each value that is not already in the multiset.
+      return (int) valuesToAcquire.stream()
+          .filter(v -> actualValues.count(v) == 0)
+          .count();
     }
 
     @Override
     public void releaseAll(Set<T> valuesToRelease) {
-      int numUniqueValuesToRelease = 0;
       synchronized (lock) {
-        for (T value : valuesToRelease) {
-          int oldCount = actualValues.remove(value, 1);
-          if (oldCount == 1) {
-            numUniqueValuesToRelease++;
-          }
-        }
+        // We need to release a permit for each value that currently has multiplicity 1.
+        int numPermitsToRelease =
+            valuesToRelease
+                .stream()
+                .mapToInt(v -> actualValues.remove(v, 1) == 1 ? 1 : 0)
+                .sum();
+        semaphore.release(numPermitsToRelease);
       }
-      semaphore.release(numUniqueValuesToRelease);
     }
 
     @Override