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

…cing metadata. (Part 5)

Change MissingDigestsFinder#findMissingDigests and RemoteExecutionClient#executeRemotely to use RemoteActionExecutionContext.

Removed all the usages of io.grpc.Context in the client code.

Fixed the regression about NetworkTime introduced by bazelbuild@bc54c64.

PiperOrigin-RevId: 354479787

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

@@ -30,13 +30,10 @@
 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;
 import io.netty.util.ReferenceCounted;
 import java.io.IOException;
@@ -161,7 +158,9 @@ private static List<PathMetadata> processQueryResult(
  */
  private ListenableFuture<ImmutableIterable<PathMetadata>> queryRemoteCache(
  ImmutableList<ListenableFuture<PathMetadata>> allPaths) throws Exception {
- Context ctx = TracingMetadataUtils.contextWithMetadata(buildRequestId, commandId, "bes-upload");
+ RequestMetadata metadata =
+ TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "bes-upload");
+ RemoteActionExecutionContext context = RemoteActionExecutionContext.create(metadata);
 
  List<PathMetadata> knownRemotePaths = new ArrayList<>(allPaths.size());
  List<PathMetadata> filesToQuery = new ArrayList<>();
@@ -181,7 +180,7 @@ private ListenableFuture<ImmutableIterable<PathMetadata>> queryRemoteCache(
  return Futures.immediateFuture(ImmutableIterable.from(knownRemotePaths));
  }
  return Futures.transform(
- ctx.call(() -> missingDigestsFinder.findMissingDigests(digestsToQuery)),
+ missingDigestsFinder.findMissingDigests(context, digestsToQuery),
  (missingDigests) -> {
  List<PathMetadata> filesToQueryUpdated = processQueryResult(missingDigests, filesToQuery);
  return ImmutableIterable.from(Iterables.concat(knownRemotePaths, filesToQueryUpdated));
@@ -197,8 +196,7 @@ private ListenableFuture<List<PathMetadata>> uploadLocalFiles(
  ImmutableIterable<PathMetadata> allPaths) {
  RequestMetadata metadata =
  TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "bes-upload");
- RemoteActionExecutionContext context =
- new RemoteActionExecutionContextImpl(metadata, new NetworkTime());
+ RemoteActionExecutionContext context = RemoteActionExecutionContext.create(metadata);
 
  ImmutableList.Builder<ListenableFuture<PathMetadata>> allPathsUploaded =
  ImmutableList.builder();

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

@@ -208,7 +208,6 @@ void shutdown() {
  * boolean)} instead.
  */
  @Deprecated
- @VisibleForTesting
  public ListenableFuture<Void> uploadBlobAsync(
  RemoteActionExecutionContext context, HashCode hash, Chunker chunker, boolean forceUpload) {
  Digest digest =

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

@@ -19,13 +19,15 @@
 import build.bazel.remote.execution.v2.ExecuteResponse;
 import build.bazel.remote.execution.v2.ExecutionGrpc;
 import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import build.bazel.remote.execution.v2.WaitExecutionRequest;
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.remote.RemoteRetrier.ProgressiveBackoff;
 import com.google.devtools.build.lib.remote.Retrier.Backoff;
 import com.google.devtools.build.lib.remote.common.OperationObserver;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
 import com.google.devtools.build.lib.remote.common.RemoteExecutionClient;
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
 import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
@@ -71,9 +73,9 @@ public ExperimentalGrpcRemoteExecutor(
  this.retrier = retrier;
  }
 
- private ExecutionBlockingStub executionBlockingStub() {
+ private ExecutionBlockingStub executionBlockingStub(RequestMetadata metadata) {
  return ExecutionGrpc.newBlockingStub(channel)
- .withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
+ .withInterceptors(TracingMetadataUtils.attachMetadataInterceptor(metadata))
  .withCallCredentials(callCredentialsProvider.getCallCredentials())
  .withDeadlineAfter(remoteOptions.remoteTimeout.getSeconds(), SECONDS);
  }
@@ -310,11 +312,16 @@ static ExecuteResponse extractResponseOrThrowIfError(Operation operation) throws
  }
 
  @Override
- public ExecuteResponse executeRemotely(ExecuteRequest request, OperationObserver observer)
+ public ExecuteResponse executeRemotely(
+ RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer)
  throws IOException, InterruptedException {
  Execution execution =
  new Execution(
- request, observer, retrier, callCredentialsProvider, this::executionBlockingStub);
+ request,
+ observer,
+ retrier,
+ callCredentialsProvider,
+ () -> this.executionBlockingStub(context.getRequestMetadata()));
  return execution.start();
  }
 

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

@@ -56,7 +56,6 @@
 import com.google.devtools.build.lib.remote.util.Utils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.protobuf.ByteString;
-import io.grpc.Context;
 import io.grpc.Status;
 import io.grpc.Status.Code;
 import io.grpc.StatusRuntimeException;
@@ -122,9 +121,11 @@ private int computeMaxMissingBlobsDigestsPerMessage() {
  return (options.maxOutboundMessageSize - overhead) / digestSize;
  }
 
- private ContentAddressableStorageFutureStub casFutureStub() {
+ private ContentAddressableStorageFutureStub casFutureStub(RemoteActionExecutionContext context) {
  return ContentAddressableStorageGrpc.newFutureStub(channel)
- .withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
+ .withInterceptors(
+ TracingMetadataUtils.attachMetadataInterceptor(context.getRequestMetadata()),
+ new NetworkTimeInterceptor(context::getNetworkTime))
  .withCallCredentials(callCredentialsProvider.getCallCredentials())
  .withDeadlineAfter(options.remoteTimeout.getSeconds(), TimeUnit.SECONDS);
  }
@@ -176,7 +177,8 @@ public static boolean isRemoteCacheOptions(RemoteOptions options) {
  }
 
  @Override
- public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(Iterable<Digest> digests) {
+ public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(
+ RemoteActionExecutionContext context, Iterable<Digest> digests) {
  if (Iterables.isEmpty(digests)) {
  return Futures.immediateFuture(ImmutableSet.of());
  }
@@ -187,13 +189,13 @@ public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(Iterable<Digest
  for (Digest digest : digests) {
  requestBuilder.addBlobDigests(digest);
  if (requestBuilder.getBlobDigestsCount() == maxMissingBlobsDigestsPerMessage) {
- getMissingDigestCalls.add(getMissingDigests(requestBuilder.build()));
+ getMissingDigestCalls.add(getMissingDigests(context, requestBuilder.build()));
  requestBuilder.clearBlobDigests();
  }
  }
 
  if (requestBuilder.getBlobDigestsCount() > 0) {
- getMissingDigestCalls.add(getMissingDigests(requestBuilder.build()));
+ getMissingDigestCalls.add(getMissingDigests(context, requestBuilder.build()));
  }
 
  ListenableFuture<ImmutableSet<Digest>> success =
@@ -209,7 +211,7 @@ public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(Iterable<Digest
  },
  MoreExecutors.directExecutor());
 
- RequestMetadata requestMetadata = TracingMetadataUtils.fromCurrentContext();
+ RequestMetadata requestMetadata = context.getRequestMetadata();
  return Futures.catchingAsync(
  success,
  RuntimeException.class,
@@ -226,10 +228,9 @@ public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(Iterable<Digest
  }
 
  private ListenableFuture<FindMissingBlobsResponse> getMissingDigests(
- FindMissingBlobsRequest request) {
- Context ctx = Context.current();
+ RemoteActionExecutionContext context, FindMissingBlobsRequest request) {
  return Utils.refreshIfUnauthenticatedAsync(
- () -> retrier.executeAsync(() -> ctx.call(() -> casFutureStub().findMissingBlobs(request))),
+ () -> retrier.executeAsync(() -> casFutureStub(context).findMissingBlobs(request)),
  callCredentialsProvider);
  }
 

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

@@ -18,11 +18,13 @@
 import build.bazel.remote.execution.v2.ExecuteResponse;
 import build.bazel.remote.execution.v2.ExecutionGrpc;
 import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import build.bazel.remote.execution.v2.WaitExecutionRequest;
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.remote.common.OperationObserver;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
 import com.google.devtools.build.lib.remote.common.RemoteExecutionClient;
 import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.remote.util.Utils;
@@ -55,9 +57,9 @@ public GrpcRemoteExecutor(
  this.retrier = retrier;
  }
 
- private ExecutionBlockingStub execBlockingStub() {
+ private ExecutionBlockingStub execBlockingStub(RequestMetadata metadata) {
  return ExecutionGrpc.newBlockingStub(channel)
- .withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
+ .withInterceptors(TracingMetadataUtils.attachMetadataInterceptor(metadata))
  .withCallCredentials(callCredentialsProvider.getCallCredentials());
  }
 
@@ -105,7 +107,8 @@ private ExecuteResponse getOperationResponse(Operation op) throws IOException {
  * trigger a retry of the Execute call, resulting in a new Operation.
  * */
  @Override
- public ExecuteResponse executeRemotely(ExecuteRequest request, OperationObserver observer)
+ public ExecuteResponse executeRemotely(
+ RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer)
  throws IOException, InterruptedException {
  // Execute has two components: the Execute call and (optionally) the WaitExecution call.
  // This is the simple flow without any errors:
@@ -149,9 +152,9 @@ public ExecuteResponse executeRemotely(ExecuteRequest request, OperationObserver
  WaitExecutionRequest.newBuilder()
  .setName(operation.get().getName())
  .build();
- replies = execBlockingStub().waitExecution(wr);
+ replies = execBlockingStub(context.getRequestMetadata()).waitExecution(wr);
  } else {
- replies = execBlockingStub().execute(request);
+ replies = execBlockingStub(context.getRequestMetadata()).execute(request);
  }
  try {
  while (replies.hasNext()) {

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

@@ -19,7 +19,6 @@
 import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.ClientInterceptor;
-import io.grpc.Context;
 import io.grpc.ForwardingClientCall;
 import io.grpc.ForwardingClientCallListener;
 import io.grpc.Metadata;
@@ -30,7 +29,6 @@
 /** The ClientInterceptor used to track network time. */
 public class NetworkTimeInterceptor implements ClientInterceptor {
 
- public static final Context.Key<NetworkTime> CONTEXT_KEY = Context.key("remote-network-time");
  private final Supplier<NetworkTime> networkTimeSupplier;
 
  public NetworkTimeInterceptor(Supplier<NetworkTime> networkTimeSupplier) {

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

@@ -33,15 +33,12 @@
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
-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.remote.util.Utils;
 import com.google.devtools.build.lib.sandbox.SandboxHelpers;
 import com.google.devtools.build.lib.vfs.Path;
-import io.grpc.Context;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -167,48 +164,42 @@ private ListenableFuture<Void> downloadFileAsync(Path path, FileArtifactValue me
  if (download == null) {
  RequestMetadata requestMetadata =
  TracingMetadataUtils.buildMetadata(buildRequestId, commandId, metadata.getActionId());
- RemoteActionExecutionContext remoteActionExecutionContext =
- new RemoteActionExecutionContextImpl(requestMetadata, new NetworkTime());
- Context ctx = TracingMetadataUtils.contextWithMetadata(requestMetadata);
- Context prevCtx = ctx.attach();
- try {
- Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());
- download = remoteCache.downloadFile(remoteActionExecutionContext, path, digest);
- downloadsInProgress.put(path, download);
- Futures.addCallback(
- download,
- new FutureCallback<Void>() {
- @Override
- public void onSuccess(Void v) {
- synchronized (lock) {
- downloadsInProgress.remove(path);
- downloadedPaths.add(path);
- }
-
- try {
- path.chmod(0755);
- } catch (IOException e) {
- logger.atWarning().withCause(e).log("Failed to chmod 755 on %s", path);
- }
+ RemoteActionExecutionContext context = RemoteActionExecutionContext.create(requestMetadata);
+
+ Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());
+ download = remoteCache.downloadFile(context, path, digest);
+ downloadsInProgress.put(path, download);
+ Futures.addCallback(
+ download,
+ new FutureCallback<Void>() {
+ @Override
+ public void onSuccess(Void v) {
+ synchronized (lock) {
+ downloadsInProgress.remove(path);
+ downloadedPaths.add(path);
  }
 
- @Override
- public void onFailure(Throwable t) {
- synchronized (lock) {
- downloadsInProgress.remove(path);
- }
- try {
- path.delete();
- } catch (IOException e) {
- logger.atWarning().withCause(e).log(
- "Failed to delete output file after incomplete download: %s", path);
- }
+ try {
+ path.chmod(0755);
+ } catch (IOException e) {
+ logger.atWarning().withCause(e).log("Failed to chmod 755 on %s", path);
  }
- },
- MoreExecutors.directExecutor());
- } finally {
- ctx.detach(prevCtx);
- }
+ }
+
+ @Override
+ public void onFailure(Throwable t) {
+ synchronized (lock) {
+ downloadsInProgress.remove(path);
+ }
+ try {
+ path.delete();
+ } catch (IOException e) {
+ logger.atWarning().withCause(e).log(
+ "Failed to delete output file after incomplete download: %s", path);
+ }
+ }
+ },
+ MoreExecutors.directExecutor());
  }
  return download;
  }

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

@@ -188,7 +188,8 @@ private void uploadOutputs(
  digests.addAll(digestToFile.keySet());
  digests.addAll(digestToBlobs.keySet());
 
- ImmutableSet<Digest> digestsToUpload = getFromFuture(cacheProtocol.findMissingDigests(digests));
+ ImmutableSet<Digest> digestsToUpload =
+ getFromFuture(cacheProtocol.findMissingDigests(context, digests));
  ImmutableList.Builder<ListenableFuture<Void>> uploads = ImmutableList.builder();
  for (Digest digest : digestsToUpload) {
  Path file = digestToFile.get(digest);

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

@@ -62,7 +62,7 @@ public void ensureInputsPresent(
  Iterable<Digest> allDigests =
  Iterables.concat(merkleTree.getAllDigests(), additionalInputs.keySet());
  ImmutableSet<Digest> missingDigests =
- getFromFuture(cacheProtocol.findMissingDigests(allDigests));
+ getFromFuture(cacheProtocol.findMissingDigests(context, allDigests));
 
  List<ListenableFuture<Void>> uploadFutures = new ArrayList<>();
  for (Digest missingDigest : missingDigests) {