public void launchContainer(
Map<String, LocalResource> localResources, Credentials credentials) {
reset(eventHandler);
@SuppressWarnings("unchecked")
Token<JobTokenIdentifier> jobToken = mock(Token.class);
TokenCache.setSessionToken(jobToken, credentials);
amContainer.handle(
new AMContainerEventLaunchRequest(
containerID,
vertexID,
new ContainerContext(
localResources, credentials, new HashMap<String, String>(), "")));
}
@VisibleForTesting
public TaskRunnerCallable(
SubmitWorkRequestProto request,
QueryFragmentInfo fragmentInfo,
Configuration conf,
ExecutionContext executionContext,
Map<String, String> envMap,
Credentials credentials,
long memoryAvailable,
AMReporter amReporter,
ConfParams confParams,
LlapDaemonExecutorMetrics metrics,
KilledTaskHandler killedTaskHandler,
FragmentCompletionHandler fragmentCompleteHandler) {
this.request = request;
this.fragmentInfo = fragmentInfo;
this.conf = conf;
this.executionContext = executionContext;
this.envMap = envMap;
this.objectRegistry = new ObjectRegistryImpl();
this.credentials = credentials;
this.memoryAvailable = memoryAvailable;
this.confParams = confParams;
this.jobToken = TokenCache.getSessionToken(credentials);
this.taskSpec = Converters.getTaskSpecfromProto(request.getFragmentSpec());
this.amReporter = amReporter;
// Register with the AMReporter when the callable is setup. Unregister once it starts running.
if (jobToken != null) {
this.amReporter.registerTask(
request.getAmHost(),
request.getAmPort(),
request.getUser(),
jobToken,
null,
request.getFragmentSpec().getDagName());
}
this.metrics = metrics;
this.requestId = request.getFragmentSpec().getFragmentIdentifierString();
// TODO Change this to the queryId/Name when that's available.
this.queryId = request.getFragmentSpec().getDagName();
this.killedTaskHandler = killedTaskHandler;
this.fragmentCompletionHanler = fragmentCompleteHandler;
}
@SuppressWarnings("unchecked")
@Test
public void testCredentialsTransfer() {
WrappedContainerMultipleDAGs wc = new WrappedContainerMultipleDAGs();
TezDAGID dagID2 = TezDAGID.getInstance("800", 500, 2);
TezDAGID dagID3 = TezDAGID.getInstance("800", 500, 3);
TezVertexID vertexID2 = TezVertexID.getInstance(dagID2, 1);
TezVertexID vertexID3 = TezVertexID.getInstance(dagID3, 1);
TezTaskID taskID2 = TezTaskID.getInstance(vertexID2, 1);
TezTaskID taskID3 = TezTaskID.getInstance(vertexID3, 1);
TezTaskAttemptID attempt11 = TezTaskAttemptID.getInstance(wc.taskID, 200);
TezTaskAttemptID attempt12 = TezTaskAttemptID.getInstance(wc.taskID, 300);
TezTaskAttemptID attempt21 = TezTaskAttemptID.getInstance(taskID2, 200);
TezTaskAttemptID attempt22 = TezTaskAttemptID.getInstance(taskID2, 300);
TezTaskAttemptID attempt31 = TezTaskAttemptID.getInstance(taskID3, 200);
TezTaskAttemptID attempt32 = TezTaskAttemptID.getInstance(taskID3, 300);
Map<String, LocalResource> LRs = new HashMap<String, LocalResource>();
AMContainerTask fetchedTask = null;
Token<TokenIdentifier> amGenToken = mock(Token.class);
Token<TokenIdentifier> token1 = mock(Token.class);
Token<TokenIdentifier> token3 = mock(Token.class);
Credentials containerCredentials = new Credentials();
TokenCache.setSessionToken(amGenToken, containerCredentials);
Text token1Name = new Text("tokenDag1");
Text token3Name = new Text("tokenDag3");
Credentials dag1Credentials = new Credentials();
dag1Credentials.addToken(new Text(token1Name), token1);
Credentials dag3Credentials = new Credentials();
dag3Credentials.addToken(new Text(token3Name), token3);
wc.launchContainer(new HashMap<String, LocalResource>(), containerCredentials);
wc.containerLaunched();
wc.assignTaskAttempt(attempt11, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNotNull(fetchedTask.getCredentials());
assertNotNull(fetchedTask.getCredentials().getToken(token1Name));
wc.taskAttemptSucceeded(attempt11);
wc.assignTaskAttempt(attempt12, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt12);
// Move to running a second DAG, with no credentials.
wc.setNewDAGID(dagID2);
wc.assignTaskAttempt(attempt21, LRs, null);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt21);
wc.assignTaskAttempt(attempt22, LRs, null);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt22);
// Move to running a third DAG, with Credentials this time
wc.setNewDAGID(dagID3);
wc.assignTaskAttempt(attempt31, LRs, dag3Credentials);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNotNull(fetchedTask.getCredentials());
assertNotNull(fetchedTask.getCredentials().getToken(token3Name));
assertNull(fetchedTask.getCredentials().getToken(token1Name));
wc.taskAttemptSucceeded(attempt31);
wc.assignTaskAttempt(attempt32, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt32);
}