Skip to content

Commit

Permalink
Expand tree outputs before eagerly prefetching them for local actions. (
Browse files Browse the repository at this point in the history 
bazelbuild#17494)

PiperOrigin-RevId: 501500046
Change-Id: I220931eff70c4a9b04468ddf1f51f6cda91dfb8a
  • Loading branch information
@coeuvre
coeuvre authored Feb 15, 2023
1 parent 4e35c02 commit 1be0ac3
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 10 deletions.
67 changes: 61 additions & 6 deletions src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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();
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)) {
Expand All @@ -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());
}
}

Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/google/devtools/build/lib/remote/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
4 changes: 3 additions & 1 deletion src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
26 changes: 23 additions & 3 deletions src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down

0 comments on commit 1be0ac3

Please sign in to comment.