[6.1.0] Handle remote cache eviction when uploading inputs for remote actions.

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);
     }
   }