From e01e7f51dd19f39ce3bc0718cec20ed6474de733 Mon Sep 17 00:00:00 2001 From: Googler Date: Tue, 18 Oct 2022 05:31:54 -0700 Subject: [PATCH] Download outputs that match the regex in the background instead of blocking spawn execution PiperOrigin-RevId: 481894591 Change-Id: I3bcfad6e878288d1b014da85108f6a2e25b303ca --- .../remote/AbstractActionInputPrefetcher.java | 41 +++++++++++++- .../lib/remote/RemoteActionInputFetcher.java | 7 ++- .../lib/remote/RemoteExecutionService.java | 53 ------------------- .../build/lib/remote/RemoteModule.java | 19 ++++++- .../build/lib/remote/RemoteOutputService.java | 4 +- .../remote/RemoteActionInputFetcherTest.java | 9 ++-- 6 files changed, 72 insertions(+), 61 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java index 5dcd60ac508624..8db09901ac7ab6 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java +++ b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java @@ -22,9 +22,12 @@ import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Sets; import com.google.common.flogger.GoogleLogger; import com.google.common.util.concurrent.ListenableFuture; +import com.google.devtools.build.lib.actions.Action; import com.google.devtools.build.lib.actions.ActionInput; import com.google.devtools.build.lib.actions.ActionInputPrefetcher; import com.google.devtools.build.lib.actions.Artifact; @@ -32,6 +35,7 @@ import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact; import com.google.devtools.build.lib.actions.FileArtifactValue; 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.remote.util.AsyncTaskCache; import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult; @@ -48,7 +52,9 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.regex.Pattern; /** * Abstract implementation of {@link ActionInputPrefetcher} which implements the orchestration of @@ -59,8 +65,10 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet private final AsyncTaskCache.NoResult downloadCache = AsyncTaskCache.NoResult.create(); private final TempPathGenerator tempPathGenerator; + protected final Set outputsAreInputs = Sets.newConcurrentHashSet(); protected final Path execRoot; + protected final ImmutableList patternsToDownload; /** Priority for the staging task. */ protected enum Priority { @@ -86,9 +94,13 @@ protected enum Priority { LOW, } - protected AbstractActionInputPrefetcher(Path execRoot, TempPathGenerator tempPathGenerator) { + protected AbstractActionInputPrefetcher( + Path execRoot, + TempPathGenerator tempPathGenerator, + ImmutableList patternsToDownload) { this.execRoot = execRoot; this.tempPathGenerator = tempPathGenerator; + this.patternsToDownload = patternsToDownload; } protected abstract boolean shouldDownloadFile(Path path, FileArtifactValue metadata); @@ -360,4 +372,31 @@ public void shutdown() { } } } + + @SuppressWarnings({"CheckReturnValue", "FutureReturnValueIgnored"}) + public void finalizeAction(Action action, MetadataHandler metadataHandler) { + List inputsToDownload = new ArrayList<>(); + List outputsToDownload = new ArrayList<>(); + + for (Artifact output : action.getOutputs()) { + if (outputsAreInputs.remove(output)) { + inputsToDownload.add(output); + } + + for (Pattern pattern : patternsToDownload) { + if (pattern.matcher(output.getExecPathString()).matches()) { + outputsToDownload.add(output); + break; + } + } + } + + if (!inputsToDownload.isEmpty()) { + prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH); + } + + if (!outputsToDownload.isEmpty()) { + prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW); + } + } } diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java index 673f6a233fdef9..413e70469a6550 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java @@ -18,6 +18,7 @@ import build.bazel.remote.execution.v2.Digest; import build.bazel.remote.execution.v2.RequestMetadata; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.ListenableFuture; import com.google.devtools.build.lib.actions.FileArtifactValue; import com.google.devtools.build.lib.actions.cache.VirtualActionInput; @@ -33,6 +34,7 @@ import com.google.devtools.build.lib.vfs.PathFragment; import io.reactivex.rxjava3.core.Completable; import java.io.IOException; +import java.util.regex.Pattern; /** * Stages output files that are stored remotely to the local filesystem. @@ -51,8 +53,9 @@ class RemoteActionInputFetcher extends AbstractActionInputPrefetcher { String commandId, RemoteCache remoteCache, Path execRoot, - TempPathGenerator tempPathGenerator) { - super(execRoot, tempPathGenerator); + TempPathGenerator tempPathGenerator, + ImmutableList patternsToDownload) { + super(execRoot, tempPathGenerator, patternsToDownload); this.buildRequestId = Preconditions.checkNotNull(buildRequestId); this.commandId = Preconditions.checkNotNull(commandId); this.remoteCache = Preconditions.checkNotNull(remoteCache); diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java index fd7ee5c75acf33..f4a6dea997bbe5 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java @@ -141,7 +141,6 @@ import java.util.concurrent.Phaser; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Predicate; -import java.util.regex.Pattern; import javax.annotation.Nullable; /** @@ -170,8 +169,6 @@ public class RemoteExecutionService { private final AtomicBoolean shutdown = new AtomicBoolean(false); private final AtomicBoolean buildInterrupted = new AtomicBoolean(false); - private final boolean shouldForceDownloads; - private final Predicate shouldForceDownloadPredicate; public RemoteExecutionService( Executor executor, @@ -215,27 +212,6 @@ public RemoteExecutionService( this.captureCorruptedOutputsDir = captureCorruptedOutputsDir; this.scheduler = Schedulers.from(executor, /*interruptibleWorker=*/ true); - - // TODO(bazel-team): Consider adding a warning or more validation if the remoteDownloadRegex is - // used without Build without the Bytes. - ImmutableList.Builder builder = ImmutableList.builder(); - if (remoteOptions.remoteOutputsMode == RemoteOutputsMode.MINIMAL - || remoteOptions.remoteOutputsMode == RemoteOutputsMode.TOPLEVEL) { - for (String regex : remoteOptions.remoteDownloadRegex) { - builder.add(Pattern.compile(regex)); - } - } - ImmutableList patterns = builder.build(); - this.shouldForceDownloads = !patterns.isEmpty(); - this.shouldForceDownloadPredicate = - path -> { - for (Pattern pattern : patterns) { - if (pattern.matcher(path).matches()) { - return true; - } - } - return false; - }; } static Command buildCommand( @@ -1004,8 +980,6 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re /* exitCode = */ result.getExitCode(), hasFilesToDownload(action.getSpawn().getOutputFiles(), filesToDownload)); - ImmutableList> forcedDownloads = ImmutableList.of(); - // Download into temporary paths, then move everything at the end. // This avoids holding the output lock while downloading, which would prevent the local branch // from completing sooner under the dynamic execution strategy. @@ -1030,15 +1004,6 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re "Symlinks in action outputs are not yet supported by " + "--experimental_remote_download_outputs=minimal"); } - if (shouldForceDownloads) { - forcedDownloads = - buildFilesToDownloadWithPredicate( - context, - progressStatusListener, - metadata, - shouldForceDownloadPredicate, - realToTmpPath); - } } FileOutErr tmpOutErr = outErr.childOutErr(); @@ -1059,17 +1024,6 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re throw e; } - // TODO(bazel-team): Unify this block with the equivalent block above. - try (SilentCloseable c = Profiler.instance().profile("Remote.forcedDownload")) { - waitForBulkTransfer(forcedDownloads, /* cancelRemainingOnInterrupt= */ true); - } catch (Exception e) { - // TODO(bazel-team): Consider adding better case-by-case exception handling instead of just - // rethrowing - captureCorruptedOutputs(e); - deletePartialDownloadedOutputs(realToTmpPath, tmpOutErr, e); - throw e; - } - FileOutErr.dump(tmpOutErr, outErr); // Ensure that we are the only ones writing to the output files when using the dynamic spawn @@ -1109,13 +1063,6 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re // might not be supported on all platforms. createSymlinks(symlinks); } else { - // TODO(bazel-team): We should unify this if-block to rely on downloadOutputs above but, as of - // 2022-07-05, downloadOuputs' semantics isn't exactly the same as build-without-the-bytes - // which is necessary for using remoteDownloadRegex. - if (!forcedDownloads.isEmpty()) { - moveOutputsToFinalLocation(forcedDownloads, realToTmpPath); - } - ActionInput inMemoryOutput = null; Digest inMemoryOutputDigest = null; PathFragment inMemoryOutputPath = getInMemoryOutputPath(action.getSpawn()); diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java index 248eda9ac6fb24..3729bbf53a0cf6 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java @@ -120,6 +120,7 @@ import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; +import java.util.regex.Pattern; import javax.annotation.Nullable; /** RemoteModule provides distributed cache and remote execution for Bazel. */ @@ -143,6 +144,7 @@ public final class RemoteModule extends BlazeModule { @Nullable private RemoteOutputService remoteOutputService; @Nullable private TempPathGenerator tempPathGenerator; @Nullable private BlockWaitingModule blockWaitingModule; + @Nullable private ImmutableList patternsToDownload; private ChannelFactory channelFactory = new ChannelFactory() { @@ -277,6 +279,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { Preconditions.checkState(actionInputFetcher == null, "actionInputFetcher must be null"); Preconditions.checkState(remoteOptions == null, "remoteOptions must be null"); Preconditions.checkState(tempPathGenerator == null, "tempPathGenerator must be null"); + Preconditions.checkState(patternsToDownload == null, "patternsToDownload must be null"); RemoteOptions remoteOptions = env.getOptions().getOptions(RemoteOptions.class); if (remoteOptions == null) { @@ -326,6 +329,16 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { FailureDetails.RemoteOptions.Code.EXECUTION_WITH_INVALID_CACHE); } + // TODO(bazel-team): Consider adding a warning or more validation if the remoteDownloadRegex is + // used without Build without the Bytes. + if (!remoteOptions.remoteOutputsMode.downloadAllOutputs()) { + ImmutableList.Builder patternsToDownloadBuilder = ImmutableList.builder(); + for (String regex : remoteOptions.remoteDownloadRegex) { + patternsToDownloadBuilder.add(Pattern.compile(regex)); + } + patternsToDownload = patternsToDownloadBuilder.build(); + } + env.getEventBus().register(this); String invocationId = env.getCommandId().toString(); String buildRequestId = env.getBuildRequestId(); @@ -914,6 +927,7 @@ public void afterCommand() throws AbruptExitException { remoteOutputService = null; tempPathGenerator = null; rpcLogFile = null; + patternsToDownload = null; } private static void afterCommandTask( @@ -1017,13 +1031,16 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB RemoteOutputsMode remoteOutputsMode = remoteOptions.remoteOutputsMode; if (!remoteOutputsMode.downloadAllOutputs() && actionContextProvider.getRemoteCache() != null) { + Preconditions.checkNotNull(patternsToDownload, "patternsToDownload must not be null"); actionInputFetcher = new RemoteActionInputFetcher( env.getBuildRequestId(), env.getCommandId().toString(), actionContextProvider.getRemoteCache(), env.getExecRoot(), - tempPathGenerator); + tempPathGenerator, + patternsToDownload); + env.getEventBus().register(actionInputFetcher); builder.setActionInputPrefetcher(actionInputFetcher); remoteOutputService.setActionInputFetcher(actionInputFetcher); actionContextProvider.setActionInputFetcher(actionInputFetcher); diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java index a1329d4c513e40..ede863b922d4ba 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java @@ -107,7 +107,9 @@ public void flushActionFileSystem(FileSystem actionFileSystem) throws IOExceptio @Override public void finalizeAction(Action action, MetadataHandler metadataHandler) { - // Intentionally left empty. + if (actionInputFetcher != null) { + actionInputFetcher.finalizeAction(action, metadataHandler); + } } @Nullable diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java index 352f258ce94095..d4a4e456a652c7 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java @@ -59,7 +59,8 @@ public void setUp() throws IOException { @Override protected AbstractActionInputPrefetcher createPrefetcher(Map cas) { RemoteCache remoteCache = newCache(options, digestUtil, cas); - return new RemoteActionInputFetcher("none", "none", remoteCache, execRoot, tempPathGenerator); + return new RemoteActionInputFetcher( + "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of()); } @Test @@ -68,7 +69,8 @@ public void testStagingVirtualActionInput() throws Exception { MetadataProvider metadataProvider = new StaticMetadataProvider(new HashMap<>()); RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>()); RemoteActionInputFetcher actionInputFetcher = - new RemoteActionInputFetcher("none", "none", remoteCache, execRoot, tempPathGenerator); + new RemoteActionInputFetcher( + "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of()); VirtualActionInput a = ActionsTestUtil.createVirtualActionInput("file1", "hello world"); // act @@ -88,7 +90,8 @@ public void testStagingEmptyVirtualActionInput() throws Exception { MetadataProvider metadataProvider = new StaticMetadataProvider(new HashMap<>()); RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>()); RemoteActionInputFetcher actionInputFetcher = - new RemoteActionInputFetcher("none", "none", remoteCache, execRoot, tempPathGenerator); + new RemoteActionInputFetcher( + "none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of()); // act wait(