GH-3683: Support programmatic KMS client factories

parquet-hadoop/README.md

@@ -455,6 +455,28 @@ If `false`, write files in encrypted footer mode, that fully encrypts the footer
 **Description:** Class implementing the KmsClient interface. "KMS" stands for “key management service”. The Client will interact with a KMS Server to wrap/unrwap encryption keys.  
 **Default value:** None
 
+KMS clients can also be supplied programmatically when they require constructor-injected dependencies:
+
+```java
+KeyToolkit.setKmsClientFactory(
+    configuration,
+    (conf, kmsInstanceID, kmsInstanceURL, accessToken) -> new CustomKmsClient(dependency));
+try {
+  // Construct and close readers and writers using configuration or its copies.
+} finally {
+  KeyToolkit.removeKmsClientFactory(configuration);
+}
+```
+
+A registered factory takes precedence over `parquet.encryption.kms.client.class`. Each invocation must return a
+distinct, uninitialized `KmsClient`; `KeyToolkit` initializes and caches it. The factory can be invoked concurrently
+for different access-token and KMS-instance combinations, so it must be thread-safe.
+
+Copies of the `Configuration` in the same JVM share the registration and its caches. A configuration deserialized in
+another JVM must register the factory there before use. Call `removeKmsClientFactory` only after all readers and writers
+using the configuration and its copies have closed. Registering another factory for the configuration or one of its
+copies replaces the previous factory and clears the registration's caches.
+
 ---
 
 **Property:** `parquet.encryption.kms.instance.id`  

parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyUnwrapper.java

@@ -19,8 +19,6 @@
 
 package org.apache.parquet.crypto.keytools;
 
-import static org.apache.parquet.crypto.keytools.KeyToolkit.KEK_READ_CACHE_PER_TOKEN;
-import static org.apache.parquet.crypto.keytools.KeyToolkit.KMS_CLIENT_CACHE_PER_TOKEN;
 import static org.apache.parquet.crypto.keytools.KeyToolkit.stringIsEmpty;
 
 import java.io.IOException;
@@ -47,6 +45,7 @@ public class FileKeyUnwrapper implements DecryptionKeyRetriever {
   private final Path parquetFilePath;
   private final String accessToken;
   private final long cacheEntryLifetime;
+  private final KeyToolkit.KmsClientCacheContext cacheContext;
 
   FileKeyUnwrapper(Configuration hadoopConfiguration, Path filePath) {
     this.hadoopConfiguration = hadoopConfiguration;
@@ -58,11 +57,13 @@ public class FileKeyUnwrapper implements DecryptionKeyRetriever {
 
     accessToken = hadoopConfiguration.getTrimmed(
         KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, KmsClient.KEY_ACCESS_TOKEN_DEFAULT);
+    cacheContext = KeyToolkit.getKmsClientCacheContext(hadoopConfiguration);
 
     // Check cache upon each file reading (clean once in cacheEntryLifetime)
-    KMS_CLIENT_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
-    KEK_READ_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
-    kekPerKekID = KEK_READ_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
+    cacheContext.getKmsClientCache().checkCacheForExpiredTokens(cacheEntryLifetime);
+    TwoLevelCacheWithExpiration<byte[]> kekReadCache = cacheContext.getKekReadCache();
+    kekReadCache.checkCacheForExpiredTokens(cacheEntryLifetime);
+    kekPerKekID = kekReadCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
 
     if (LOG.isDebugEnabled()) {
       LOG.debug(
@@ -168,7 +169,7 @@ KeyToolkit.KmsClientAndDetails getKmsClientFromConfigOrKeyMaterial(KeyMaterial k
     }
 
     KmsClient kmsClient = KeyToolkit.getKmsClient(
-        kmsInstanceID, kmsInstanceURL, hadoopConfiguration, accessToken, cacheEntryLifetime);
+        kmsInstanceID, kmsInstanceURL, hadoopConfiguration, accessToken, cacheEntryLifetime, cacheContext);
     if (null == kmsClient) {
       throw new ParquetCryptoRuntimeException(
           "KMSClient was not successfully created for reading encrypted data.");

parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/FileKeyWrapper.java

@@ -19,12 +19,10 @@
 
 package org.apache.parquet.crypto.keytools;
 
-import static org.apache.parquet.crypto.keytools.KeyToolkit.KEK_WRITE_CACHE_PER_TOKEN;
-import static org.apache.parquet.crypto.keytools.KeyToolkit.KMS_CLIENT_CACHE_PER_TOKEN;
-
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.util.Arrays;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
@@ -73,25 +71,32 @@ public class FileKeyWrapper {
     accessToken = hadoopConfiguration.getTrimmed(
         KeyToolkit.KEY_ACCESS_TOKEN_PROPERTY_NAME, KmsClient.KEY_ACCESS_TOKEN_DEFAULT);
 
+    KeyToolkit.KmsClientCacheContext cacheContext = KeyToolkit.getKmsClientCacheContext(configuration);
+
     // Check caches upon each file writing (clean once in cacheEntryLifetime)
-    KMS_CLIENT_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
+    cacheContext.getKmsClientCache().checkCacheForExpiredTokens(cacheEntryLifetime);
 
     if (null == kmsClientAndDetails) {
       kmsInstanceID = hadoopConfiguration.getTrimmed(
           KeyToolkit.KMS_INSTANCE_ID_PROPERTY_NAME, KmsClient.KMS_INSTANCE_ID_DEFAULT);
       kmsInstanceURL = hadoopConfiguration.getTrimmed(
           KeyToolkit.KMS_INSTANCE_URL_PROPERTY_NAME, KmsClient.KMS_INSTANCE_URL_DEFAULT);
       kmsClient = KeyToolkit.getKmsClient(
-          kmsInstanceID, kmsInstanceURL, configuration, accessToken, cacheEntryLifetime);
+          kmsInstanceID, kmsInstanceURL, configuration, accessToken, cacheEntryLifetime, cacheContext);
     } else {
       kmsInstanceID = kmsClientAndDetails.getKmsInstanceID();
       kmsInstanceURL = kmsClientAndDetails.getKmsInstanceURL();
       kmsClient = kmsClientAndDetails.getKmsClient();
     }
 
     if (doubleWrapping) {
-      KEK_WRITE_CACHE_PER_TOKEN.checkCacheForExpiredTokens(cacheEntryLifetime);
-      KEKPerMasterKeyID = KEK_WRITE_CACHE_PER_TOKEN.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
+      TwoLevelCacheWithExpiration<ConcurrentMap<String, KeyEncryptionKey>> kekWriteCache =
+          cacheContext.getKekWriteCache();
+      kekWriteCache.checkCacheForExpiredTokens(cacheEntryLifetime);
+      ConcurrentMap<String, ConcurrentMap<String, KeyEncryptionKey>> kekPerKmsInstanceID =
+          kekWriteCache.getOrCreateInternalCache(accessToken, cacheEntryLifetime);
+      KEKPerMasterKeyID =
+          kekPerKmsInstanceID.computeIfAbsent(kmsInstanceID, ignored -> new ConcurrentHashMap<>());
       int kekLengthBits =
           configuration.getInt(KeyToolkit.KEK_LENGTH_PROPERTY_NAME, KeyToolkit.KEK_LENGTH_DEFAULT);
       if (Arrays.binarySearch(ACCEPTABLE_KEK_LENGTHS, kekLengthBits) < 0) {