Skip to content

Commit 92955e6

Browse files
Googlercopybara-github
Googler
authored and
copybara-github
committed
Remote: Use parameters instead of thread-local storage to provide tracing metadata. (Part 2)
Change RemoteCacheClient#uploadActionResult to use RemoteActionExecutionContext. PiperOrigin-RevId: 354233348
1 parent f0b0c39 commit 92955e6

16 files changed

Lines changed: 143 additions & 34 deletions

File tree

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,11 @@ private ByteStreamStub bsAsyncStub() {
136136
.withDeadlineAfter(options.remoteTimeout.getSeconds(), TimeUnit.SECONDS);
137137
}
138138

139-
private ActionCacheBlockingStub acBlockingStub() {
139+
private ActionCacheBlockingStub acBlockingStub(RemoteActionExecutionContext context) {
140140
return ActionCacheGrpc.newBlockingStub(channel)
141-
.withInterceptors(TracingMetadataUtils.attachMetadataFromContextInterceptor())
141+
.withInterceptors(
142+
TracingMetadataUtils.attachMetadataInterceptor(context.getRequestMetadata()),
143+
new NetworkTimeInterceptor(context::getNetworkTime))
142144
.withCallCredentials(callCredentialsProvider.getCallCredentials())
143145
.withDeadlineAfter(options.remoteTimeout.getSeconds(), TimeUnit.SECONDS);
144146
}
@@ -259,14 +261,15 @@ public ListenableFuture<ActionResult> downloadActionResult(
259261
}
260262

261263
@Override
262-
public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
264+
public void uploadActionResult(
265+
RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
263266
throws IOException, InterruptedException {
264267
try {
265268
Utils.refreshIfUnauthenticated(
266269
() ->
267270
retrier.execute(
268271
() ->
269-
acBlockingStub()
272+
acBlockingStub(context)
270273
.updateActionResult(
271274
UpdateActionResultRequest.newBuilder()
272275
.setInstanceName(options.remoteInstanceName)

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ public ActionResult downloadActionResult(
129129
* @throws ExecException if uploading any of the action outputs is not supported
130130
*/
131131
public ActionResult upload(
132+
RemoteActionExecutionContext context,
132133
ActionKey actionKey,
133134
Action action,
134135
Command command,
@@ -142,20 +143,22 @@ public ActionResult upload(
142143
resultBuilder.setExitCode(exitCode);
143144
ActionResult result = resultBuilder.build();
144145
if (exitCode == 0 && !action.getDoNotCache()) {
145-
cacheProtocol.uploadActionResult(actionKey, result);
146+
cacheProtocol.uploadActionResult(context, actionKey, result);
146147
}
147148
return result;
148149
}
149150

150151
public ActionResult upload(
152+
RemoteActionExecutionContext context,
151153
ActionKey actionKey,
152154
Action action,
153155
Command command,
154156
Path execRoot,
155157
Collection<Path> outputs,
156158
FileOutErr outErr)
157159
throws ExecException, IOException, InterruptedException {
158-
return upload(actionKey, action, command, execRoot, outputs, outErr, /* exitCode= */ 0);
160+
return upload(
161+
context, actionKey, action, command, execRoot, outputs, outErr, /* exitCode= */ 0);
159162
}
160163

161164
private void uploadOutputs(

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,13 @@ public void store(SpawnResult result) throws ExecException, InterruptedException
291291
RemoteSpawnRunner.resolveActionInputs(execRoot, spawn.getOutputFiles());
292292
try (SilentCloseable c = prof.profile(ProfilerTask.UPLOAD_TIME, "upload outputs")) {
293293
remoteCache.upload(
294-
actionKey, action, command, execRoot, files, context.getFileOutErr());
294+
remoteActionExecutionContext,
295+
actionKey,
296+
action,
297+
command,
298+
execRoot,
299+
files,
300+
context.getFileOutErr());
295301
} catch (IOException e) {
296302
String errorMessage;
297303
if (!verboseFailures) {

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

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,15 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
298298
}
299299
} catch (IOException e) {
300300
return execLocallyAndUploadOrFail(
301-
spawn, context, inputMap, actionKey, action, command, uploadLocalResults, e);
301+
remoteActionExecutionContext,
302+
spawn,
303+
context,
304+
inputMap,
305+
actionKey,
306+
action,
307+
command,
308+
uploadLocalResults,
309+
e);
302310
}
303311

304312
ExecuteRequest.Builder requestBuilder =
@@ -383,7 +391,15 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
383391
});
384392
} catch (IOException e) {
385393
return execLocallyAndUploadOrFail(
386-
spawn, context, inputMap, actionKey, action, command, uploadLocalResults, e);
394+
remoteActionExecutionContext,
395+
spawn,
396+
context,
397+
inputMap,
398+
actionKey,
399+
action,
400+
command,
401+
uploadLocalResults,
402+
e);
387403
}
388404
} finally {
389405
withMetadata.detach(previous);
@@ -556,6 +572,7 @@ private SpawnResult execLocally(Spawn spawn, SpawnExecutionContext context)
556572
}
557573

558574
private SpawnResult execLocallyAndUploadOrFail(
575+
RemoteActionExecutionContext remoteActionExecutionContext,
559576
Spawn spawn,
560577
SpawnExecutionContext context,
561578
SortedMap<PathFragment, ActionInput> inputMap,
@@ -572,7 +589,14 @@ private SpawnResult execLocallyAndUploadOrFail(
572589
}
573590
if (remoteOptions.remoteLocalFallback && !RemoteRetrierUtils.causedByExecTimeout(cause)) {
574591
return execLocallyAndUpload(
575-
spawn, context, inputMap, actionKey, action, command, uploadLocalResults);
592+
remoteActionExecutionContext,
593+
spawn,
594+
context,
595+
inputMap,
596+
actionKey,
597+
action,
598+
command,
599+
uploadLocalResults);
576600
}
577601
return handleError(cause, context.getFileOutErr(), actionKey, context);
578602
}
@@ -724,6 +748,7 @@ private Map<Path, Long> getInputCtimes(SortedMap<PathFragment, ActionInput> inpu
724748

725749
@VisibleForTesting
726750
SpawnResult execLocallyAndUpload(
751+
RemoteActionExecutionContext remoteActionExecutionContext,
727752
Spawn spawn,
728753
SpawnExecutionContext context,
729754
SortedMap<PathFragment, ActionInput> inputMap,
@@ -751,7 +776,13 @@ SpawnResult execLocallyAndUpload(
751776
Collection<Path> outputFiles = resolveActionInputs(execRoot, spawn.getOutputFiles());
752777
try (SilentCloseable c = Profiler.instance().profile(UPLOAD_TIME, "upload outputs")) {
753778
remoteCache.upload(
754-
actionKey, action, command, execRoot, outputFiles, context.getFileOutErr());
779+
remoteActionExecutionContext,
780+
actionKey,
781+
action,
782+
command,
783+
execRoot,
784+
outputFiles,
785+
context.getFileOutErr());
755786
} catch (IOException e) {
756787
if (verboseFailures) {
757788
report(Event.debug("Upload to remote cache failed: " + e.getMessage()));

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ ListenableFuture<ActionResult> downloadActionResult(
8383
* @throws IOException If there is an error uploading the action result.
8484
* @throws InterruptedException In case the thread
8585
*/
86-
void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
86+
void uploadActionResult(
87+
RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
8788
throws IOException, InterruptedException;
8889

8990
/**

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,12 @@ public DiskAndRemoteCacheClient(
4949
}
5050

5151
@Override
52-
public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
52+
public void uploadActionResult(
53+
RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
5354
throws IOException, InterruptedException {
54-
diskCache.uploadActionResult(actionKey, actionResult);
55+
diskCache.uploadActionResult(context, actionKey, actionResult);
5556
if (!options.incompatibleRemoteResultsIgnoreDisk || options.remoteUploadLocalResults) {
56-
remoteCache.uploadActionResult(actionKey, actionResult);
57+
remoteCache.uploadActionResult(context, actionKey, actionResult);
5758
}
5859
}
5960

@@ -177,7 +178,7 @@ public ListenableFuture<ActionResult> downloadActionResult(
177178
if (actionResult == null) {
178179
return Futures.immediateFuture(null);
179180
} else {
180-
diskCache.uploadActionResult(actionKey, actionResult);
181+
diskCache.uploadActionResult(context, actionKey, actionResult);
181182
return Futures.immediateFuture(actionResult);
182183
}
183184
},

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,8 @@ public ListenableFuture<ActionResult> downloadActionResult(
108108
}
109109

110110
@Override
111-
public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
111+
public void uploadActionResult(
112+
RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
112113
throws IOException {
113114
try (InputStream data = actionResult.toByteString().newInput()) {
114115
saveFile(actionKey.getDigest().getHash(), data, /* actionResult= */ true);

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -705,7 +705,8 @@ private boolean reset(InputStream in) throws IOException {
705705
}
706706

707707
@Override
708-
public void uploadActionResult(ActionKey actionKey, ActionResult actionResult)
708+
public void uploadActionResult(
709+
RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult)
709710
throws IOException, InterruptedException {
710711
ByteString serialized = actionResult.toByteString();
711712
ListenableFuture<Void> uploadFuture =

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -700,7 +700,8 @@ private ActionResult uploadDirectory(RemoteCache remoteCache, List<Path> outputs
700700
Action action = Action.getDefaultInstance();
701701
ActionKey actionKey = DIGEST_UTIL.computeActionKey(action);
702702
Command cmd = Command.getDefaultInstance();
703-
return remoteCache.upload(actionKey, action, cmd, execRoot, outputs, outErr);
703+
return remoteCache.upload(
704+
remoteActionExecutionContext, actionKey, action, cmd, execRoot, outputs, outErr);
704705
}
705706

706707
@Test
@@ -826,6 +827,7 @@ public void updateActionResult(
826827

827828
ActionResult result =
828829
remoteCache.upload(
830+
remoteActionExecutionContext,
829831
DIGEST_UTIL.asActionKey(actionDigest),
830832
action,
831833
command,
@@ -888,6 +890,7 @@ public void updateActionResult(
888890

889891
ActionResult result =
890892
remoteCache.upload(
893+
remoteActionExecutionContext,
891894
DIGEST_UTIL.asActionKey(actionDigest),
892895
action,
893896
command,
@@ -1036,6 +1039,7 @@ public void onError(Throwable t) {
10361039
.when(mockByteStreamImpl)
10371040
.queryWriteStatus(any(), any());
10381041
remoteCache.upload(
1042+
remoteActionExecutionContext,
10391043
actionKey,
10401044
Action.getDefaultInstance(),
10411045
Command.getDefaultInstance(),

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,14 @@
5555
import com.google.devtools.build.lib.clock.JavaClock;
5656
import com.google.devtools.build.lib.remote.RemoteCache.OutputFilesLocker;
5757
import com.google.devtools.build.lib.remote.RemoteCache.UploadManifest;
58+
import com.google.devtools.build.lib.remote.common.NetworkTime;
59+
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
60+
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContextImpl;
5861
import com.google.devtools.build.lib.remote.common.RemoteCacheClient.ActionKey;
5962
import com.google.devtools.build.lib.remote.options.RemoteOptions;
6063
import com.google.devtools.build.lib.remote.util.DigestUtil;
6164
import com.google.devtools.build.lib.remote.util.InMemoryCacheClient;
65+
import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
6266
import com.google.devtools.build.lib.remote.util.Utils;
6367
import com.google.devtools.build.lib.remote.util.Utils.InMemoryOutput;
6468
import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
@@ -99,6 +103,7 @@ public class RemoteCacheTests {
99103

100104
@Mock private OutputFilesLocker outputFilesLocker;
101105

106+
private RemoteActionExecutionContext remoteActionExecutionContext;
102107
private FileSystem fs;
103108
private Path execRoot;
104109
ArtifactRoot artifactRoot;
@@ -110,6 +115,11 @@ public class RemoteCacheTests {
110115
@Before
111116
public void setUp() throws Exception {
112117
MockitoAnnotations.initMocks(this);
118+
remoteActionExecutionContext =
119+
new RemoteActionExecutionContextImpl(
120+
TracingMetadataUtils.buildMetadata(
121+
"none", "none", Digest.getDefaultInstance().getHash()),
122+
new NetworkTime());
113123
fs = new InMemoryFileSystem(new JavaClock(), DigestHashFunction.SHA256);
114124
execRoot = fs.getPath("/execroot");
115125
execRoot.createDirectoryAndParents();
@@ -1426,6 +1436,7 @@ public void testUploadDirectory() throws Exception {
14261436
InMemoryRemoteCache remoteCache = newRemoteCache();
14271437
ActionResult result =
14281438
remoteCache.upload(
1439+
remoteActionExecutionContext,
14291440
digestUtil.asActionKey(actionDigest),
14301441
action,
14311442
cmd,
@@ -1462,6 +1473,7 @@ public void testUploadEmptyDirectory() throws Exception {
14621473
InMemoryRemoteCache remoteCache = newRemoteCache();
14631474
ActionResult result =
14641475
remoteCache.upload(
1476+
remoteActionExecutionContext,
14651477
actionDigest,
14661478
action,
14671479
cmd,
@@ -1518,6 +1530,7 @@ public void testUploadNestedDirectory() throws Exception {
15181530
InMemoryRemoteCache remoteCache = newRemoteCache();
15191531
ActionResult result =
15201532
remoteCache.upload(
1533+
remoteActionExecutionContext,
15211534
actionDigest,
15221535
action,
15231536
cmd,

0 commit comments

Comments
 (0)