Cleanup stale state when remote cache evicted

Currently, when building without the bytes, if Bazel failed to download blobs from CAS when fetching them as inputs to local actions, Bazel fails the build with message like `... --remote_download_outputs=minimal does not work if your remote cache evicts files during builds.` and this message keep showing up until a manually `bazel clean`.

This PR fixes that by cleaning up stale state in skyframe and action cache upon remote cache eviction so that a following build can continue without `bazel shutdown` or `bazel clean`.

Fixes bazelbuild#17366.
Part of bazelbuild#16660.

Closes bazelbuild#17462.

PiperOrigin-RevId: 510952745
Change-Id: I4fc59a21195565c68375a19ead76738d2208c4ac

Author: coeuvre

src/main/java/com/google/devtools/build/lib/actions/ActionCacheChecker.java

@@ -145,19 +145,11 @@ private ActionCache.Entry getCacheEntry(Action action) {
     if (!cacheConfig.enabled()) {
       return null; // ignore existing cache when disabled.
     }
-    for (Artifact output : action.getOutputs()) {
-      ActionCache.Entry entry = actionCache.get(output.getExecPathString());
-      if (entry != null) {
-        return entry;
-      }
-    }
-    return null;
+    return ActionCacheUtils.getCacheEntry(actionCache, action);
   }
 
   private void removeCacheEntry(Action action) {
-    for (Artifact output : action.getOutputs()) {
-      actionCache.remove(output.getExecPathString());
-    }
+    ActionCacheUtils.removeCacheEntry(actionCache, action);
   }
 
   @Nullable

src/main/java/com/google/devtools/build/lib/actions/ActionCacheUtils.java

@@ -0,0 +1,40 @@
+// 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.actions;
+
+import com.google.devtools.build.lib.actions.cache.ActionCache;
+import javax.annotation.Nullable;
+
+/** Utility functions for {@link ActionCache}. */
+public class ActionCacheUtils {
+  private ActionCacheUtils() {}
+
+  /** Checks whether one of existing output paths is already used as a key. */
+  @Nullable
+  public static ActionCache.Entry getCacheEntry(ActionCache actionCache, Action action) {
+    for (Artifact output : action.getOutputs()) {
+      ActionCache.Entry entry = actionCache.get(output.getExecPathString());
+      if (entry != null) {
+        return entry;
+      }
+    }
+    return null;
+  }
+
+  public static void removeCacheEntry(ActionCache actionCache, Action action) {
+    for (Artifact output : action.getOutputs()) {
+      actionCache.remove(output.getExecPathString());
+    }
+  }
+}

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

@@ -42,6 +42,7 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
 import com.google.devtools.build.lib.events.Reporter;
+import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
 import com.google.devtools.build.lib.remote.util.AsyncTaskCache;
 import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult;
 import com.google.devtools.build.lib.remote.util.TempPathGenerator;
@@ -61,6 +62,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 
 /**
  * Abstract implementation of {@link ActionInputPrefetcher} which implements the orchestration of
@@ -78,6 +80,8 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet
   protected final Path execRoot;
   protected final ImmutableList<Pattern> patternsToDownload;
 
+  private final Set<ActionInput> missingActionInputs = Sets.newConcurrentHashSet();
+
   private static class Context {
     private final Set<Path> nonWritableDirs = Sets.newConcurrentHashSet();
 
@@ -399,7 +403,8 @@ private Completable prefetchInputFileOrSymlink(
     PathFragment prefetchExecPath = metadata.getMaterializationExecPath().orElse(execPath);
 
     Completable prefetch =
-        downloadFileNoCheckRx(context, execRoot.getRelative(prefetchExecPath), metadata, priority);
+        downloadFileNoCheckRx(
+            context, execRoot.getRelative(prefetchExecPath), input, metadata, priority);
 
     // If prefetching to a different path, plant a symlink into it.
     if (!prefetchExecPath.equals(execPath)) {
@@ -418,15 +423,23 @@ private Completable prefetchInputFileOrSymlink(
    * download finished.
    */
   private Completable downloadFileRx(
-      Context context, Path path, FileArtifactValue metadata, Priority priority) {
+      Context context,
+      Path path,
+      @Nullable ActionInput actionInput,
+      FileArtifactValue metadata,
+      Priority priority) {
     if (!canDownloadFile(path, metadata)) {
       return Completable.complete();
     }
-    return downloadFileNoCheckRx(context, path, metadata, priority);
+    return downloadFileNoCheckRx(context, path, actionInput, metadata, priority);
   }
 
   private Completable downloadFileNoCheckRx(
-      Context context, Path path, FileArtifactValue metadata, Priority priority) {
+      Context context,
+      Path path,
+      @Nullable ActionInput actionInput,
+      FileArtifactValue metadata,
+      Priority priority) {
     if (path.isSymbolicLink()) {
       try {
         path = path.getRelative(path.readSymbolicLink());
@@ -440,26 +453,32 @@ private Completable downloadFileNoCheckRx(
     AtomicBoolean completed = new AtomicBoolean(false);
     Completable download =
         Completable.using(
-            tempPathGenerator::generateTempPath,
-            tempPath ->
-                toCompletable(
-                        () ->
-                            doDownloadFile(
-                                tempPath, finalPath.relativeTo(execRoot), metadata, priority),
-                        directExecutor())
-                    .doOnComplete(
-                        () -> {
-                          finalizeDownload(context, tempPath, finalPath);
-                          completed.set(true);
-                        }),
-            tempPath -> {
-              if (!completed.get()) {
-                deletePartialDownload(tempPath);
-              }
-            },
-            // Set eager=false here because we want cleanup the download *after* upstream is
-            // disposed.
-            /* eager= */ false);
+                tempPathGenerator::generateTempPath,
+                tempPath ->
+                    toCompletable(
+                            () ->
+                                doDownloadFile(
+                                    tempPath, finalPath.relativeTo(execRoot), metadata, priority),
+                            directExecutor())
+                        .doOnComplete(
+                            () -> {
+                              finalizeDownload(context, tempPath, finalPath);
+                              completed.set(true);
+                            }),
+                tempPath -> {
+                  if (!completed.get()) {
+                    deletePartialDownload(tempPath);
+                  }
+                },
+                // Set eager=false here because we want cleanup the download *after* upstream is
+                // disposed.
+                /* eager= */ false)
+            .doOnError(
+                error -> {
+                  if (error instanceof CacheNotFoundException && actionInput != null) {
+                    missingActionInputs.add(actionInput);
+                  }
+                });
 
     return downloadCache.executeIfNot(
         finalPath,
@@ -479,19 +498,27 @@ private Completable downloadFileNoCheckRx(
    * <p>The file will be written into a temporary file and moved to the final destination after the
    * download finished.
    */
-  public void downloadFile(Path path, FileArtifactValue metadata)
+  public void downloadFile(Path path, @Nullable ActionInput actionInput, FileArtifactValue metadata)
       throws IOException, InterruptedException {
-    getFromFuture(downloadFileAsync(path.asFragment(), metadata, Priority.CRITICAL));
+    getFromFuture(downloadFileAsync(path.asFragment(), actionInput, metadata, Priority.CRITICAL));
   }
 
   protected ListenableFuture<Void> downloadFileAsync(
-      PathFragment path, FileArtifactValue metadata, Priority priority) {
+      PathFragment path,
+      @Nullable ActionInput actionInput,
+      FileArtifactValue metadata,
+      Priority priority) {
     Context context = new Context();
     return toListenableFuture(
         Completable.using(
             () -> context,
             ctx ->
-                downloadFileRx(context, execRoot.getFileSystem().getPath(path), metadata, priority),
+                downloadFileRx(
+                    context,
+                    execRoot.getFileSystem().getPath(path),
+                    actionInput,
+                    metadata,
+                    priority),
             Context::finalizeContext));
   }
 
@@ -648,4 +675,8 @@ private boolean outputMatchesPattern(Artifact output) {
   public void flushOutputTree() throws InterruptedException {
     downloadCache.awaitInProgressTasks();
   }
+
+  public ImmutableSet<ActionInput> getMissingActionInputs() {
+    return ImmutableSet.copyOf(missingActionInputs);
+  }
 }

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

@@ -53,6 +53,7 @@ java_library(
         "//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",
@@ -185,11 +186,13 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//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/util",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:flogger",
         "//third_party:guava",
+        "//third_party:jsr305",
         "//third_party:rxjava3",
     ],
 )

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

@@ -24,6 +24,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputMap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
@@ -533,6 +534,12 @@ public RemoteFileArtifactValue getRemoteMetadata() {
     };
   }
 
+  @Nullable
+  protected ActionInput getActionInput(PathFragment path) {
+    PathFragment execPath = path.relativeTo(execRoot);
+    return inputArtifactData.getInput(execPath.getPathString());
+  }
+
   @Nullable
   protected RemoteFileArtifactValue getRemoteMetadata(PathFragment path) {
     if (!isOutput(path)) {
@@ -572,7 +579,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
     FileArtifactValue m = getRemoteMetadata(path);
     if (m != null) {
       try {
-        inputFetcher.downloadFile(delegateFs.getPath(path), m);
+        inputFetcher.downloadFile(delegateFs.getPath(path), getActionInput(path), m);
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
         throw new IOException(

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

@@ -151,15 +151,21 @@ public ListenableFuture<Void> uploadActionResult(
    */
   public ListenableFuture<Void> uploadFile(
       RemoteActionExecutionContext context, Digest digest, Path file) {
+    return uploadFile(context, digest, file, /* force= */ false);
+  }
+
+  protected ListenableFuture<Void> uploadFile(
+      RemoteActionExecutionContext context, Digest digest, Path file, boolean force) {
     if (digest.getSizeBytes() == 0) {
       return COMPLETED_SUCCESS;
     }
 
     Completable upload =
-        casUploadCache.executeIfNot(
+        casUploadCache.execute(
             digest,
             RxFutures.toCompletable(
-                () -> cacheProtocol.uploadFile(context, digest, file), directExecutor()));
+                () -> cacheProtocol.uploadFile(context, digest, file), directExecutor()),
+            force);
 
     return RxFutures.toListenableFuture(upload);
   }
@@ -176,15 +182,21 @@ public ListenableFuture<Void> uploadFile(
    */
   public ListenableFuture<Void> uploadBlob(
       RemoteActionExecutionContext context, Digest digest, ByteString data) {
+    return uploadBlob(context, digest, data, /* force= */ false);
+  }
+
+  protected ListenableFuture<Void> uploadBlob(
+      RemoteActionExecutionContext context, Digest digest, ByteString data, boolean force) {
     if (digest.getSizeBytes() == 0) {
       return COMPLETED_SUCCESS;
     }
 
     Completable upload =
-        casUploadCache.executeIfNot(
+        casUploadCache.execute(
             digest,
             RxFutures.toCompletable(
-                () -> cacheProtocol.uploadBlob(context, digest, data), directExecutor()));
+                () -> cacheProtocol.uploadBlob(context, digest, data), directExecutor()),
+            force);
 
     return RxFutures.toListenableFuture(upload);
   }

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

@@ -34,7 +34,9 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
+import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.analysis.AnalysisResult;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
@@ -59,6 +61,7 @@
 import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
 import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
 import com.google.devtools.build.lib.remote.RemoteServerCapabilities.ServerCapabilitiesRequirement;
+import com.google.devtools.build.lib.remote.ToplevelArtifactsDownloader.PathToMetadataConverter;
 import com.google.devtools.build.lib.remote.common.RemoteCacheClient;
 import com.google.devtools.build.lib.remote.common.RemoteExecutionClient;
 import com.google.devtools.build.lib.remote.downloader.GrpcRemoteDownloader;
@@ -987,7 +990,6 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
               remoteOptions.useNewExitCodeForLostInputs);
       env.getEventBus().register(actionInputFetcher);
       builder.setActionInputPrefetcher(actionInputFetcher);
-      remoteOutputService.setActionInputFetcher(actionInputFetcher);
       actionContextProvider.setActionInputFetcher(actionInputFetcher);
 
       toplevelArtifactsDownloader =
@@ -996,15 +998,35 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
               remoteOutputsMode.downloadToplevelOutputsOnly(),
               env.getSkyframeExecutor().getEvaluator(),
               actionInputFetcher,
-              (path) -> {
-                FileSystem fileSystem = path.getFileSystem();
-                if (fileSystem instanceof RemoteActionFileSystem) {
-                  return ((RemoteActionFileSystem) path.getFileSystem())
-                      .getRemoteMetadata(path.asFragment());
+              new PathToMetadataConverter() {
+                @Nullable
+                @Override
+                public FileArtifactValue getMetadata(Path path) {
+                  FileSystem fileSystem = path.getFileSystem();
+                  if (fileSystem instanceof RemoteActionFileSystem) {
+                    return ((RemoteActionFileSystem) path.getFileSystem())
+                        .getRemoteMetadata(path.asFragment());
+                  }
+                  return null;
+                }
+
+                @Nullable
+                @Override
+                public ActionInput getActionInput(Path path) {
+                  FileSystem fileSystem = path.getFileSystem();
+                  if (fileSystem instanceof RemoteActionFileSystem) {
+                    return ((RemoteActionFileSystem) path.getFileSystem())
+                        .getActionInput(path.asFragment());
+                  }
+                  return null;
                 }
-                return null;
               });
       env.getEventBus().register(toplevelArtifactsDownloader);
+
+      remoteOutputService.setActionInputFetcher(actionInputFetcher);
+      remoteOutputService.setMemoizingEvaluator(env.getSkyframeExecutor().getEvaluator());
+      remoteOutputService.setActionCache(env.getBlazeWorkspace().getPersistentActionCache());
+      env.getEventBus().register(remoteOutputService);
     }
   }