From 6d6fa81e47591f41b9dce5e2a9a8512ec13d3f4a Mon Sep 17 00:00:00 2001
From: Tiago Quelhas <tjgq@google.com>
Date: Thu, 6 Apr 2023 06:17:12 -0700
Subject: [PATCH] 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 #17923.

Closes #17995.

PiperOrigin-RevId: 522319291
Change-Id: I68ab952ed6357027ec71a67a104f91a684a7a040
---
 .../lib/remote/RemoteExecutionService.java    | 38 +++++++++++++++----
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
index f38c1e3dd1f620..d3b40f5cf00ac0 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
+++ b/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