Add Bazel support for --rewind_lost_inputs

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

@@ -80,7 +80,7 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet
   protected final Path execRoot;
   protected final RemoteOutputChecker remoteOutputChecker;
 
-  @Nullable private final ActionOutputDirectoryHelper outputDirectoryHelper;
+  @Nullable protected final ActionOutputDirectoryHelper outputDirectoryHelper;
 
   /** The state of a directory tracked by {@link DirectoryTracker}, as explained below. */
   enum DirectoryState {
@@ -142,8 +142,9 @@ private void setWritable(Path dir, DirectoryState newState) throws IOException {
       directoryStateMap.compute(
           dir,
           (unusedKey, oldState) -> {
-            if (oldState == DirectoryState.TEMPORARILY_WRITABLE
-                || oldState == DirectoryState.PERMANENTLY_WRITABLE) {
+            if (!forceRefetch(dir)
+                && (oldState == DirectoryState.TEMPORARILY_WRITABLE
+                    || oldState == DirectoryState.PERMANENTLY_WRITABLE)) {
               // Already writable, but must potentially upgrade from temporary to permanent.
               return newState == DirectoryState.PERMANENTLY_WRITABLE ? newState : oldState;
             }
@@ -179,8 +180,9 @@ void setOutputPermissions(Path dir) throws IOException {
       directoryStateMap.compute(
           dir,
           (unusedKey, oldState) -> {
-            if (oldState == DirectoryState.OUTPUT_PERMISSIONS
-                || oldState == DirectoryState.PERMANENTLY_WRITABLE) {
+            if (!forceRefetch(dir)
+                && (oldState == DirectoryState.OUTPUT_PERMISSIONS
+                    || oldState == DirectoryState.PERMANENTLY_WRITABLE)) {
               // Either the output permissions have already been set, or we're not changing the
               // permissions ever again.
               return oldState;
@@ -258,6 +260,12 @@ private static boolean shouldDownloadFile(Path path, FileArtifactValue metadata)
 
   protected abstract boolean canDownloadFile(Path path, FileArtifactValue metadata);
 
+  /**
+   * If true, then all previously acquired knowledge of the file system state of this path (e.g. the
+   * existence of tree artifact directories or previously downloaded files) must be discarded.
+   */
+  protected abstract boolean forceRefetch(Path path);
+
   /**
    * Downloads file to the given path via its metadata.
    *
@@ -608,15 +616,16 @@ private Completable downloadFileNoCheckRx(
                           alreadyDeleted.set(true);
                         }));
 
-    return downloadCache.executeIfNot(
+    return downloadCache.execute(
         finalPath,
         Completable.defer(
             () -> {
               if (shouldDownloadFile(finalPath, metadata)) {
                 return download;
               }
               return Completable.complete();
-            }));
+            }),
+        forceRefetch(finalPath));
   }
 
   private void finalizeDownload(
@@ -693,7 +702,7 @@ private static void deletePartialDownload(Path path) {
   }
 
   private Completable plantSymlink(Symlink symlink) {
-    return downloadCache.executeIfNot(
+    return downloadCache.execute(
         symlink.linkPath(),
         Completable.defer(
             () -> {
@@ -702,7 +711,8 @@ private Completable plantSymlink(Symlink symlink) {
               symlink.linkPath().delete();
               symlink.linkPath().createSymbolicLink(symlink.targetPath());
               return Completable.complete();
-            }));
+            }),
+        forceRefetch(symlink.linkPath));
   }
 
   public ImmutableSet<Path> downloadedFiles() {
@@ -739,7 +749,7 @@ public void finalizeAction(Action action, OutputMetadataStore outputMetadataStor
       }
 
       var metadata = outputMetadataStore.getOutputMetadata(output);
-      if (!metadata.isRemote()) {
+      if (!canDownloadFile(output.getPath(), metadata)) {
         continue;
       }
 

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

@@ -257,6 +257,7 @@ java_library(
     deps = [
         ":abstract_action_input_prefetcher",
         ":combined_cache",
+        ":concurrent_artifact_path_trie",
         ":remote_output_checker",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:action_output_directory_helper",
@@ -566,6 +567,7 @@ java_library(
         ":remote_action_file_system",
         ":remote_action_input_fetcher",
         ":remote_output_checker",
+        ":remote_rewound_action_synchronizer",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
@@ -795,6 +797,7 @@ java_library(
     name = "remote_output_checker",
     srcs = ["RemoteOutputChecker.java"],
     deps = [
+        ":concurrent_artifact_path_trie",
         "//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/analysis:analysis_cluster",
@@ -871,18 +874,48 @@ java_library(
     srcs = ["Store.java"],
 )
 
+java_library(
+    name = "concurrent_artifact_path_trie",
+    srcs = ["ConcurrentArtifactPathTrie.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//third_party:guava",
+    ],
+)
+
+java_library(
+    name = "remote_rewound_action_synchronizer",
+    srcs = ["RemoteRewoundActionSynchronizer.java"],
+    deps = [
+        ":remote_action_input_fetcher",
+        "//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/lib/actions:file_metadata",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
+        "//src/main/java/com/google/devtools/build/lib/vfs:output_service",
+        "//third_party:caffeine",
+        "//third_party:guava",
+        "//third_party:jsr305",
+    ],
+)
+
 java_library(
     name = "remote_important_output_handler",
     srcs = ["RemoteImportantOutputHandler.java"],
     deps = [
         ":remote_output_checker",
+        ":remote_rewound_action_synchronizer",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//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/actions:important_output_handler",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/remote/common:bulk_transfer_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:output_service",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe",
         "//src/main/protobuf:failure_details_java_proto",
@@ -907,6 +940,8 @@ java_library(
         ":remote_action_file_system",
         ":remote_execution_cache",
         ":remote_output_checker",
+        ":remote_output_service",
+        ":remote_rewound_action_synchronizer",
         ":scrubber",
         ":upload_manifest",
         "//src/main/java/com/google/devtools/build/lib/actions",

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

@@ -0,0 +1,44 @@
+package com.google.devtools.build.lib.remote;
+
+import com.google.common.base.Preconditions;
+import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import java.util.concurrent.ConcurrentSkipListSet;
+
+/**
+ * A specialized concurrent trie that stores paths of artifacts and allows checking whether a given
+ * path is contained in (in the case of a tree artifact) or exactly matches (in any other case) an
+ * artifact in the trie.
+ */
+final class ConcurrentArtifactPathTrie {
+  // Invariant: no path in this set is a prefix of another path.
+  private final ConcurrentSkipListSet<PathFragment> paths =
+      new ConcurrentSkipListSet<>(PathFragment.HIERARCHICAL_COMPARATOR);
+
+  /**
+   * Adds the given {@link ActionInput} to the trie.
+   *
+   * <p>The caller must ensure that no object's path passed to this method is a prefix of any
+   * previously added object's path. Bazel enforces this for non-aggregate artifacts. Callers must
+   * not pass in {@link Artifact.TreeFileArtifact}s (which have exec paths that have their parent
+   * tree artifact's exec path as a prefix) or non-Artifact {@link ActionInput}s that violate this
+   * invariant.
+   */
+  void add(ActionInput input) {
+    Preconditions.checkArgument(
+        !(input instanceof Artifact.TreeFileArtifact),
+        "TreeFileArtifacts should not be added to the trie: %s",
+        input);
+    paths.add(input.getExecPath());
+  }
+
+  /** Checks whether the given {@link PathFragment} is contained in an artifact in the trie. */
+  boolean contains(PathFragment execPath) {
+    // By the invariant of this set, there is at most one prefix of execPath in the set. Since the
+    // comparator sorts all children of a path right after the path itself, if such a prefix
+    // exists, it must thus sort right before execPath (or be equal to it).
+    var floorPath = paths.floor(execPath);
+    return floorPath != null && execPath.startsWith(floorPath);
+  }
+}

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

@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.remote;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 
@@ -25,6 +24,7 @@
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionOutputDirectoryHelper;
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.VirtualActionInput;
 import com.google.devtools.build.lib.events.Reporter;
@@ -35,7 +35,9 @@
 import com.google.devtools.build.lib.util.TempPathGenerator;
 import com.google.devtools.build.lib.vfs.OutputPermissions;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.Symlinks;
 import java.io.IOException;
+import java.util.Collection;
 import javax.annotation.Nullable;
 
 /**
@@ -50,6 +52,7 @@ public class RemoteActionInputFetcher extends AbstractActionInputPrefetcher {
   private final String buildRequestId;
   private final String commandId;
   private final CombinedCache combinedCache;
+  private final ConcurrentArtifactPathTrie rewoundActionOutputs = new ConcurrentArtifactPathTrie();
 
   RemoteActionInputFetcher(
       Reporter reporter,
@@ -80,7 +83,18 @@ protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOExc
 
   @Override
   protected boolean canDownloadFile(Path path, FileArtifactValue metadata) {
-    return metadata.isRemote();
+    // When action rewinding is enabled, an action that had remote metadata at some point during the
+    // build may have been re-executed locally to regenerate lost inputs, but may then be rewound
+    // again and thus have its (now local) outputs deleted. In this case, we need to download the
+    // outputs again, even if they are now considered local.
+    return metadata.isRemote() || (forceRefetch(path) && !path.exists(Symlinks.NOFOLLOW));
+  }
+
+  @Override
+  protected boolean forceRefetch(Path path) {
+    // Caches for download operations and output directory creation need to be disregarded for the
+    // outputs of rewound actions as they may have been deleted after they were first created.
+    return path.startsWith(execRoot) && rewoundActionOutputs.contains(path.relativeTo(execRoot));
   }
 
   @Override
@@ -93,7 +107,6 @@ protected ListenableFuture<Void> doDownloadFile(
       Priority priority,
       Reason reason)
       throws IOException {
-    checkArgument(metadata.isRemote(), "Cannot download file that is not a remote file.");
     RequestMetadata requestMetadata =
         TracingMetadataUtils.buildMetadata(
             buildRequestId,
@@ -140,4 +153,22 @@ protected ListenableFuture<Void> doDownloadFile(
                 }),
         directExecutor());
   }
+
+  public void handleRewoundActionOutputs(Collection<Artifact> outputs) {
+    // SkyframeActionExecutor#prepareForRewinding does *not* call this method because the
+    // RemoteActionFileSystem corresponds to an ActionFileSystemType with inMemoryFileSystem() ==
+    // true. While it is true that resetting outputDirectoryHelper isn't necessary to undo the
+    // caching of output directory creation during action preparation, we still need to reset here
+    // since outputDirectoryHelper is also used by AbstractActionInputPrefetcher.
+    outputDirectoryHelper.invalidateTreeArtifactDirectoryCreation(outputs);
+    for (Artifact output : outputs) {
+      // Action templates have TreeFileArtifacts as outputs, which isn't supported by the trie. We
+      // only need to track the tree artifacts themselves.
+      if (output instanceof Artifact.TreeFileArtifact) {
+        rewoundActionOutputs.add(output.getParent());
+      } else {
+        rewoundActionOutputs.add(output);
+      }
+    }
+  }
 }

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

@@ -64,6 +64,8 @@
 import com.google.common.eventbus.Subscribe;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
@@ -151,8 +153,8 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.Semaphore;
@@ -192,9 +194,10 @@ public class RemoteExecutionService {
   private final Set<String> reportedErrors = new HashSet<>();
 
   @SuppressWarnings("AllowVirtualThreads")
-  private final ExecutorService backgroundTaskExecutor =
-      Executors.newThreadPerTaskExecutor(
-          Thread.ofVirtual().name("remote-execution-bg-", 0).factory());
+  private final ListeningExecutorService backgroundTaskExecutor =
+      MoreExecutors.listeningDecorator(
+          Executors.newThreadPerTaskExecutor(
+              Thread.ofVirtual().name("remote-execution-bg-", 0).factory()));
 
   private final AtomicBoolean shutdown = new AtomicBoolean(false);
   private final AtomicBoolean buildInterrupted = new AtomicBoolean(false);
@@ -1761,16 +1764,31 @@ public void uploadOutputs(
 
     if (remoteOptions.remoteCacheAsync
         && !action.getSpawn().getResourceOwner().mayModifySpawnOutputsAfterExecution()) {
-      backgroundTaskExecutor.execute(
-          () -> {
-            try {
-              doUploadOutputs(action, spawnResult, onUploadComplete);
-            } catch (ExecException e) {
-              reportUploadError(e);
-            } catch (InterruptedException ignored) {
-              // ThreadPerTaskExecutor does not care about interrupt status.
-            }
-          });
+      var uploadDone = new CountDownLatch(1);
+      var future =
+          backgroundTaskExecutor.submit(
+              () -> {
+                try {
+                  doUploadOutputs(action, spawnResult, onUploadComplete);
+                } catch (ExecException e) {
+                  reportUploadError(e);
+                } catch (InterruptedException ignored) {
+                  // ThreadPerTaskExecutor does not care about interrupt status.
+                } finally {
+                  uploadDone.countDown();
+                }
+              });
+
+      if (outputService instanceof RemoteOutputService remoteOutputService
+          && remoteOutputService.getRewoundActionSynchronizer()
+              instanceof RemoteRewoundActionSynchronizer remoteRewoundActionSynchronizer) {
+        remoteRewoundActionSynchronizer.registerOutputUploadTask(
+            action.getRemoteActionExecutionContext().getSpawnOwner(),
+            () -> {
+              future.cancel(true);
+              uploadDone.await();
+            });
+      }
     } else {
       doUploadOutputs(action, spawnResult, onUploadComplete);
     }