Skip to content

Commit

Permalink
Remote: Use parameters instead of thread-local storage to provide tra…
Browse files Browse the repository at this point in the history
…cing metadata. (Part 2)

Change RemoteCacheClient#uploadActionResult to use RemoteActionExecutionContext.

PiperOrigin-RevId: 354233348
  • Loading branch information
Googler authored and copybara-github committed Jan 28, 2021
1 parent f0b0c39 commit 92955e6
Show file tree
Hide file tree
Showing 16 changed files with 143 additions and 34 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -556,6 +572,7 @@ private SpawnResult execLocally(Spawn spawn, SpawnExecutionContext context)
}

private SpawnResult execLocallyAndUploadOrFail(
RemoteActionExecutionContext remoteActionExecutionContext,
Spawn spawn,
SpawnExecutionContext context,
SortedMap<PathFragment, ActionInput> inputMap,
Expand All @@ -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);
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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);
}
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -826,6 +827,7 @@ public void updateActionResult(

ActionResult result =
remoteCache.upload(
remoteActionExecutionContext,
DIGEST_UTIL.asActionKey(actionDigest),
action,
command,
Expand Down Expand Up @@ -888,6 +890,7 @@ public void updateActionResult(

ActionResult result =
remoteCache.upload(
remoteActionExecutionContext,
DIGEST_UTIL.asActionKey(actionDigest),
action,
command,
Expand Down Expand Up @@ -1036,6 +1039,7 @@ public void onError(Throwable t) {
.when(mockByteStreamImpl)
.queryWriteStatus(any(), any());
remoteCache.upload(
remoteActionExecutionContext,
actionKey,
Action.getDefaultInstance(),
Command.getDefaultInstance(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -99,6 +103,7 @@ public class RemoteCacheTests {

@Mock private OutputFilesLocker outputFilesLocker;

private RemoteActionExecutionContext remoteActionExecutionContext;
private FileSystem fs;
private Path execRoot;
ArtifactRoot artifactRoot;
Expand All @@ -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();
Expand Down Expand Up @@ -1426,6 +1436,7 @@ public void testUploadDirectory() throws Exception {
InMemoryRemoteCache remoteCache = newRemoteCache();
ActionResult result =
remoteCache.upload(
remoteActionExecutionContext,
digestUtil.asActionKey(actionDigest),
action,
cmd,
Expand Down Expand Up @@ -1462,6 +1473,7 @@ public void testUploadEmptyDirectory() throws Exception {
InMemoryRemoteCache remoteCache = newRemoteCache();
ActionResult result =
remoteCache.upload(
remoteActionExecutionContext,
actionDigest,
action,
cmd,
Expand Down Expand Up @@ -1518,6 +1530,7 @@ public void testUploadNestedDirectory() throws Exception {
InMemoryRemoteCache remoteCache = newRemoteCache();
ActionResult result =
remoteCache.upload(
remoteActionExecutionContext,
actionDigest,
action,
cmd,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -309,6 +310,7 @@ public Void answer(InvocationOnMock invocation) {
})
.when(remoteCache)
.upload(
any(RemoteActionExecutionContext.class),
any(ActionKey.class),
any(Action.class),
any(Command.class),
Expand All @@ -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),
Expand Down Expand Up @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand Down Expand Up @@ -566,6 +572,7 @@ public Void answer(InvocationOnMock invocation) {
})
.when(remoteCache)
.upload(
any(RemoteActionExecutionContext.class),
any(ActionKey.class),
any(Action.class),
any(Command.class),
Expand All @@ -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),
Expand Down Expand Up @@ -637,6 +645,7 @@ public Void answer(InvocationOnMock invocation) {
})
.when(remoteCache)
.upload(
any(RemoteActionExecutionContext.class),
any(ActionKey.class),
any(Action.class),
any(Command.class),
Expand All @@ -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),
Expand Down
Loading

0 comments on commit 92955e6

Please sign in to comment.