Skip to content

Commit 1be0ac3

Browse files
coeuvrecoeuvre
authored
Expand tree outputs before eagerly prefetching them for local actions. (bazelbuild#17494)
PiperOrigin-RevId: 501500046 Change-Id: I220931eff70c4a9b04468ddf1f51f6cda91dfb8a
1 parent 4e35c02 commit 1be0ac3

5 files changed

Lines changed: 89 additions & 10 deletions

File tree

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

Lines changed: 61 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
package com.google.devtools.build.lib.remote;
1515

1616
import static com.google.common.base.Preconditions.checkState;
17+
import static com.google.common.util.concurrent.Futures.addCallback;
1718
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
1819
import static com.google.devtools.build.lib.remote.util.RxFutures.toCompletable;
1920
import static com.google.devtools.build.lib.remote.util.RxFutures.toListenableFuture;
@@ -26,6 +27,7 @@
2627
import com.google.common.collect.ImmutableSet;
2728
import com.google.common.collect.Sets;
2829
import com.google.common.flogger.GoogleLogger;
30+
import com.google.common.util.concurrent.FutureCallback;
2931
import com.google.common.util.concurrent.ListenableFuture;
3032
import com.google.devtools.build.lib.actions.Action;
3133
import com.google.devtools.build.lib.actions.ActionInput;
@@ -37,6 +39,9 @@
3739
import com.google.devtools.build.lib.actions.MetadataProvider;
3840
import com.google.devtools.build.lib.actions.cache.MetadataHandler;
3941
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
42+
import com.google.devtools.build.lib.events.Event;
43+
import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
44+
import com.google.devtools.build.lib.events.Reporter;
4045
import com.google.devtools.build.lib.remote.util.AsyncTaskCache;
4146
import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult;
4247
import com.google.devtools.build.lib.remote.util.TempPathGenerator;
@@ -63,6 +68,7 @@
6368
public abstract class AbstractActionInputPrefetcher implements ActionInputPrefetcher {
6469
private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
6570

71+
private final Reporter reporter;
6672
private final AsyncTaskCache.NoResult<Path> downloadCache = AsyncTaskCache.NoResult.create();
6773
private final TempPathGenerator tempPathGenerator;
6874
protected final Set<Artifact> outputsAreInputs = Sets.newConcurrentHashSet();
@@ -109,9 +115,11 @@ protected enum Priority {
109115
}
110116

111117
protected AbstractActionInputPrefetcher(
118+
Reporter reporter,
112119
Path execRoot,
113120
TempPathGenerator tempPathGenerator,
114121
ImmutableList<Pattern> patternsToDownload) {
122+
this.reporter = reporter;
115123
this.execRoot = execRoot;
116124
this.tempPathGenerator = tempPathGenerator;
117125
this.patternsToDownload = patternsToDownload;
@@ -538,17 +546,33 @@ public void shutdown() {
538546
}
539547
}
540548

549+
/** Event which is fired when inputs for local action are eagerly prefetched. */
550+
public static class InputsEagerlyPrefetched implements Postable {
551+
private final List<Artifact> artifacts;
552+
553+
public InputsEagerlyPrefetched(List<Artifact> artifacts) {
554+
this.artifacts = artifacts;
555+
}
556+
557+
public List<Artifact> getArtifacts() {
558+
return artifacts;
559+
}
560+
}
561+
541562
@SuppressWarnings({"CheckReturnValue", "FutureReturnValueIgnored"})
542563
public void finalizeAction(Action action, MetadataHandler metadataHandler) {
543564
List<Artifact> inputsToDownload = new ArrayList<>();
544565
List<Artifact> outputsToDownload = new ArrayList<>();
545566

546567
for (Artifact output : action.getOutputs()) {
547568
if (outputsAreInputs.remove(output)) {
548-
inputsToDownload.add(output);
549-
}
550-
551-
if (output.isTreeArtifact()) {
569+
if (output.isTreeArtifact()) {
570+
var children = metadataHandler.getTreeArtifactChildren((SpecialArtifact) output);
571+
inputsToDownload.addAll(children);
572+
} else {
573+
inputsToDownload.add(output);
574+
}
575+
} else if (output.isTreeArtifact()) {
552576
var children = metadataHandler.getTreeArtifactChildren((SpecialArtifact) output);
553577
for (var file : children) {
554578
if (outputMatchesPattern(file)) {
@@ -561,11 +585,42 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) {
561585
}
562586

563587
if (!inputsToDownload.isEmpty()) {
564-
prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH);
588+
var future = prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH);
589+
addCallback(
590+
future,
591+
new FutureCallback<Void>() {
592+
@Override
593+
public void onSuccess(Void unused) {
594+
reporter.post(new InputsEagerlyPrefetched(inputsToDownload));
595+
}
596+
597+
@Override
598+
public void onFailure(Throwable throwable) {
599+
reporter.handle(
600+
Event.warn(
601+
String.format(
602+
"Failed to eagerly prefetch inputs: %s", throwable.getMessage())));
603+
}
604+
},
605+
directExecutor());
565606
}
566607

567608
if (!outputsToDownload.isEmpty()) {
568-
prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW);
609+
var future = prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW);
610+
addCallback(
611+
future,
612+
new FutureCallback<Void>() {
613+
@Override
614+
public void onSuccess(Void unused) {}
615+
616+
@Override
617+
public void onFailure(Throwable throwable) {
618+
reporter.handle(
619+
Event.warn(
620+
String.format("Failed to download outputs: %s", throwable.getMessage())));
621+
}
622+
},
623+
directExecutor());
569624
}
570625
}
571626

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ java_library(
182182
"//src/main/java/com/google/devtools/build/lib/actions",
183183
"//src/main/java/com/google/devtools/build/lib/actions:artifacts",
184184
"//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
185+
"//src/main/java/com/google/devtools/build/lib/events",
185186
"//src/main/java/com/google/devtools/build/lib/remote/util",
186187
"//src/main/java/com/google/devtools/build/lib/vfs",
187188
"//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.devtools.build.lib.actions.FileArtifactValue;
2424
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
2525
import com.google.devtools.build.lib.actions.cache.VirtualActionInput.EmptyActionInput;
26+
import com.google.devtools.build.lib.events.Reporter;
2627
import com.google.devtools.build.lib.remote.common.BulkTransferException;
2728
import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
2829
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
@@ -49,13 +50,14 @@ class RemoteActionInputFetcher extends AbstractActionInputPrefetcher {
4950
private final RemoteCache remoteCache;
5051

5152
RemoteActionInputFetcher(
53+
Reporter reporter,
5254
String buildRequestId,
5355
String commandId,
5456
RemoteCache remoteCache,
5557
Path execRoot,
5658
TempPathGenerator tempPathGenerator,
5759
ImmutableList<Pattern> patternsToDownload) {
58-
super(execRoot, tempPathGenerator, patternsToDownload);
60+
super(reporter, execRoot, tempPathGenerator, patternsToDownload);
5961
this.buildRequestId = Preconditions.checkNotNull(buildRequestId);
6062
this.commandId = Preconditions.checkNotNull(commandId);
6163
this.remoteCache = Preconditions.checkNotNull(remoteCache);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -911,6 +911,7 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
911911
Preconditions.checkNotNull(patternsToDownload, "patternsToDownload must not be null");
912912
actionInputFetcher =
913913
new RemoteActionInputFetcher(
914+
env.getReporter(),
914915
env.getBuildRequestId(),
915916
env.getCommandId().toString(),
916917
actionContextProvider.getRemoteCache(),

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

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@
1919
import build.bazel.remote.execution.v2.Digest;
2020
import com.google.common.collect.ImmutableList;
2121
import com.google.common.collect.Maps;
22+
import com.google.common.eventbus.EventBus;
2223
import com.google.common.hash.HashCode;
2324
import com.google.devtools.build.lib.actions.MetadataProvider;
2425
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
2526
import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
27+
import com.google.devtools.build.lib.events.Reporter;
2628
import com.google.devtools.build.lib.remote.options.RemoteOptions;
2729
import com.google.devtools.build.lib.remote.util.DigestUtil;
2830
import com.google.devtools.build.lib.remote.util.InMemoryCacheClient;
@@ -60,7 +62,13 @@ public void setUp() throws IOException {
6062
protected AbstractActionInputPrefetcher createPrefetcher(Map<HashCode, byte[]> cas) {
6163
RemoteCache remoteCache = newCache(options, digestUtil, cas);
6264
return new RemoteActionInputFetcher(
63-
"none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
65+
new Reporter(new EventBus()),
66+
"none",
67+
"none",
68+
remoteCache,
69+
execRoot,
70+
tempPathGenerator,
71+
ImmutableList.of());
6472
}
6573

6674
@Test
@@ -70,7 +78,13 @@ public void testStagingVirtualActionInput() throws Exception {
7078
RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
7179
RemoteActionInputFetcher actionInputFetcher =
7280
new RemoteActionInputFetcher(
73-
"none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
81+
new Reporter(new EventBus()),
82+
"none",
83+
"none",
84+
remoteCache,
85+
execRoot,
86+
tempPathGenerator,
87+
ImmutableList.of());
7488
VirtualActionInput a = ActionsTestUtil.createVirtualActionInput("file1", "hello world");
7589

7690
// act
@@ -91,7 +105,13 @@ public void testStagingEmptyVirtualActionInput() throws Exception {
91105
RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
92106
RemoteActionInputFetcher actionInputFetcher =
93107
new RemoteActionInputFetcher(
94-
"none", "none", remoteCache, execRoot, tempPathGenerator, ImmutableList.of());
108+
new Reporter(new EventBus()),
109+
"none",
110+
"none",
111+
remoteCache,
112+
execRoot,
113+
tempPathGenerator,
114+
ImmutableList.of());
95115

96116
// act
97117
wait(

0 commit comments

Comments
 (0)