[incubator-gobblin] branch master updated: [GOBBLIN-967] Change token refresh method in YarnContainerSecirityManager

[su...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

suvasude pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-gobblin.git


The following commit(s) were added to refs/heads/master by this push:
     new 6c099ff  [GOBBLIN-967] Change token refresh method in YarnContainerSecirityManager
6c099ff is described below

commit 6c099ff7c82c66ff683c7a7f3e96d3d03c13f9a6
Author: Zihan Li <zi...@zihli-mn1.linkedin.biz>
AuthorDate: Sat Nov 16 14:18:56 2019 -0800

    [GOBBLIN-967] Change token refresh method in YarnContainerSecirityManager
    
    Closes #2813 from ZihanLi58/GOBBLIN-967
---
 .../gobblin/yarn/YarnContainerSecurityManager.java   | 20 +++++++++-----------
 .../apache/gobblin/yarn/YarnSecurityManagerTest.java |  7 ++++---
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/YarnContainerSecurityManager.java b/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/YarnContainerSecurityManager.java
index 12f8d74..060da6a 100644
--- a/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/YarnContainerSecurityManager.java
+++ b/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/YarnContainerSecurityManager.java
@@ -18,10 +18,10 @@
 package org.apache.gobblin.yarn;
 
 import java.io.IOException;
-import java.util.Collection;
 
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.security.token.Token;
 import org.apache.hadoop.security.token.TokenIdentifier;
@@ -71,7 +71,7 @@ public class YarnContainerSecurityManager extends AbstractIdleService {
   @Subscribe
   public void handleTokenFileUpdatedEvent(DelegationTokenUpdatedEvent delegationTokenUpdatedEvent) {
     try {
-      addDelegationTokens(readDelegationTokens(this.tokenFilePath));
+      addCredentials(readCredentials(this.tokenFilePath));
     } catch (IOException ioe) {
       throw Throwables.propagate(ioe);
     }
@@ -91,18 +91,16 @@ public class YarnContainerSecurityManager extends AbstractIdleService {
    * Read the {@link Token}s stored in the token file.
    */
   @VisibleForTesting
-  Collection<Token<? extends TokenIdentifier>> readDelegationTokens(Path tokenFilePath) throws IOException {
-    LOGGER.info("Reading updated token from token file: " + tokenFilePath);
-    return YarnHelixUtils.readTokensFromFile(tokenFilePath, this.fs.getConf());
+  Credentials readCredentials(Path tokenFilePath) throws IOException {
+    LOGGER.info("Reading updated credentials from token file: " + tokenFilePath);
+    return Credentials.readTokenStorageFile(tokenFilePath, this.fs.getConf());
   }
 
   @VisibleForTesting
-  void addDelegationTokens(Collection<Token<? extends TokenIdentifier>> tokens) throws IOException {
-    for (Token<? extends TokenIdentifier> token : tokens) {
-      if (!UserGroupInformation.getCurrentUser().addToken(token)) {
-        LOGGER.error(String.format("Failed to add token %s to user %s",
-            token.toString(), UserGroupInformation.getLoginUser().getShortUserName()));
-      }
+  void addCredentials(Credentials credentials) throws IOException {
+    for (Token<? extends TokenIdentifier> token : credentials.getAllTokens()) {
+      LOGGER.info("updating "+token.toString());
     }
+    UserGroupInformation.getCurrentUser().addCredentials(credentials);
   }
 }
diff --git a/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnSecurityManagerTest.java b/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnSecurityManagerTest.java
index ee15a21..14a8d5b 100644
--- a/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnSecurityManagerTest.java
+++ b/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnSecurityManagerTest.java
@@ -29,6 +29,7 @@ import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.Text;
+import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.security.token.Token;
 import org.apache.helix.HelixManager;
@@ -183,9 +184,9 @@ public class YarnSecurityManagerTest {
 
   @Test(dependsOnMethods = "testWriteDelegationTokenToFile")
   public void testYarnContainerSecurityManager() throws IOException {
-    Collection<Token<?>> tokens = this.yarnContainerSecurityManager.readDelegationTokens(this.tokenFilePath);
-    assertToken(tokens);
-    this.yarnContainerSecurityManager.addDelegationTokens(tokens);
+    Credentials credentials = this.yarnContainerSecurityManager.readCredentials(this.tokenFilePath);
+    assertToken(credentials.getAllTokens());
+    this.yarnContainerSecurityManager.addCredentials(credentials);
     assertToken(UserGroupInformation.getCurrentUser().getTokens());
   }