From fb1c5491398bbdac181e867022881fe2ff73c884 Mon Sep 17 00:00:00 2001
From: Inigo Goiri <inigoiri@apache.org>
Date: Mon, 22 Apr 2019 13:31:20 -0700
Subject: [PATCH] HDFS-14374. Expose total number of delegation tokens in
 AbstractDelegationTokenSecretManager. Contributed by CR Hota.

---
 .../AbstractDelegationTokenSecretManager.java |  9 +++++++-
 .../token/delegation/TestDelegationToken.java | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/AbstractDelegationTokenSecretManager.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/AbstractDelegationTokenSecretManager.java
index 9259d877e1..f329accec7 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/AbstractDelegationTokenSecretManager.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/AbstractDelegationTokenSecretManager.java
@@ -145,7 +145,14 @@ public synchronized void reset() {
     setDelegationTokenSeqNum(0);
     currentTokens.clear();
   }
-  
+
+  /**
+   * Total count of active delegation tokens.
+   */
+  public long getCurrentTokensSize() {
+    return currentTokens.size();
+  }
+
   /** 
    * Add a previously used master key to cache (when NN restarts), 
    * should be called before activate(). 
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/delegation/TestDelegationToken.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/delegation/TestDelegationToken.java
index ad12f0baba..df685cf681 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/delegation/TestDelegationToken.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/delegation/TestDelegationToken.java
@@ -259,6 +259,29 @@ public void testGetUserWithOwnerAndReal() {
                  ugi.getRealUser().getAuthenticationMethod());
   }
 
+  @Test
+  public void testDelegationTokenCount() throws Exception {
+    final TestDelegationTokenSecretManager dtSecretManager =
+        new TestDelegationTokenSecretManager(24*60*60*1000,
+            3*1000, 1*1000, 3600000);
+    try {
+      dtSecretManager.startThreads();
+      Assert.assertEquals(dtSecretManager.getCurrentTokensSize(), 0);
+      final Token<TestDelegationTokenIdentifier> token1 =
+          generateDelegationToken(dtSecretManager, "SomeUser", "JobTracker");
+      Assert.assertEquals(dtSecretManager.getCurrentTokensSize(), 1);
+      final Token<TestDelegationTokenIdentifier> token2 =
+          generateDelegationToken(dtSecretManager, "SomeUser", "JobTracker");
+      Assert.assertEquals(dtSecretManager.getCurrentTokensSize(), 2);
+      dtSecretManager.cancelToken(token1, "JobTracker");
+      Assert.assertEquals(dtSecretManager.getCurrentTokensSize(), 1);
+      dtSecretManager.cancelToken(token2, "JobTracker");
+      Assert.assertEquals(dtSecretManager.getCurrentTokensSize(), 0);
+    } finally {
+      dtSecretManager.stopThreads();
+    }
+  }
+
   @Test
   public void testDelegationTokenSecretManager() throws Exception {
     final TestDelegationTokenSecretManager dtSecretManager =