[6.1.0] Handle remote cache eviction when uploading inputs for remote actions. (bazelbuild#17605)

* Extract code for cleaning stale state when remote CAS evicted blobs into another class.

PiperOrigin-RevId: 512049304
Change-Id: I87e9e6bcd4d4c4d8be27be13cfb8ba70b199b6e6

* Handle remote cache eviction when uploading inputs for remote actions.

Previously, we handle remote cache eviction when downloading inputs for local actions. However, it's possible that when Bazel need to re-execute an remote action, it detected that some inputs are missing from remote CAS. In this case, Bazel will try to upload inputs by reading from local filesystem. Since the inputs were generated remotely, not downloaded and evicted remotely, the upload will fail with FileNotFoundException.

This CL changes the code to correctly handles above case by reading through ActionFS when uploading inputs and propagate CacheNotFoundException.

Related to bazelbuild#16660.

PiperOrigin-RevId: 512568547
Change-Id: I3a28cadbb6285fa3727e1603f37abf8843c093c9

* Fix tests

---------

Co-authored-by: kshyanashree <[email protected]>

Author: coeuvre

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -34,6 +34,7 @@ java_library(
             "Retrier.java",
             "AbstractActionInputPrefetcher.java",
             "ToplevelArtifactsDownloader.java",
+            "LeaseService.java",
         ],
     ),
     exports = [
@@ -46,13 +47,13 @@ java_library(
         ":ReferenceCountedChannel",
         ":Retrier",
         ":abstract_action_input_prefetcher",
+        ":lease_service",
         ":toplevel_artifacts_downloader",
         "//src/main/java/com/google/devtools/build/lib:build-request-options",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib:runtime/command_line_path_factory",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:action_input_helper",
-        "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_data",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
@@ -83,6 +84,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/remote/common",
+        "//src/main/java/com/google/devtools/build/lib/remote/common:cache_not_found_exception",
         "//src/main/java/com/google/devtools/build/lib/remote/disk",
         "//src/main/java/com/google/devtools/build/lib/remote/downloader",
         "//src/main/java/com/google/devtools/build/lib/remote/grpc",
@@ -185,6 +187,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/remote/common",
+        "//src/main/java/com/google/devtools/build/lib/remote/common:cache_not_found_exception",
         "//src/main/java/com/google/devtools/build/lib/remote/util",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
@@ -220,3 +223,15 @@ java_library(
         "//third_party:jsr305",
     ],
 )
+
+java_library(
+    name = "lease_service",
+    srcs = ["LeaseService.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/actions",
+        "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_data",
+        "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/skyframe",
+        "//third_party:jsr305",
+    ],
+)

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

@@ -34,7 +34,6 @@
 import java.io.PushbackInputStream;
 import java.util.NoSuchElementException;
 import java.util.Objects;
-import java.util.function.Supplier;
 
 /**
  * Splits a data source into one or more {@link Chunk}s of at most {@code chunkSize} bytes.
@@ -92,8 +91,7 @@ public boolean equals(Object o) {
         return false;
       }
       Chunk other = (Chunk) o;
-      return other.offset == offset
-          && other.data.equals(data);
+      return other.offset == offset && other.data.equals(data);
     }
 
     @Override
@@ -102,7 +100,12 @@ public int hashCode() {
     }
   }
 
-  private final Supplier<InputStream> dataSupplier;
+  /** A supplier that provide data as {@link InputStream}. */
+  public interface ChunkDataSupplier {
+    InputStream get() throws IOException;
+  }
+
+  private final ChunkDataSupplier dataSupplier;
   private final long size;
   private final int chunkSize;
   private final Chunk emptyChunk;
@@ -117,7 +120,7 @@ public int hashCode() {
   // lazily on the first call to next(), as opposed to opening it in the constructor or on reset().
   private boolean initialized;
 
-  Chunker(Supplier<InputStream> dataSupplier, long size, int chunkSize, boolean compressed) {
+  Chunker(ChunkDataSupplier dataSupplier, long size, int chunkSize, boolean compressed) {
     this.dataSupplier = checkNotNull(dataSupplier);
     this.size = size;
     this.chunkSize = chunkSize;
@@ -287,7 +290,7 @@ public static class Builder {
     private int chunkSize = getDefaultChunkSize();
     protected long size;
     private boolean compressed;
-    protected Supplier<InputStream> inputStream;
+    protected ChunkDataSupplier inputStream;
 
     @CanIgnoreReturnValue
     public Builder setInput(byte[] data) {
@@ -310,14 +313,7 @@ public Builder setInput(long size, InputStream in) {
     public Builder setInput(long size, Path file) {
       checkState(inputStream == null);
       this.size = size;
-      inputStream =
-          () -> {
-            try {
-              return file.getInputStream();
-            } catch (IOException e) {
-              throw new RuntimeException(e);
-            }
-          };
+      inputStream = file::getInputStream;
       return this;
     }
 
@@ -326,30 +322,16 @@ public Builder setInput(long size, ActionInput actionInput, Path execRoot) {
       checkState(inputStream == null);
       this.size = size;
       if (actionInput instanceof VirtualActionInput) {
-        inputStream =
-            () -> {
-              try {
-                return ((VirtualActionInput) actionInput).getBytes().newInput();
-              } catch (IOException e) {
-                throw new RuntimeException(e);
-              }
-            };
+        inputStream = () -> ((VirtualActionInput) actionInput).getBytes().newInput();
       } else {
-        inputStream =
-            () -> {
-              try {
-                return ActionInputHelper.toInputPath(actionInput, execRoot).getInputStream();
-              } catch (IOException e) {
-                throw new RuntimeException(e);
-              }
-            };
+        inputStream = () -> ActionInputHelper.toInputPath(actionInput, execRoot).getInputStream();
       }
       return this;
     }
 
     @CanIgnoreReturnValue
     @VisibleForTesting
-    protected final Builder setInputSupplier(Supplier<InputStream> inputStream) {
+    protected final Builder setInputSupplier(ChunkDataSupplier inputStream) {
       this.inputStream = inputStream;
       return this;
     }

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

@@ -0,0 +1,76 @@
+// Copyright 2023 The Bazel Authors. 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.remote;
+
+import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.ActionCacheUtils;
+import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ActionLookupData;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.cache.ActionCache;
+import com.google.devtools.build.skyframe.MemoizingEvaluator;
+import java.util.HashMap;
+import java.util.Set;
+import javax.annotation.Nullable;
+
+/** A lease service that manages the lease of remote blobs. */
+public class LeaseService {
+  private final MemoizingEvaluator memoizingEvaluator;
+  @Nullable private final ActionCache actionCache;
+
+  public LeaseService(MemoizingEvaluator memoizingEvaluator, @Nullable ActionCache actionCache) {
+    this.memoizingEvaluator = memoizingEvaluator;
+    this.actionCache = actionCache;
+  }
+
+  /** Clean up internal state when files are evicted from remote CAS. */
+  public void handleMissingInputs(Set<ActionInput> missingActionInputs) {
+    if (missingActionInputs.isEmpty()) {
+      return;
+    }
+
+    var actions = new HashMap<ActionLookupData, Action>();
+
+    try {
+      for (ActionInput actionInput : missingActionInputs) {
+        if (actionInput instanceof Artifact.DerivedArtifact) {
+          Artifact.DerivedArtifact output = (Artifact.DerivedArtifact) actionInput;
+          ActionLookupData actionLookupData = output.getGeneratingActionKey();
+          var actionLookupValue =
+              memoizingEvaluator.getExistingValue(actionLookupData.getActionLookupKey());
+          if (actionLookupValue instanceof ActionLookupValue) {
+            Action action =
+                ((ActionLookupValue) actionLookupValue)
+                    .getAction(actionLookupData.getActionIndex());
+            actions.put(actionLookupData, action);
+          }
+        }
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+
+    if (!actions.isEmpty()) {
+      var actionKeys = actions.keySet();
+      memoizingEvaluator.delete(key -> key instanceof ActionLookupData && actionKeys.contains(key));
+
+      if (actionCache != null) {
+        for (var action : actions.values()) {
+          ActionCacheUtils.removeCacheEntry(actionCache, action);
+        }
+      }
+    }
+  }
+}

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

@@ -64,6 +64,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionRequirements;
@@ -371,10 +372,14 @@ private MerkleTree buildInputMerkleTree(
           spawn,
           context,
           (Object nodeKey, InputWalker walker) -> {
-            subMerkleTrees.add(buildMerkleTreeVisitor(nodeKey, walker, metadataProvider));
+            subMerkleTrees.add(
+                buildMerkleTreeVisitor(
+                    nodeKey, walker, metadataProvider, context.getPathResolver()));
           });
       if (!outputDirMap.isEmpty()) {
-        subMerkleTrees.add(MerkleTree.build(outputDirMap, metadataProvider, execRoot, digestUtil));
+        subMerkleTrees.add(
+            MerkleTree.build(
+                outputDirMap, metadataProvider, execRoot, context.getPathResolver(), digestUtil));
       }
       return MerkleTree.merge(subMerkleTrees, digestUtil);
     } else {
@@ -394,31 +399,44 @@ private MerkleTree buildInputMerkleTree(
           toolSignature == null ? ImmutableSet.of() : toolSignature.toolInputs,
           context.getMetadataProvider(),
           execRoot,
+          context.getPathResolver(),
           digestUtil);
     }
   }
 
   private MerkleTree buildMerkleTreeVisitor(
-      Object nodeKey, InputWalker walker, MetadataProvider metadataProvider)
+      Object nodeKey,
+      InputWalker walker,
+      MetadataProvider metadataProvider,
+      ArtifactPathResolver artifactPathResolver)
       throws IOException, ForbiddenActionInputException {
     MerkleTree result = merkleTreeCache.getIfPresent(nodeKey);
     if (result == null) {
-      result = uncachedBuildMerkleTreeVisitor(walker, metadataProvider);
+      result = uncachedBuildMerkleTreeVisitor(walker, metadataProvider, artifactPathResolver);
       merkleTreeCache.put(nodeKey, result);
     }
     return result;
   }
 
   @VisibleForTesting
   public MerkleTree uncachedBuildMerkleTreeVisitor(
-      InputWalker walker, MetadataProvider metadataProvider)
+      InputWalker walker,
+      MetadataProvider metadataProvider,
+      ArtifactPathResolver artifactPathResolver)
       throws IOException, ForbiddenActionInputException {
     ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue<>();
     subMerkleTrees.add(
-        MerkleTree.build(walker.getLeavesInputMapping(), metadataProvider, execRoot, digestUtil));
+        MerkleTree.build(
+            walker.getLeavesInputMapping(),
+            metadataProvider,
+            execRoot,
+            artifactPathResolver,
+            digestUtil));
     walker.visitNonLeaves(
         (Object subNodeKey, InputWalker subWalker) -> {
-          subMerkleTrees.add(buildMerkleTreeVisitor(subNodeKey, subWalker, metadataProvider));
+          subMerkleTrees.add(
+              buildMerkleTreeVisitor(
+                  subNodeKey, subWalker, metadataProvider, artifactPathResolver));
         });
     return MerkleTree.merge(subMerkleTrees, digestUtil);
   }

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

@@ -957,9 +957,13 @@ public ActionInput getActionInput(Path path) {
               });
       env.getEventBus().register(toplevelArtifactsDownloader);
 
+      var leaseService =
+          new LeaseService(
+              env.getSkyframeExecutor().getEvaluator(),
+              env.getBlazeWorkspace().getPersistentActionCache());
+
       remoteOutputService.setActionInputFetcher(actionInputFetcher);
-      remoteOutputService.setMemoizingEvaluator(env.getSkyframeExecutor().getEvaluator());
-      remoteOutputService.setActionCache(env.getBlazeWorkspace().getPersistentActionCache());
+      remoteOutputService.setLeaseService(leaseService);
       env.getEventBus().register(remoteOutputService);
     }
   }