Expand tree outputs before eagerly prefetching them for local actions. (bazelbuild#17494)

PiperOrigin-RevId: 501500046
Change-Id: I220931eff70c4a9b04468ddf1f51f6cda91dfb8a

Author: coeuvre

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

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.remote;
 
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.util.concurrent.Futures.addCallback;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static com.google.devtools.build.lib.remote.util.RxFutures.toCompletable;
 import static com.google.devtools.build.lib.remote.util.RxFutures.toListenableFuture;
@@ -26,6 +27,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import com.google.common.flogger.GoogleLogger;
+import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionInput;
@@ -37,6 +39,9 @@
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.MetadataHandler;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
+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.util.AsyncTaskCache;
 import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult;
 import com.google.devtools.build.lib.remote.util.TempPathGenerator;
@@ -63,6 +68,7 @@
 public abstract class AbstractActionInputPrefetcher implements ActionInputPrefetcher {
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
+  private final Reporter reporter;
   private final AsyncTaskCache.NoResult<Path> downloadCache = AsyncTaskCache.NoResult.create();
   private final TempPathGenerator tempPathGenerator;
   protected final Set<Artifact> outputsAreInputs = Sets.newConcurrentHashSet();
@@ -109,9 +115,11 @@ protected enum Priority {
   }
 
   protected AbstractActionInputPrefetcher(
+      Reporter reporter,
       Path execRoot,
       TempPathGenerator tempPathGenerator,
       ImmutableList<Pattern> patternsToDownload) {
+    this.reporter = reporter;
     this.execRoot = execRoot;
     this.tempPathGenerator = tempPathGenerator;
     this.patternsToDownload = patternsToDownload;
@@ -538,17 +546,33 @@ public void shutdown() {
     }
   }
 
+  /** Event which is fired when inputs for local action are eagerly prefetched. */
+  public static class InputsEagerlyPrefetched implements Postable {
+    private final List<Artifact> artifacts;
+
+    public InputsEagerlyPrefetched(List<Artifact> artifacts) {
+      this.artifacts = artifacts;
+    }
+
+    public List<Artifact> getArtifacts() {
+      return artifacts;
+    }
+  }
+
   @SuppressWarnings({"CheckReturnValue", "FutureReturnValueIgnored"})
   public void finalizeAction(Action action, MetadataHandler metadataHandler) {
     List<Artifact> inputsToDownload = new ArrayList<>();
     List<Artifact> outputsToDownload = new ArrayList<>();
 
     for (Artifact output : action.getOutputs()) {
       if (outputsAreInputs.remove(output)) {
-        inputsToDownload.add(output);
-      }
-
-      if (output.isTreeArtifact()) {
+        if (output.isTreeArtifact()) {
+          var children = metadataHandler.getTreeArtifactChildren((SpecialArtifact) output);
+          inputsToDownload.addAll(children);
+        } else {
+          inputsToDownload.add(output);
+        }
+      } else if (output.isTreeArtifact()) {
         var children = metadataHandler.getTreeArtifactChildren((SpecialArtifact) output);
         for (var file : children) {
           if (outputMatchesPattern(file)) {
@@ -561,11 +585,42 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) {
     }
 
     if (!inputsToDownload.isEmpty()) {
-      prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH);
+      var future = prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH);
+      addCallback(
+          future,
+          new FutureCallback<Void>() {
+            @Override
+            public void onSuccess(Void unused) {
+              reporter.post(new InputsEagerlyPrefetched(inputsToDownload));
+            }
+
+            @Override
+            public void onFailure(Throwable throwable) {
+              reporter.handle(
+                  Event.warn(
+                      String.format(
+                          "Failed to eagerly prefetch inputs: %s", throwable.getMessage())));
+            }
+          },
+          directExecutor());
     }
 
     if (!outputsToDownload.isEmpty()) {
-      prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW);
+      var future = prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW);
+      addCallback(
+          future,
+          new FutureCallback<Void>() {
+            @Override
+            public void onSuccess(Void unused) {}
+
+            @Override
+            public void onFailure(Throwable throwable) {
+              reporter.handle(
+                  Event.warn(
+                      String.format("Failed to download outputs: %s", throwable.getMessage())));
+            }
+          },
+          directExecutor());
     }
   }
 

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

@@ -182,6 +182,7 @@ java_library(
         "//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/events",
         "//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",

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

@@ -23,6 +23,7 @@
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput.EmptyActionInput;
+import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.remote.common.BulkTransferException;
 import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
@@ -49,13 +50,14 @@ class RemoteActionInputFetcher extends AbstractActionInputPrefetcher {
   private final RemoteCache remoteCache;
 
   RemoteActionInputFetcher(
+      Reporter reporter,
       String buildRequestId,
       String commandId,
       RemoteCache remoteCache,
       Path execRoot,
       TempPathGenerator tempPathGenerator,
       ImmutableList<Pattern> patternsToDownload) {
-    super(execRoot, tempPathGenerator, patternsToDownload);
+    super(reporter, execRoot, tempPathGenerator, patternsToDownload);
     this.buildRequestId = Preconditions.checkNotNull(buildRequestId);
     this.commandId = Preconditions.checkNotNull(commandId);
     this.remoteCache = Preconditions.checkNotNull(remoteCache);

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

@@ -911,6 +911,7 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
       Preconditions.checkNotNull(patternsToDownload, "patternsToDownload must not be null");
       actionInputFetcher =
           new RemoteActionInputFetcher(
+              env.getReporter(),
               env.getBuildRequestId(),
               env.getCommandId().toString(),
               actionContextProvider.getRemoteCache(),

src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java

@@ -19,10 +19,12 @@
 import build.bazel.remote.execution.v2.Digest;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
+import com.google.common.eventbus.EventBus;
 import com.google.common.hash.HashCode;
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
+import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.InMemoryCacheClient;
@@ -60,7 +62,13 @@ public void setUp() throws IOException {
   protected AbstractActionInputPrefetcher createPrefetcher(Map<HashCode, byte[]> cas) {
     RemoteCache remoteCache = newCache(options, digestUtil, cas);
     return new RemoteActionInputFetcher(
-        "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
+        new Reporter(new EventBus()),
+        "none",
+        "none",
+        remoteCache,
+        execRoot,
+        tempPathGenerator,
+        ImmutableList.of());
   }
 
   @Test
@@ -70,7 +78,13 @@ public void testStagingVirtualActionInput() throws Exception {
     RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
     RemoteActionInputFetcher actionInputFetcher =
         new RemoteActionInputFetcher(
-            "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
+            new Reporter(new EventBus()),
+            "none",
+            "none",
+            remoteCache,
+            execRoot,
+            tempPathGenerator,
+            ImmutableList.of());
     VirtualActionInput a = ActionsTestUtil.createVirtualActionInput("file1", "hello world");
 
     // act
@@ -91,7 +105,13 @@ public void testStagingEmptyVirtualActionInput() throws Exception {
     RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
     RemoteActionInputFetcher actionInputFetcher =
         new RemoteActionInputFetcher(
-            "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
+            new Reporter(new EventBus()),
+            "none",
+            "none",
+            remoteCache,
+            execRoot,
+            tempPathGenerator,
+            ImmutableList.of());
 
     // act
     wait(