Remote: Use parameters instead of thread-local storage to provide tracing metadata. (Part 3)
Change RemoteCacheClient#downloadBlob to use RemoteActionExecutionContext.
PiperOrigin-RevId: 354239205
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
index 8e68aab..3950814 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
@@ -261,14 +261,15 @@
* @return a future that completes after the download completes (succeeds / fails). If successful,
* the content is stored in the future's {@code byte[]}.
*/
- public ListenableFuture<byte[]> downloadBlob(Digest digest) {
+ public ListenableFuture<byte[]> downloadBlob(
+ RemoteActionExecutionContext context, Digest digest) {
if (digest.getSizeBytes() == 0) {
return EMPTY_BYTES;
}
ByteArrayOutputStream bOut = new ByteArrayOutputStream((int) digest.getSizeBytes());
SettableFuture<byte[]> outerF = SettableFuture.create();
Futures.addCallback(
- cacheProtocol.downloadBlob(digest, bOut),
+ cacheProtocol.downloadBlob(context, digest, bOut),
new FutureCallback<Void>() {
@Override
public void onSuccess(Void aVoid) {
@@ -305,12 +306,13 @@
* @throws ExecException in case clean up after a failed download failed.
*/
public void download(
+ RemoteActionExecutionContext context,
ActionResult result,
Path execRoot,
FileOutErr origOutErr,
OutputFilesLocker outputFilesLocker)
throws ExecException, IOException, InterruptedException {
- ActionResultMetadata metadata = parseActionResultMetadata(result, execRoot);
+ ActionResultMetadata metadata = parseActionResultMetadata(context, result, execRoot);
List<ListenableFuture<FileMetadata>> downloads =
Stream.concat(
@@ -321,7 +323,7 @@
(file) -> {
try {
ListenableFuture<Void> download =
- downloadFile(toTmpDownloadPath(file.path()), file.digest());
+ downloadFile(context, toTmpDownloadPath(file.path()), file.digest());
return Futures.transform(download, (d) -> file, directExecutor());
} catch (IOException e) {
return Futures.<FileMetadata>immediateFailedFuture(e);
@@ -337,7 +339,7 @@
if (origOutErr != null) {
tmpOutErr = origOutErr.childOutErr();
}
- downloads.addAll(downloadOutErr(result, tmpOutErr));
+ downloads.addAll(downloadOutErr(context, result, tmpOutErr));
try {
waitForBulkTransfer(downloads, /* cancelRemainingOnInterrupt=*/ true);
@@ -449,7 +451,8 @@
}
/** Downloads a file (that is not a directory). The content is fetched from the digest. */
- public ListenableFuture<Void> downloadFile(Path path, Digest digest) throws IOException {
+ public ListenableFuture<Void> downloadFile(
+ RemoteActionExecutionContext context, Path path, Digest digest) throws IOException {
Preconditions.checkNotNull(path.getParentDirectory()).createDirectoryAndParents();
if (digest.getSizeBytes() == 0) {
// Handle empty file locally.
@@ -472,7 +475,7 @@
OutputStream out = new LazyFileOutputStream(path);
SettableFuture<Void> outerF = SettableFuture.create();
- ListenableFuture<Void> f = cacheProtocol.downloadBlob(digest, out);
+ ListenableFuture<Void> f = cacheProtocol.downloadBlob(context, digest, out);
Futures.addCallback(
f,
new FutureCallback<Void>() {
@@ -509,7 +512,8 @@
return outerF;
}
- private List<ListenableFuture<FileMetadata>> downloadOutErr(ActionResult result, OutErr outErr) {
+ private List<ListenableFuture<FileMetadata>> downloadOutErr(
+ RemoteActionExecutionContext context, ActionResult result, OutErr outErr) {
List<ListenableFuture<FileMetadata>> downloads = new ArrayList<>();
if (!result.getStdoutRaw().isEmpty()) {
try {
@@ -521,7 +525,8 @@
} else if (result.hasStdoutDigest()) {
downloads.add(
Futures.transform(
- cacheProtocol.downloadBlob(result.getStdoutDigest(), outErr.getOutputStream()),
+ cacheProtocol.downloadBlob(
+ context, result.getStdoutDigest(), outErr.getOutputStream()),
(d) -> null,
directExecutor()));
}
@@ -535,7 +540,8 @@
} else if (result.hasStderrDigest()) {
downloads.add(
Futures.transform(
- cacheProtocol.downloadBlob(result.getStderrDigest(), outErr.getErrorStream()),
+ cacheProtocol.downloadBlob(
+ context, result.getStderrDigest(), outErr.getErrorStream()),
(d) -> null,
directExecutor()));
}
@@ -549,6 +555,7 @@
* <p>This method only downloads output directory metadata, stdout and stderr as well as the
* contents of {@code inMemoryOutputPath} if specified.
*
+ * @param context the context this action running with
* @param result the action result metadata of a successfully executed action (exit code = 0).
* @param outputs the action's declared output files
* @param inMemoryOutputPath the path of an output file whose contents should be returned in
@@ -564,6 +571,7 @@
*/
@Nullable
public InMemoryOutput downloadMinimal(
+ RemoteActionExecutionContext context,
String actionId,
ActionResult result,
Collection<? extends ActionInput> outputs,
@@ -579,7 +587,7 @@
ActionResultMetadata metadata;
try (SilentCloseable c = Profiler.instance().profile("Remote.parseActionResultMetadata")) {
- metadata = parseActionResultMetadata(result, execRoot);
+ metadata = parseActionResultMetadata(context, result, execRoot);
}
if (!metadata.symlinks().isEmpty()) {
@@ -614,9 +622,10 @@
try (SilentCloseable c = Profiler.instance().profile("Remote.download")) {
ListenableFuture<byte[]> inMemoryOutputDownload = null;
if (inMemoryOutput != null) {
- inMemoryOutputDownload = downloadBlob(inMemoryOutputDigest);
+ inMemoryOutputDownload = downloadBlob(context, inMemoryOutputDigest);
}
- waitForBulkTransfer(downloadOutErr(result, outErr), /* cancelRemainingOnInterrupt=*/ true);
+ waitForBulkTransfer(
+ downloadOutErr(context, result, outErr), /* cancelRemainingOnInterrupt=*/ true);
if (inMemoryOutputDownload != null) {
waitForBulkTransfer(
ImmutableList.of(inMemoryOutputDownload), /* cancelRemainingOnInterrupt=*/ true);
@@ -708,7 +717,8 @@
return new DirectoryMetadata(filesBuilder.build(), symlinksBuilder.build());
}
- private ActionResultMetadata parseActionResultMetadata(ActionResult actionResult, Path execRoot)
+ private ActionResultMetadata parseActionResultMetadata(
+ RemoteActionExecutionContext context, ActionResult actionResult, Path execRoot)
throws IOException, InterruptedException {
Preconditions.checkNotNull(actionResult, "actionResult");
Map<Path, ListenableFuture<Tree>> dirMetadataDownloads =
@@ -717,7 +727,7 @@
dirMetadataDownloads.put(
execRoot.getRelative(dir.getPath()),
Futures.transform(
- downloadBlob(dir.getTreeDigest()),
+ downloadBlob(context, dir.getTreeDigest()),
(treeBytes) -> {
try {
return Tree.parseFrom(treeBytes);
