From 6adf8462bafc27e3f4bf83948a974c089bd70682 Mon Sep 17 00:00:00 2001
From: Xiaoyu Yao <xyao@apache.org>
Date: Thu, 17 Sep 2020 10:39:19 -0700
Subject: [PATCH] HADOOP-17208. LoadBalanceKMSClientProvider#deleteKey should
 invalidateCache via all KMSClientProvider instances. (#2259)

---
 .../kms/LoadBalancingKMSClientProvider.java   |  1 +
 .../hadoop/crypto/key/kms/ValueQueue.java     | 21 +++++++++----------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
index ee2295cff7..818069423e 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
@@ -502,6 +502,7 @@ public Void call(KMSClientProvider provider) throws IOException {
         return null;
       }
     }, nextIdx(), false);
+    invalidateCache(name);
   }
 
   @Override
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/ValueQueue.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/ValueQueue.java
index 7d26acbf21..13aa0f7fd8 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/ValueQueue.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/ValueQueue.java
@@ -299,19 +299,18 @@ public E getNext(String keyName)
    * @param keyName the key to drain the Queue for
    */
   public void drain(String keyName) {
+    Runnable e;
+    while ((e = queue.deleteByName(keyName)) != null) {
+      executor.remove(e);
+    }
+    writeLock(keyName);
     try {
-      Runnable e;
-      while ((e = queue.deleteByName(keyName)) != null) {
-        executor.remove(e);
+      LinkedBlockingQueue kq = keyQueues.getIfPresent(keyName);
+      if (kq != null) {
+        kq.clear();
       }
-      writeLock(keyName);
-      try {
-        keyQueues.get(keyName).clear();
-      } finally {
-        writeUnlock(keyName);
-      }
-    } catch (ExecutionException ex) {
-      //NOP
+    } finally {
+      writeUnlock(keyName);
     }
   }