From 4f85c9a73bbc9dc5a849d287aa1d92094e4960e3 Mon Sep 17 00:00:00 2001
From: Lei Yang <leyang@linkedin.com>
Date: Wed, 23 Mar 2022 12:47:32 -0700
Subject: [PATCH] HDFS-16518: Add shutdownhook to invalidate the KeyProviders
 in the cache

Fixes #4100
Signed-off-by: Owen O'Malley <oomalley@linkedin.com>
---
 .../apache/hadoop/hdfs/KeyProviderCache.java  | 25 +++++++++++++++++++
 .../hadoop/hdfs/TestKeyProviderCache.java     |  9 +++++++
 2 files changed, 34 insertions(+)

diff --git a/hadoop-hdfs-project/hadoop-hdfs-client/src/main/java/org/apache/hadoop/hdfs/KeyProviderCache.java b/hadoop-hdfs-project/hadoop-hdfs-client/src/main/java/org/apache/hadoop/hdfs/KeyProviderCache.java
index 29073a1b380..62902f2a7eb 100644
--- a/hadoop-hdfs-project/hadoop-hdfs-client/src/main/java/org/apache/hadoop/hdfs/KeyProviderCache.java
+++ b/hadoop-hdfs-project/hadoop-hdfs-client/src/main/java/org/apache/hadoop/hdfs/KeyProviderCache.java
@@ -26,6 +26,7 @@ import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.crypto.key.KeyProvider;
 import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.util.KMSUtil;
 
 import org.apache.hadoop.thirdparty.com.google.common.annotations.VisibleForTesting;
@@ -34,6 +35,7 @@ import org.apache.hadoop.thirdparty.com.google.common.cache.CacheBuilder;
 import org.apache.hadoop.thirdparty.com.google.common.cache.RemovalListener;
 import org.apache.hadoop.thirdparty.com.google.common.cache.RemovalNotification;
 
+import org.apache.hadoop.util.ShutdownHookManager;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -65,6 +67,9 @@ public class KeyProviderCache {
           }
         })
         .build();
+
+    ShutdownHookManager.get().addShutdownHook(new KeyProviderCacheFinalizer(),
+        SHUTDOWN_HOOK_PRIORITY);
   }
 
   public KeyProvider get(final Configuration conf,
@@ -85,6 +90,26 @@ public class KeyProviderCache {
     }
   }
 
+  public static final int SHUTDOWN_HOOK_PRIORITY = FileSystem.SHUTDOWN_HOOK_PRIORITY - 1;
+
+  private class KeyProviderCacheFinalizer implements Runnable {
+    @Override
+    public synchronized void run() {
+      invalidateCache();
+    }
+  }
+
+  /**
+   * Invalidate cache. KeyProviders in the cache will be closed by cache hook.
+   */
+  @VisibleForTesting
+  synchronized void invalidateCache() {
+    LOG.debug("Invalidating all cached KeyProviders.");
+    if (cache != null) {
+      cache.invalidateAll();
+    }
+  }
+
   private URI createKeyProviderURI(Configuration conf) {
     final String providerUriStr = conf.getTrimmed(
         CommonConfigurationKeysPublic.HADOOP_SECURITY_KEY_PROVIDER_PATH);
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestKeyProviderCache.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestKeyProviderCache.java
index 9fc6b389494..58011d7f15f 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestKeyProviderCache.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestKeyProviderCache.java
@@ -32,6 +32,8 @@ public class TestKeyProviderCache {
 
   public static class DummyKeyProvider extends KeyProvider {
 
+    public static int CLOSE_CALL_COUNT = 0;
+
     public DummyKeyProvider(Configuration conf) {
       super(conf);
     }
@@ -76,6 +78,10 @@ public class TestKeyProviderCache {
     public void flush() throws IOException {
     }
 
+    @Override
+    public void close() {
+      CLOSE_CALL_COUNT += 1;
+    }
   }
 
   public static class Factory extends KeyProviderFactory {
@@ -124,6 +130,9 @@ public class TestKeyProviderCache {
     Assert.assertFalse("Same KeyProviders returned !!",
         keyProvider1 == keyProvider4);
 
+    kpCache.invalidateCache();
+    Assert.assertEquals("Expected number of closing calls doesn't match",
+        3, DummyKeyProvider.CLOSE_CALL_COUNT);
   }
 
   private URI getKeyProviderUriFromConf(Configuration conf) {