2 files changed, 38 insertions, 10 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 dae49c9e46..76debaa8a1 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
@@ -60,6 +60,8 @@ public abstract class MultisetSemaphore<T> {
    */
   public abstract void releaseAll(Set<T> valuesToRelease);
 
+  public abstract int estimateCurrentNumUniqueValues();
+
   /**
    * Returns a {@link MultisetSemaphore} with a backing {@link Semaphore} that has an unbounded
    * number of permits; that is, {@link #acquireAll} will never block.
@@ -122,6 +124,12 @@ public abstract class MultisetSemaphore<T> {
     @Override
     public void releaseAll(Set<T> valuesToRelease) {
     }
+
+    @Override
+    public int estimateCurrentNumUniqueValues() {
+      // We can't give a good estimate since we don't track values at all.
+      return 0;
+    }
   }
 
   private static class NaiveMultisetSemaphore<T> extends MultisetSemaphore<T> {
@@ -162,5 +170,14 @@ public abstract class MultisetSemaphore<T> {
       }
       semaphore.release(numUniqueValuesToRelease);
     }
+
+    @Override
+    public int estimateCurrentNumUniqueValues() {
+      // Notes:
+      //   (1) The race here is completely benign; we're just supposed to return an estimate.
+      //   (2) See the javadoc for Multiset#size, which explains to use entrySet().size() to get the
+      //       number of unique values.
+      return actualValues.entrySet().size();
+    }
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/concurrent/MultisetSemaphoreTest.java b/src/test/java/com/google/devtools/build/lib/concurrent/MultisetSemaphoreTest.java
index 6e75c0d26b..8d659fd0f5 100644
--- a/src/test/java/com/google/devtools/build/lib/concurrent/MultisetSemaphoreTest.java
+++ b/src/test/java/com/google/devtools/build/lib/concurrent/MultisetSemaphoreTest.java
@@ -45,17 +45,27 @@ public class MultisetSemaphoreTest {
         .maxNumUniqueValues(3)
         .build();
 
-    // And we serially acquire permits for 3 unique values
+    // And we serially acquire permits for 3 unique values,
     multisetSemaphore.acquireAll(ImmutableSet.of("a", "b", "c"));
+    // Then the MultisetSemaphore thinks it currently has 3 unique values.
+    assertThat(multisetSemaphore.estimateCurrentNumUniqueValues()).isEqualTo(3);
     // And then attempt to acquire permits for 2 of those same unique values,
-    // Then we don't deadlock.
+    // Then we don't deadlock,
     multisetSemaphore.acquireAll(ImmutableSet.of("b", "c"));
+    // And the MultisetSemaphore still thinks it currently has 3 unique values.
+    assertThat(multisetSemaphore.estimateCurrentNumUniqueValues()).isEqualTo(3);
     // And then we release one of the permit for one of those unique values,
     multisetSemaphore.releaseAll(ImmutableSet.of("c"));
-    // And then we release the other permit,
+    // Then the MultisetSemaphore still thinks it currently has 3 unique values.
+    assertThat(multisetSemaphore.estimateCurrentNumUniqueValues()).isEqualTo(3);
+    // And then we release the final permit for that unique value,
     multisetSemaphore.releaseAll(ImmutableSet.of("c"));
-    // We are able to acquire a permit for a 4th unique value.
+    // Then the MultisetSemaphore thinks it currently has 2 unique values.
+    assertThat(multisetSemaphore.estimateCurrentNumUniqueValues()).isEqualTo(2);
+    // And we are able to acquire a permit for a 4th unique value,
     multisetSemaphore.acquireAll(ImmutableSet.of("d"));
+    // And the MultisetSemaphore thinks it currently has 3 unique values.
+    assertThat(multisetSemaphore.estimateCurrentNumUniqueValues()).isEqualTo(3);
   }
 
   @Test
@@ -171,12 +181,13 @@ public class MultisetSemaphoreTest {
                     public void run() {
                       try {
                         Set<String> vals = ImmutableSet.of(sameVal, differentVal);
-                        // Tries to acquire a permit for a set of two values, one of which is the same for all
-                        // the N Runnables and one of which is unique across all N Runnables.
+                        // Tries to acquire a permit for a set of two values, one of which is the
+                        // same for all the N Runnables and one of which is unique across all N
+                        // Runnables,
                         multisetSemaphore.acquireAll(vals);
-                        // And then sleeps
+                        // And then sleeps,
                         Thread.sleep(napTimeMs);
-                        // And then releases its permits
+                        // And then releases its permits,
                         multisetSemaphore.releaseAll(vals);
                         // And then counts down the done latch,
                         allDoneLatch.countDown();
@@ -230,8 +241,8 @@ public class MultisetSemaphoreTest {
                     @Override
                     public void run() {
                       try {
-                        // Tries to acquire a permit for the set of N values, with a unique iteration order
-                        // (across all the N! different permutations)
+                        // Tries to acquire a permit for the set of N values, with a unique
+                        // iteration order (across all the N! different permutations),
                         multisetSemaphore.acquireAll(orderedSet);
                         // And then immediately releases the permit.
                         multisetSemaphore.releaseAll(orderedSet);