Deduplicate concurrent computations of the same Merkle tree.

Currently, it's possible for concurrent actions to end up computing the same Merkle tree, even when the cache is enabled. This change makes it so that a later action waits for the completion of the computation started by an earlier action.

Progress on bazelbuild#17923.

Closes bazelbuild#17995.

PiperOrigin-RevId: 522319291
Change-Id: I68ab952ed6357027ec71a67a104f91a684a7a040

src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java

@@ -143,6 +143,8 @@
 import java.util.SortedMap;
 import java.util.TreeMap;
 import java.util.TreeSet;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Phaser;
@@ -168,7 +170,7 @@ public class RemoteExecutionService {
   @Nullable private final RemoteExecutionClient remoteExecutor;
   private final TempPathGenerator tempPathGenerator;
   @Nullable private final Path captureCorruptedOutputsDir;
-  private final Cache<Object, MerkleTree> merkleTreeCache;
+  private final Cache<Object, CompletableFuture<MerkleTree>> merkleTreeCache;
   private final Set<String> reportedErrors = new HashSet<>();
   private final Phaser backgroundTaskPhaser = new Phaser(1);
 
@@ -344,7 +346,7 @@ public boolean mayBeExecutedRemotely(Spawn spawn) {
   }
 
   @VisibleForTesting
-  Cache<Object, MerkleTree> getMerkleTreeCache() {
+  Cache<Object, CompletableFuture<MerkleTree>> getMerkleTreeCache() {
     return merkleTreeCache;
   }
 
@@ -418,12 +420,34 @@ private MerkleTree buildMerkleTreeVisitor(
       MetadataProvider metadataProvider,
       ArtifactPathResolver artifactPathResolver)
       throws IOException, ForbiddenActionInputException {
-    MerkleTree result = merkleTreeCache.getIfPresent(nodeKey);
-    if (result == null) {
-      result = uncachedBuildMerkleTreeVisitor(walker, metadataProvider, artifactPathResolver);
-      merkleTreeCache.put(nodeKey, result);
+    // Deduplicate concurrent computations for the same node. It's not possible to use
+    // MerkleTreeCache#get(key, loader) because the loading computation may cause other nodes to be
+    // recursively looked up, which is not allowed. Instead, use a future as described at
+    // https://github.com/ben-manes/caffeine/wiki/Faq#recursive-computations.
+    var freshFuture = new CompletableFuture<MerkleTree>();
+    var priorFuture = merkleTreeCache.asMap().putIfAbsent(nodeKey, freshFuture);
+    if (priorFuture == null) {
+      // No preexisting cache entry, so we must do the computation ourselves.
+      try {
+        freshFuture.complete(
+            uncachedBuildMerkleTreeVisitor(walker, metadataProvider, artifactPathResolver));
+      } catch (Exception e) {
+        freshFuture.completeExceptionally(e);
+      }
+    }
+    try {
+      return (priorFuture != null ? priorFuture : freshFuture).join();
+    } catch (CompletionException e) {
+      Throwable cause = checkNotNull(e.getCause());
+      if (cause instanceof IOException) {
+        throw (IOException) cause;
+      } else if (cause instanceof ForbiddenActionInputException) {
+        throw (ForbiddenActionInputException) cause;
+      } else {
+        checkState(cause instanceof RuntimeException);
+        throw (RuntimeException) cause;
+      }
     }
-    return result;
   }
 
   @VisibleForTesting