diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java index f1a3b6e..def6da5 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java @@ -110,7 +110,8 @@ private AHSClient historyClient; private boolean historyServiceEnabled; protected TimelineClient timelineClient; - protected Text timelineService; + @VisibleForTesting + Text timelineService; protected boolean timelineServiceEnabled; private static final String ROOT = "root"; @@ -272,12 +273,6 @@ public YarnClientApplication createApplication() private void addTimelineDelegationToken( ContainerLaunchContext clc) throws YarnException, IOException { - org.apache.hadoop.security.token.Token timelineDelegationToken = - timelineClient.getDelegationToken( - UserGroupInformation.getCurrentUser().getUserName()); - if (timelineDelegationToken == null) { - return; - } Credentials credentials = new Credentials(); DataInputByteBuffer dibb = new DataInputByteBuffer(); ByteBuffer tokens = clc.getTokens(); @@ -290,11 +285,15 @@ private void addTimelineDelegationToken( // one more for (org.apache.hadoop.security.token.Token token : credentials .getAllTokens()) { - TokenIdentifier tokenIdentifier = token.decodeIdentifier(); - if (tokenIdentifier instanceof TimelineDelegationTokenIdentifier) { + if (token.getKind().equals(TimelineDelegationTokenIdentifier.KIND_NAME)) { return; } } + org.apache.hadoop.security.token.Token + timelineDelegationToken = getTimelineDelegationToken(); + if (timelineDelegationToken == null) { + return; + } credentials.addToken(timelineService, timelineDelegationToken); if (LOG.isDebugEnabled()) { LOG.debug("Add timline delegation token into credentials: " @@ -306,6 +305,13 @@ private void addTimelineDelegationToken( clc.setTokens(tokens); } + @VisibleForTesting + org.apache.hadoop.security.token.Token + getTimelineDelegationToken() throws IOException, YarnException { + return timelineClient.getDelegationToken( + UserGroupInformation.getCurrentUser().getUserName()); + } + @Private @VisibleForTesting protected boolean isSecurityEnabled() { diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClient.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClient.java index 8259893..3c1b1c1 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClient.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClient.java @@ -20,6 +20,8 @@ import static org.mockito.Matchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -39,6 +41,7 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.io.DataInputByteBuffer; import org.apache.hadoop.io.DataOutputBuffer; +import org.apache.hadoop.io.Text; import org.apache.hadoop.security.Credentials; import org.apache.hadoop.security.SecurityUtil; import org.apache.hadoop.security.UserGroupInformation; @@ -745,10 +748,13 @@ public void testAutomaticTimelineDelegationTokenLoading() Configuration conf = new YarnConfiguration(); conf.setBoolean(YarnConfiguration.TIMELINE_SERVICE_ENABLED, true); SecurityUtil.setAuthenticationMethod(AuthenticationMethod.KERBEROS, conf); + TimelineDelegationTokenIdentifier timelineDT = + new TimelineDelegationTokenIdentifier(); final Token dToken = - new Token(); + new Token( + timelineDT.getBytes(), new byte[0], timelineDT.getKind(), new Text()); // crate a mock client - YarnClientImpl client = new YarnClientImpl() { + YarnClientImpl client = spy(new YarnClientImpl() { @Override protected void serviceInit(Configuration conf) throws Exception { if (getConfig().getBoolean(YarnConfiguration.TIMELINE_SERVICE_ENABLED, @@ -784,34 +790,48 @@ public ApplicationReport getApplicationReport(ApplicationId appId) { public boolean isSecurityEnabled() { return true; } - }; + }); client.init(conf); client.start(); - ApplicationSubmissionContext context = - mock(ApplicationSubmissionContext.class); - ApplicationId applicationId = ApplicationId.newInstance(0, 1); - when(context.getApplicationId()).thenReturn(applicationId); - DataOutputBuffer dob = new DataOutputBuffer(); - Credentials credentials = new Credentials(); - credentials.writeTokenStorageToStream(dob); - ByteBuffer tokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength()); - ContainerLaunchContext clc = ContainerLaunchContext.newInstance( - null, null, null, null, tokens, null); - when(context.getAMContainerSpec()).thenReturn(clc); - client.submitApplication(context); - // Check whether token is added or not - credentials = new Credentials(); - DataInputByteBuffer dibb = new DataInputByteBuffer(); - tokens = clc.getTokens(); - if (tokens != null) { - dibb.reset(tokens); - credentials.readTokenStorageStream(dibb); - tokens.rewind(); + try { + // when i == 0, timeline DT already exists, no need to get one more + // when i == 1, timeline DT doesn't exist, need to get one more + for (int i = 0; i < 2; ++i) { + ApplicationSubmissionContext context = + mock(ApplicationSubmissionContext.class); + ApplicationId applicationId = ApplicationId.newInstance(0, i + 1); + when(context.getApplicationId()).thenReturn(applicationId); + DataOutputBuffer dob = new DataOutputBuffer(); + Credentials credentials = new Credentials(); + if (i == 0) { + credentials.addToken(client.timelineService, dToken); + } + credentials.writeTokenStorageToStream(dob); + ByteBuffer tokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength()); + ContainerLaunchContext clc = ContainerLaunchContext.newInstance( + null, null, null, null, tokens, null); + when(context.getAMContainerSpec()).thenReturn(clc); + client.submitApplication(context); + if (i == 0) { + // GetTimelineDelegationToken shouldn't be called + verify(client, never()).getTimelineDelegationToken(); + } + // In either way, token should be there + credentials = new Credentials(); + DataInputByteBuffer dibb = new DataInputByteBuffer(); + tokens = clc.getTokens(); + if (tokens != null) { + dibb.reset(tokens); + credentials.readTokenStorageStream(dibb); + tokens.rewind(); + } + Collection> dTokens = + credentials.getAllTokens(); + Assert.assertEquals(1, dTokens.size()); + Assert.assertEquals(dToken, dTokens.iterator().next()); + } + } finally { + client.stop(); } - Collection> dTokens = - credentials.getAllTokens(); - Assert.assertEquals(1, dTokens.size()); - Assert.assertEquals(dToken, dTokens.iterator().next()); - client.stop(); } }