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

…cing metadata. (Part 2)

Change RemoteCacheClient#uploadActionResult to use RemoteActionExecutionContext.

PiperOrigin-RevId: 354233348

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

@@ -136,9 +136,11 @@ private ByteStreamStub bsAsyncStub() {
         .withDeadlineAfter(options.remoteTimeout.getSeconds(), TimeUnit.SECONDS);
   }
 
-  private ActionCacheBlockingStub acBlockingStub() {
+  private ActionCacheBlockingStub acBlockingStub(RemoteActionExecutionContext context) {
     return ActionCacheGrpc.newBlockingStub(channel)
-        .withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
+        .withInterceptors(
+            TracingMetadataUtils.attachMetadataInterceptor(context.getRequestMetadata()),
+            new NetworkTimeInterceptor(context::getNetworkTime))
         .withCallCredentials(callCredentialsProvider.getCallCredentials())
         .withDeadlineAfter(options.remoteTimeout.getSeconds(), TimeUnit.SECONDS);
   }
@@ -259,14 +261,15 @@ public ListenableFuture<ActionResult> downloadActionResult(
   }
 
   @Override
-  public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
+  public void uploadActionResult(
+      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
       throws IOException, InterruptedException {
     try {
       Utils.refreshIfUnauthenticated(
           () ->
               retrier.execute(
                   () ->
-                      acBlockingStub()
+                      acBlockingStub(context)
                           .updateActionResult(
                               UpdateActionResultRequest.newBuilder()
                                   .setInstanceName(options.remoteInstanceName)

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

@@ -129,6 +129,7 @@ public ActionResult downloadActionResult(
    * @throws ExecException if uploading any of the action outputs is not supported
    */
   public ActionResult upload(
+      RemoteActionExecutionContext context,
       ActionKey actionKey,
       Action action,
       Command command,
@@ -142,20 +143,22 @@ public ActionResult upload(
     resultBuilder.setExitCode(exitCode);
     ActionResult result = resultBuilder.build();
     if (exitCode == 0 && !action.getDoNotCache()) {
-      cacheProtocol.uploadActionResult(actionKey, result);
+      cacheProtocol.uploadActionResult(context, actionKey, result);
     }
     return result;
   }
 
   public ActionResult upload(
+      RemoteActionExecutionContext context,
       ActionKey actionKey,
       Action action,
       Command command,
       Path execRoot,
       Collection<Path> outputs,
       FileOutErr outErr)
       throws ExecException, IOException, InterruptedException {
-    return upload(actionKey, action, command, execRoot, outputs, outErr, /* exitCode= */ 0);
+    return upload(
+        context, actionKey, action, command, execRoot, outputs, outErr, /* exitCode= */ 0);
   }
 
   private void uploadOutputs(

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

@@ -291,7 +291,13 @@ public void store(SpawnResult result) throws ExecException, InterruptedException
               RemoteSpawnRunner.resolveActionInputs(execRoot, spawn.getOutputFiles());
           try (SilentCloseable c = prof.profile(ProfilerTask.UPLOAD_TIME, "upload outputs")) {
             remoteCache.upload(
-                actionKey, action, command, execRoot, files, context.getFileOutErr());
+                remoteActionExecutionContext,
+                actionKey,
+                action,
+                command,
+                execRoot,
+                files,
+                context.getFileOutErr());
           } catch (IOException e) {
             String errorMessage;
             if (!verboseFailures) {

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

@@ -298,7 +298,15 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
         }
       } catch (IOException e) {
         return execLocallyAndUploadOrFail(
-            spawn, context, inputMap, actionKey, action, command, uploadLocalResults, e);
+            remoteActionExecutionContext,
+            spawn,
+            context,
+            inputMap,
+            actionKey,
+            action,
+            command,
+            uploadLocalResults,
+            e);
       }
 
       ExecuteRequest.Builder requestBuilder =
@@ -383,7 +391,15 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
             });
       } catch (IOException e) {
         return execLocallyAndUploadOrFail(
-            spawn, context, inputMap, actionKey, action, command, uploadLocalResults, e);
+            remoteActionExecutionContext,
+            spawn,
+            context,
+            inputMap,
+            actionKey,
+            action,
+            command,
+            uploadLocalResults,
+            e);
       }
     } finally {
       withMetadata.detach(previous);
@@ -556,6 +572,7 @@ private SpawnResult execLocally(Spawn spawn, SpawnExecutionContext context)
   }
 
   private SpawnResult execLocallyAndUploadOrFail(
+      RemoteActionExecutionContext remoteActionExecutionContext,
       Spawn spawn,
       SpawnExecutionContext context,
       SortedMap<PathFragment, ActionInput> inputMap,
@@ -572,7 +589,14 @@ private SpawnResult execLocallyAndUploadOrFail(
     }
     if (remoteOptions.remoteLocalFallback && !RemoteRetrierUtils.causedByExecTimeout(cause)) {
       return execLocallyAndUpload(
-          spawn, context, inputMap, actionKey, action, command, uploadLocalResults);
+          remoteActionExecutionContext,
+          spawn,
+          context,
+          inputMap,
+          actionKey,
+          action,
+          command,
+          uploadLocalResults);
     }
     return handleError(cause, context.getFileOutErr(), actionKey, context);
   }
@@ -724,6 +748,7 @@ private Map<Path, Long> getInputCtimes(SortedMap<PathFragment, ActionInput> inpu
 
   @VisibleForTesting
   SpawnResult execLocallyAndUpload(
+      RemoteActionExecutionContext remoteActionExecutionContext,
       Spawn spawn,
       SpawnExecutionContext context,
       SortedMap<PathFragment, ActionInput> inputMap,
@@ -751,7 +776,13 @@ SpawnResult execLocallyAndUpload(
     Collection<Path> outputFiles = resolveActionInputs(execRoot, spawn.getOutputFiles());
     try (SilentCloseable c = Profiler.instance().profile(UPLOAD_TIME, "upload outputs")) {
       remoteCache.upload(
-          actionKey, action, command, execRoot, outputFiles, context.getFileOutErr());
+          remoteActionExecutionContext,
+          actionKey,
+          action,
+          command,
+          execRoot,
+          outputFiles,
+          context.getFileOutErr());
     } catch (IOException e) {
       if (verboseFailures) {
         report(Event.debug("Upload to remote cache failed: " + e.getMessage()));

src/main/java/com/google/devtools/build/lib/remote/common/RemoteCacheClient.java

@@ -83,7 +83,8 @@ ListenableFuture<ActionResult> downloadActionResult(
    * @throws IOException If there is an error uploading the action result.
    * @throws InterruptedException In case the thread
    */
-  void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
+  void uploadActionResult(
+      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
       throws IOException, InterruptedException;
 
   /**

src/main/java/com/google/devtools/build/lib/remote/disk/DiskAndRemoteCacheClient.java

@@ -49,11 +49,12 @@ public DiskAndRemoteCacheClient(
   }
 
   @Override
-  public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
+  public void uploadActionResult(
+      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
       throws IOException, InterruptedException {
-    diskCache.uploadActionResult(actionKey, actionResult);
+    diskCache.uploadActionResult(context, actionKey, actionResult);
     if (!options.incompatibleRemoteResultsIgnoreDisk || options.remoteUploadLocalResults) {
-      remoteCache.uploadActionResult(actionKey, actionResult);
+      remoteCache.uploadActionResult(context, actionKey, actionResult);
     }
   }
 
@@ -177,7 +178,7 @@ public ListenableFuture<ActionResult> downloadActionResult(
             if (actionResult == null) {
               return Futures.immediateFuture(null);
             } else {
-              diskCache.uploadActionResult(actionKey, actionResult);
+              diskCache.uploadActionResult(context, actionKey, actionResult);
               return Futures.immediateFuture(actionResult);
             }
           },

src/main/java/com/google/devtools/build/lib/remote/disk/DiskCacheClient.java

@@ -108,7 +108,8 @@ public ListenableFuture<ActionResult> downloadActionResult(
   }
 
   @Override
-  public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
+  public void uploadActionResult(
+      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
       throws IOException {
     try (InputStream data = actionResult.toByteString().newInput()) {
       saveFile(actionKey.getDigest().getHash(), data, /* actionResult= */ true);

src/main/java/com/google/devtools/build/lib/remote/http/HttpCacheClient.java

@@ -705,7 +705,8 @@ private boolean reset(InputStream in) throws IOException {
   }
 
   @Override
-  public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
+  public void uploadActionResult(
+      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
       throws IOException, InterruptedException {
     ByteString serialized = actionResult.toByteString();
     ListenableFuture<Void> uploadFuture =

src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java

@@ -700,7 +700,8 @@ private ActionResult uploadDirectory(RemoteCache remoteCache, List<Path> outputs
     Action action = Action.getDefaultInstance();
     ActionKey actionKey = DIGEST_UTIL.computeActionKey(action);
     Command cmd = Command.getDefaultInstance();
-    return remoteCache.upload(actionKey, action, cmd, execRoot, outputs, outErr);
+    return remoteCache.upload(
+        remoteActionExecutionContext, actionKey, action, cmd, execRoot, outputs, outErr);
   }
 
   @Test
@@ -826,6 +827,7 @@ public void updateActionResult(
 
     ActionResult result =
         remoteCache.upload(
+            remoteActionExecutionContext,
             DIGEST_UTIL.asActionKey(actionDigest),
             action,
             command,
@@ -888,6 +890,7 @@ public void updateActionResult(
 
     ActionResult result =
         remoteCache.upload(
+            remoteActionExecutionContext,
             DIGEST_UTIL.asActionKey(actionDigest),
             action,
             command,
@@ -1036,6 +1039,7 @@ public void onError(Throwable t) {
         .when(mockByteStreamImpl)
         .queryWriteStatus(any(), any());
     remoteCache.upload(
+        remoteActionExecutionContext,
         actionKey,
         Action.getDefaultInstance(),
         Command.getDefaultInstance(),

src/test/java/com/google/devtools/build/lib/remote/RemoteCacheTests.java

@@ -55,10 +55,14 @@
 import com.google.devtools.build.lib.clock.JavaClock;
 import com.google.devtools.build.lib.remote.RemoteCache.OutputFilesLocker;
 import com.google.devtools.build.lib.remote.RemoteCache.UploadManifest;
+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.common.RemoteCacheClient.ActionKey;
 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;
+import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.remote.util.Utils;
 import com.google.devtools.build.lib.remote.util.Utils.InMemoryOutput;
 import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
@@ -99,6 +103,7 @@ public class RemoteCacheTests {
 
   @Mock private OutputFilesLocker outputFilesLocker;
 
+  private RemoteActionExecutionContext remoteActionExecutionContext;
   private FileSystem fs;
   private Path execRoot;
   ArtifactRoot artifactRoot;
@@ -110,6 +115,11 @@ public class RemoteCacheTests {
   @Before
   public void setUp() throws Exception {
     MockitoAnnotations.initMocks(this);
+    remoteActionExecutionContext =
+        new RemoteActionExecutionContextImpl(
+            TracingMetadataUtils.buildMetadata(
+                "none", "none", Digest.getDefaultInstance().getHash()),
+            new NetworkTime());
     fs = new InMemoryFileSystem(new JavaClock(), DigestHashFunction.SHA256);
     execRoot = fs.getPath("/execroot");
     execRoot.createDirectoryAndParents();
@@ -1426,6 +1436,7 @@ public void testUploadDirectory() throws Exception {
     InMemoryRemoteCache remoteCache = newRemoteCache();
     ActionResult result =
         remoteCache.upload(
+            remoteActionExecutionContext,
             digestUtil.asActionKey(actionDigest),
             action,
             cmd,
@@ -1462,6 +1473,7 @@ public void testUploadEmptyDirectory() throws Exception {
     InMemoryRemoteCache remoteCache = newRemoteCache();
     ActionResult result =
         remoteCache.upload(
+            remoteActionExecutionContext,
             actionDigest,
             action,
             cmd,
@@ -1518,6 +1530,7 @@ public void testUploadNestedDirectory() throws Exception {
     InMemoryRemoteCache remoteCache = newRemoteCache();
     ActionResult result =
         remoteCache.upload(
+            remoteActionExecutionContext,
             actionDigest,
             action,
             cmd,

src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnCacheTest.java

@@ -270,6 +270,7 @@ public Void answer(InvocationOnMock invocation) {
     verify(remoteCache).download(eq(actionResult), eq(execRoot), eq(outErr), any());
     verify(remoteCache, never())
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -309,6 +310,7 @@ public Void answer(InvocationOnMock invocation) {
             })
         .when(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -318,6 +320,7 @@ public Void answer(InvocationOnMock invocation) {
     entry.store(result);
     verify(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -485,6 +488,7 @@ public void failedActionsAreNotUploaded() throws Exception {
     entry.store(result);
     verify(remoteCache, never())
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -510,6 +514,7 @@ public void printWarningIfUploadFails() throws Exception {
     doThrow(new IOException("cache down"))
         .when(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -520,6 +525,7 @@ public void printWarningIfUploadFails() throws Exception {
     entry.store(result);
     verify(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -566,6 +572,7 @@ public Void answer(InvocationOnMock invocation) {
             })
         .when(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -575,6 +582,7 @@ public Void answer(InvocationOnMock invocation) {
     entry.store(result);
     verify(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -637,6 +645,7 @@ public Void answer(InvocationOnMock invocation) {
             })
         .when(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),
@@ -646,6 +655,7 @@ public Void answer(InvocationOnMock invocation) {
     entry.store(result);
     verify(remoteCache)
         .upload(
+            any(RemoteActionExecutionContext.class),
             any(ActionKey.class),
             any(Action.class),
             any(Command.class),