Remote: Use parameters instead of thread-local storage to provide tra…

…cing metadata. (Part 4)

Change RemoteCacheClient#upload{File,Blob} to use RemoteActionExecutionContext.

PiperOrigin-RevId: 354472775

src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.remote;
 
 import build.bazel.remote.execution.v2.Digest;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -29,7 +30,11 @@
 import com.google.devtools.build.lib.buildeventstream.PathConverter;
 import com.google.devtools.build.lib.collect.ImmutableIterable;
 import com.google.devtools.build.lib.remote.common.MissingDigestsFinder;
+import com.google.devtools.build.lib.remote.common.NetworkTime;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContextImpl;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
+import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import io.grpc.Context;
 import io.netty.util.AbstractReferenceCounted;
@@ -50,7 +55,8 @@ class ByteStreamBuildEventArtifactUploader extends AbstractReferenceCounted
  implements BuildEventArtifactUploader {
 
  private final ListeningExecutorService uploadExecutor;
- private final Context ctx;
+ private final String buildRequestId;
+ private final String commandId;
  private final ByteStreamUploader uploader;
  private final String remoteServerInstanceName;
  private final MissingDigestsFinder missingDigestsFinder;
@@ -61,15 +67,17 @@ class ByteStreamBuildEventArtifactUploader extends AbstractReferenceCounted
  ByteStreamUploader uploader,
  MissingDigestsFinder missingDigestsFinder,
  String remoteServerName,
- Context ctx,
+ String buildRequestId,
+ String commandId,
  @Nullable String remoteInstanceName,
  int maxUploadThreads) {
  this.uploader = Preconditions.checkNotNull(uploader);
  String remoteServerInstanceName = Preconditions.checkNotNull(remoteServerName);
  if (!Strings.isNullOrEmpty(remoteInstanceName)) {
  remoteServerInstanceName += "/" + remoteInstanceName;
  }
- this.ctx = ctx;
+ this.buildRequestId = buildRequestId;
+ this.commandId = commandId;
  this.remoteServerInstanceName = remoteServerInstanceName;
  // Limit the maximum threads number to 1000 (chosen arbitrarily)
  this.uploadExecutor =
@@ -153,6 +161,8 @@ private static List<PathMetadata> processQueryResult(
  */
  private ListenableFuture<ImmutableIterable<PathMetadata>> queryRemoteCache(
  ImmutableList<ListenableFuture<PathMetadata>> allPaths) throws Exception {
+ Context ctx = TracingMetadataUtils.contextWithMetadata(buildRequestId, commandId, "bes-upload");
+
  List<PathMetadata> knownRemotePaths = new ArrayList<>(allPaths.size());
  List<PathMetadata> filesToQuery = new ArrayList<>();
  Set<Digest> digestsToQuery = new HashSet<>();
@@ -185,19 +195,20 @@ private ListenableFuture<ImmutableIterable<PathMetadata>> queryRemoteCache(
  */
  private ListenableFuture<List<PathMetadata>> uploadLocalFiles(
  ImmutableIterable<PathMetadata> allPaths) {
+ RequestMetadata metadata =
+ TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "bes-upload");
+ RemoteActionExecutionContext context =
+ new RemoteActionExecutionContextImpl(metadata, new NetworkTime());
+
  ImmutableList.Builder<ListenableFuture<PathMetadata>> allPathsUploaded =
  ImmutableList.builder();
  for (PathMetadata path : allPaths) {
  if (!path.isRemote() && !path.isDirectory()) {
  Chunker chunker =
  Chunker.builder().setInput(path.getDigest().getSizeBytes(), path.getPath()).build();
  final ListenableFuture<Void> upload;
- Context prevCtx = ctx.attach();
- try {
- upload = uploader.uploadBlobAsync(path.getDigest(), chunker, /* forceUpload=*/ false);
- } finally {
- ctx.detach(prevCtx);
- }
+ upload =
+ uploader.uploadBlobAsync(context, path.getDigest(), chunker, /* forceUpload= */ false);
  allPathsUploaded.add(Futures.transform(upload, unused -> path, uploadExecutor));
  } else {
  allPathsUploaded.add(Futures.immediateFuture(path));

src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploaderFactory.java

@@ -18,7 +18,6 @@
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
 import com.google.devtools.build.lib.runtime.BuildEventArtifactUploaderFactory;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
-import io.grpc.Context;
 import javax.annotation.Nullable;
 
 /**
@@ -29,20 +28,23 @@ class ByteStreamBuildEventArtifactUploaderFactory implements
 
  private final ByteStreamUploader uploader;
  private final String remoteServerName;
- private final Context ctx;
+ private final String buildRequestId;
+ private final String commandId;
  private final MissingDigestsFinder missingDigestsFinder;
  @Nullable private final String remoteInstanceName;
 
  ByteStreamBuildEventArtifactUploaderFactory(
  ByteStreamUploader uploader,
  MissingDigestsFinder missingDigestsFinder,
  String remoteServerName,
- Context ctx,
+ String buildRequestId,
+ String commandId,
  @Nullable String remoteInstanceName) {
  this.uploader = uploader;
  this.missingDigestsFinder = missingDigestsFinder;
  this.remoteServerName = remoteServerName;
- this.ctx = ctx;
+ this.buildRequestId = buildRequestId;
+ this.commandId = commandId;
  this.remoteInstanceName = remoteInstanceName;
  }
 
@@ -52,7 +54,8 @@ public BuildEventArtifactUploader create(CommandEnvironment env) {
  uploader.retain(),
  missingDigestsFinder,
  remoteServerName,
- ctx,
+ buildRequestId,
+ commandId,
  remoteInstanceName,
  env.getOptions().getOptions(RemoteOptions.class).buildEventUploadMaxThreads);
  }

src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java

@@ -36,12 +36,12 @@
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
 import com.google.devtools.build.lib.remote.RemoteRetrier.ProgressiveBackoff;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
 import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.remote.util.Utils;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
-import io.grpc.Context;
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.Status.Code;
@@ -134,9 +134,10 @@ public ByteStreamUploader(
  * uploaded, if {@code true} the blob is uploaded.
  * @throws IOException when reading of the {@link Chunker}s input source fails
  */
- public void uploadBlob(HashCode hash, Chunker chunker, boolean forceUpload)
+ public void uploadBlob(
+ RemoteActionExecutionContext context, HashCode hash, Chunker chunker, boolean forceUpload)
  throws IOException, InterruptedException {
- uploadBlobs(singletonMap(hash, chunker), forceUpload);
+ uploadBlobs(context, singletonMap(hash, chunker), forceUpload);
  }
 
  /**
@@ -156,12 +157,14 @@ public void uploadBlob(HashCode hash, Chunker chunker, boolean forceUpload)
  * uploaded, if {@code true} the blob is uploaded.
  * @throws IOException when reading of the {@link Chunker}s input source or uploading fails
  */
- public void uploadBlobs(Map<HashCode, Chunker> chunkers, boolean forceUpload)
+ public void uploadBlobs(
+ RemoteActionExecutionContext context, Map<HashCode, Chunker> chunkers, boolean forceUpload)
  throws IOException, InterruptedException {
  List<ListenableFuture<Void>> uploads = new ArrayList<>();
 
  for (Map.Entry<HashCode, Chunker> chunkerEntry : chunkers.entrySet()) {
- uploads.add(uploadBlobAsync(chunkerEntry.getKey(), chunkerEntry.getValue(), forceUpload));
+ uploads.add(
+ uploadBlobAsync(context, chunkerEntry.getKey(), chunkerEntry.getValue(), forceUpload));
  }
 
  try {
@@ -200,14 +203,17 @@ void shutdown() {
  }
  }
 
- /** @deprecated Use {@link #uploadBlobAsync(Digest, Chunker, boolean)} instead. */
+ /**
+ * @deprecated Use {@link #uploadBlobAsync(RemoteActionExecutionContext, Digest, Chunker,
+ * boolean)} instead.
+ */
  @Deprecated
  @VisibleForTesting
  public ListenableFuture<Void> uploadBlobAsync(
- HashCode hash, Chunker chunker, boolean forceUpload) {
+ RemoteActionExecutionContext context, HashCode hash, Chunker chunker, boolean forceUpload) {
  Digest digest =
  Digest.newBuilder().setHash(hash.toString()).setSizeBytes(chunker.getSize()).build();
- return uploadBlobAsync(digest, chunker, forceUpload);
+ return uploadBlobAsync(context, digest, chunker, forceUpload);
  }
 
  /**
@@ -227,7 +233,7 @@ public ListenableFuture<Void> uploadBlobAsync(
  * @throws IOException when reading of the {@link Chunker}s input source fails
  */
  public ListenableFuture<Void> uploadBlobAsync(
- Digest digest, Chunker chunker, boolean forceUpload) {
+ RemoteActionExecutionContext context, Digest digest, Chunker chunker, boolean forceUpload) {
  synchronized (lock) {
  checkState(!isShutdown, "Must not call uploadBlobs after shutdown.");
 
@@ -242,7 +248,7 @@ public ListenableFuture<Void> uploadBlobAsync(
 
  ListenableFuture<Void> uploadResult =
  Futures.transform(
- startAsyncUpload(digest, chunker),
+ startAsyncUpload(context, digest, chunker),
  (v) -> {
  synchronized (lock) {
  uploadedBlobs.add(HashCode.fromString(digest.getHash()));
@@ -294,7 +300,8 @@ private static String buildUploadResourceName(String instanceName, UUID uuid, Di
  }
 
  /** Starts a file upload and returns a future representing the upload. */
- private ListenableFuture<Void> startAsyncUpload(Digest digest, Chunker chunker) {
+ private ListenableFuture<Void> startAsyncUpload(
+ RemoteActionExecutionContext context, Digest digest, Chunker chunker) {
  try {
  chunker.reset();
  } catch (IOException e) {
@@ -313,7 +320,13 @@ private ListenableFuture<Void> startAsyncUpload(Digest digest, Chunker chunker)
  String resourceName = buildUploadResourceName(instanceName, uploadId, digest);
  AsyncUpload newUpload =
  new AsyncUpload(
- channel, callCredentialsProvider, callTimeoutSecs, retrier, resourceName, chunker);
+ context,
+ channel,
+ callCredentialsProvider,
+ callTimeoutSecs,
+ retrier,
+ resourceName,
+ chunker);
  ListenableFuture<Void> currUpload = newUpload.start();
  currUpload.addListener(
  () -> {
@@ -348,6 +361,7 @@ public ReferenceCounted touch(Object o) {
 
  private static class AsyncUpload {
 
+ private final RemoteActionExecutionContext context;
  private final Channel channel;
  private final CallCredentialsProvider callCredentialsProvider;
  private final long callTimeoutSecs;
@@ -358,12 +372,14 @@ private static class AsyncUpload {
  private ClientCall<WriteRequest, WriteResponse> call;
 
  AsyncUpload(
+ RemoteActionExecutionContext context,
  Channel channel,
  CallCredentialsProvider callCredentialsProvider,
  long callTimeoutSecs,
  Retrier retrier,
  String resourceName,
  Chunker chunker) {
+ this.context = context;
  this.channel = channel;
  this.callCredentialsProvider = callCredentialsProvider;
  this.callTimeoutSecs = callTimeoutSecs;
@@ -373,7 +389,6 @@ private static class AsyncUpload {
  }
 
  ListenableFuture<Void> start() {
- Context ctx = Context.current();
  ProgressiveBackoff progressiveBackoff = new ProgressiveBackoff(retrier::newBackoff);
  AtomicLong committedOffset = new AtomicLong(0);
 
@@ -383,8 +398,7 @@ ListenableFuture<Void> start() {
  retrier.executeAsync(
  () -> {
  if (committedOffset.get() < chunker.getSize()) {
- return ctx.call(
- () -> callAndQueryOnFailure(committedOffset, progressiveBackoff));
+ return callAndQueryOnFailure(committedOffset, progressiveBackoff);
  }
  return Futures.immediateFuture(null);
  },
@@ -409,7 +423,8 @@ ListenableFuture<Void> start() {
 
  private ByteStreamFutureStub bsFutureStub() {
  return ByteStreamGrpc.newFutureStub(channel)
- .withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
+ .withInterceptors(
+ TracingMetadataUtils.attachMetadataInterceptor(context.getRequestMetadata()))
  .withCallCredentials(callCredentialsProvider.getCallCredentials())
  .withDeadlineAfter(callTimeoutSecs, SECONDS);
  }
@@ -420,7 +435,7 @@ private ListenableFuture<Void> callAndQueryOnFailure(
  call(committedOffset),
  Exception.class,
  (e) -> guardQueryWithSuppression(e, committedOffset, progressiveBackoff),
- Context.current().fixedContextExecutor(MoreExecutors.directExecutor()));
+ MoreExecutors.directExecutor());
  }
 
  private ListenableFuture<Void> guardQueryWithSuppression(
@@ -584,7 +599,9 @@ public void onReady() {
  }
  }
  };
- call.start(callListener, TracingMetadataUtils.headersFromCurrentContext());
+ call.start(
+ callListener,
+ TracingMetadataUtils.headersFromRequestMetadata(context.getRequestMetadata()));
  call.request(1);
  return uploadResult;
  }

src/main/java/com/google/devtools/build/lib/remote/GrpcCacheClient.java

@@ -393,16 +393,22 @@ public void onCompleted() {
  }
 
  @Override
- public ListenableFuture<Void> uploadFile(Digest digest, Path path) {
+ public ListenableFuture<Void> uploadFile(
+ RemoteActionExecutionContext context, Digest digest, Path path) {
  return uploader.uploadBlobAsync(
+ context,
  digest,
  Chunker.builder().setInput(digest.getSizeBytes(), path).build(),
  /* forceUpload= */ true);
  }
 
  @Override
- public ListenableFuture<Void> uploadBlob(Digest digest, ByteString data) {
+ public ListenableFuture<Void> uploadBlob(
+ RemoteActionExecutionContext context, Digest digest, ByteString data) {
  return uploader.uploadBlobAsync(
- digest, Chunker.builder().setInput(data.toByteArray()).build(), /* forceUpload= */ true);
+ context,
+ digest,
+ Chunker.builder().setInput(data.toByteArray()).build(),
+ /* forceUpload= */ true);
  }
 }

src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java

@@ -139,7 +139,7 @@ public ActionResult upload(
  int exitCode)
  throws ExecException, IOException, InterruptedException {
  ActionResult.Builder resultBuilder = ActionResult.newBuilder();
- uploadOutputs(execRoot, actionKey, action, command, outputs, outErr, resultBuilder);
+ uploadOutputs(context, execRoot, actionKey, action, command, outputs, outErr, resultBuilder);
  resultBuilder.setExitCode(exitCode);
  ActionResult result = resultBuilder.build();
  if (exitCode == 0 && !action.getDoNotCache()) {
@@ -162,6 +162,7 @@ public ActionResult upload(
  }
 
  private void uploadOutputs(
+ RemoteActionExecutionContext context,
  Path execRoot,
  ActionKey actionKey,
  Action action,
@@ -192,14 +193,14 @@ private void uploadOutputs(
  for (Digest digest : digestsToUpload) {
  Path file = digestToFile.get(digest);
  if (file != null) {
- uploads.add(cacheProtocol.uploadFile(digest, file));
+ uploads.add(cacheProtocol.uploadFile(context, digest, file));
  } else {
  ByteString blob = digestToBlobs.get(digest);
  if (blob == null) {
  String message = "FindMissingBlobs call returned an unknown digest: " + digest;
  throw new IOException(message);
  }
- uploads.add(cacheProtocol.uploadBlob(digest, blob));
+ uploads.add(cacheProtocol.uploadBlob(context, digest, blob));
  }
  }