From 1be0ac3e73698e31a349ece629c887b06e102a0b Mon Sep 17 00:00:00 2001 From: Chi Wang Date: Wed, 15 Feb 2023 16:36:01 +0100 Subject: [PATCH] Expand tree outputs before eagerly prefetching them for local actions. (#17494) PiperOrigin-RevId: 501500046 Change-Id: I220931eff70c4a9b04468ddf1f51f6cda91dfb8a --- .../remote/AbstractActionInputPrefetcher.java | 67 +++++++++++++++++-- .../google/devtools/build/lib/remote/BUILD | 1 + .../lib/remote/RemoteActionInputFetcher.java | 4 +- .../build/lib/remote/RemoteModule.java | 1 + .../remote/RemoteActionInputFetcherTest.java | 26 ++++++- 5 files changed, 89 insertions(+), 10 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 9590f2c57cfa5a..6f2be02509b237 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 @@ -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 downloadCache = AsyncTaskCache.NoResult.create(); private final TempPathGenerator tempPathGenerator; protected final Set outputsAreInputs = Sets.newConcurrentHashSet(); @@ -109,9 +115,11 @@ protected enum Priority { } protected AbstractActionInputPrefetcher( + Reporter reporter, Path execRoot, TempPathGenerator tempPathGenerator, ImmutableList patternsToDownload) { + this.reporter = reporter; this.execRoot = execRoot; this.tempPathGenerator = tempPathGenerator; this.patternsToDownload = patternsToDownload; @@ -538,6 +546,19 @@ public void shutdown() { } } + /** Event which is fired when inputs for local action are eagerly prefetched. */ + public static class InputsEagerlyPrefetched implements Postable { + private final List artifacts; + + public InputsEagerlyPrefetched(List artifacts) { + this.artifacts = artifacts; + } + + public List getArtifacts() { + return artifacts; + } + } + @SuppressWarnings({"CheckReturnValue", "FutureReturnValueIgnored"}) public void finalizeAction(Action action, MetadataHandler metadataHandler) { List inputsToDownload = new ArrayList<>(); @@ -545,10 +566,13 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) { 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() { + @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() { + @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()); } } diff --git a/src/main/java/com/google/devtools/build/lib/remote/BUILD b/src/main/java/com/google/devtools/build/lib/remote/BUILD index 692df86d825b78..dfbe7adc482023 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/BUILD +++ b/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", 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 d128e83a6f6bad..696fa80394c225 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 @@ -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 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); 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 eb0a6eee5a7a2e..46d7a183193fc2 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 @@ -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(), 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 d4a4e456a652c7..ce034b84045aa2 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 @@ -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 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(