From 06c56ff79b4cdf82f733498d3edfa0b6e531a2db Mon Sep 17 00:00:00 2001
From: Varun Saxena <varunsaxena@apache.org>
Date: Wed, 13 Jul 2016 07:55:34 +0530
Subject: [PATCH] YARN-5353. ResourceManager can leak delegation tokens when
 they are shared across apps. (Jason Lowe via Varun Saxena).

---
 .../resourcemanager/security/DelegationTokenRenewer.java   | 7 +------
 .../security/TestDelegationTokenRenewer.java               | 3 +++
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/security/DelegationTokenRenewer.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/security/DelegationTokenRenewer.java
index 4177ee21034..dfbf33397d1 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/security/DelegationTokenRenewer.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/security/DelegationTokenRenewer.java
@@ -773,7 +773,7 @@ public void updateKeepAliveApplications(List<ApplicationId> appIds) {
 
   private void removeApplicationFromRenewal(ApplicationId applicationId) {
     rmContext.getSystemCredentialsForApps().remove(applicationId);
-    Set<DelegationTokenToRenew> tokens = appTokens.get(applicationId);
+    Set<DelegationTokenToRenew> tokens = appTokens.remove(applicationId);
 
     if (tokens != null && !tokens.isEmpty()) {
       synchronized (tokens) {
@@ -798,15 +798,10 @@ private void removeApplicationFromRenewal(ApplicationId applicationId) {
           // cancel the token
           cancelToken(dttr);
 
-          it.remove();
           allTokens.remove(dttr.token);
         }
       }
     }
-
-    if(tokens != null && tokens.isEmpty()) {
-      appTokens.remove(applicationId);
-    }
   }
 
   /**
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestDelegationTokenRenewer.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestDelegationTokenRenewer.java
index 4aee3c1ec91..5dfee89bc6f 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestDelegationTokenRenewer.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestDelegationTokenRenewer.java
@@ -1254,6 +1254,9 @@ public void testCancelWithMultipleAppSubmissions() throws Exception{
     Assert.assertTrue(dttr.referringAppIds.isEmpty());
     Assert.assertTrue(dttr.isTimerCancelled());
     Assert.assertTrue(Renewer.cancelled);
+
+    // make sure the token also has been removed from appTokens
+    Assert.assertFalse(renewer.getDelegationTokens().contains(token1));
   }
 
 }