Added support for allocated resources management

Author: Zgoda91

xds/src/main/java/io/grpc/xds/internal/TlsXdsCredentialsProvider.java

@@ -16,15 +16,18 @@
 
 package io.grpc.xds.internal;
 
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.Durations;
 import io.grpc.ChannelCredentials;
+import io.grpc.ResourceAllocatingChannelCredentials;
 import io.grpc.TlsChannelCredentials;
 import io.grpc.internal.GrpcUtil;
 import io.grpc.internal.JsonUtil;
 import io.grpc.util.AdvancedTlsX509KeyManager;
 import io.grpc.util.AdvancedTlsX509TrustManager;
 import io.grpc.xds.XdsCredentialsProvider;
+import java.io.Closeable;
 import java.io.File;
 import java.text.ParseException;
 import java.util.Map;
@@ -51,10 +54,10 @@ public final class TlsXdsCredentialsProvider extends XdsCredentialsProvider {
 
   @Override
   protected ChannelCredentials newChannelCredentials(Map<String, ?> jsonConfig) {
-    TlsChannelCredentials.Builder builder = TlsChannelCredentials.newBuilder();
+    TlsChannelCredentials.Builder tlsChannelCredsBuilder = TlsChannelCredentials.newBuilder();
 
     if (jsonConfig == null) {
-      return builder.build();
+      return tlsChannelCredsBuilder.build();
     }
 
     // use refresh interval from bootstrap config if provided; else defaults to 600s
@@ -70,17 +73,22 @@ protected ChannelCredentials newChannelCredentials(Map<String, ?> jsonConfig) {
       }
     }
 
+    ImmutableList.Builder<Closeable> resourcesBuilder = ImmutableList.builder();
+    ScheduledExecutorService scheduledExecutorService = null;
+
     // use trust certificate file path from bootstrap config if provided; else use system default
     String rootCertPath = JsonUtil.getString(jsonConfig, ROOT_FILE_KEY);
     if (rootCertPath != null) {
       try {
+        scheduledExecutorService = scheduledExecutorServiceFactory.create();
         AdvancedTlsX509TrustManager trustManager = AdvancedTlsX509TrustManager.newBuilder().build();
-        trustManager.updateTrustCredentials(
+        Closeable trustManagerFuture = trustManager.updateTrustCredentials(
             new File(rootCertPath),
             refreshIntervalSeconds,
             TimeUnit.SECONDS,
-            scheduledExecutorServiceFactory.create());
-        builder.trustManager(trustManager);
+            scheduledExecutorService);
+        resourcesBuilder.add(trustManagerFuture);
+        tlsChannelCredsBuilder.trustManager(trustManager);
       } catch (Exception e) {
         logger.log(Level.WARNING, "Unable to read root certificates", e);
         return null;
@@ -93,14 +101,18 @@ protected ChannelCredentials newChannelCredentials(Map<String, ?> jsonConfig) {
     String privateKeyPath = JsonUtil.getString(jsonConfig, KEY_FILE_KEY);
     if (certChainPath != null && privateKeyPath != null) {
       try {
+        if (scheduledExecutorService == null) {
+          scheduledExecutorService = scheduledExecutorServiceFactory.create();
+        }
         AdvancedTlsX509KeyManager keyManager = new AdvancedTlsX509KeyManager();
-        keyManager.updateIdentityCredentials(
+        Closeable keyManagerFuture = keyManager.updateIdentityCredentials(
             new File(certChainPath),
             new File(privateKeyPath),
             refreshIntervalSeconds,
             TimeUnit.SECONDS,
-            scheduledExecutorServiceFactory.create());
-        builder.keyManager(keyManager);
+            scheduledExecutorService);
+        resourcesBuilder.add(keyManagerFuture);
+        tlsChannelCredsBuilder.keyManager(keyManager);
       } catch (Exception e) {
         logger.log(Level.WARNING, "Unable to read certificate chain or private key", e);
         return null;
@@ -110,7 +122,17 @@ protected ChannelCredentials newChannelCredentials(Map<String, ?> jsonConfig) {
       return null;
     }
 
-    return builder.build();
+    // if executor was initialized, add it to allocated resource list
+    if (scheduledExecutorService != null) {
+      resourcesBuilder.add(asCloseable(scheduledExecutorService));
+    }
+
+    return ResourceAllocatingChannelCredentials.create(
+      tlsChannelCredsBuilder.build(), resourcesBuilder.build());
+  }
+
+  private static Closeable asCloseable(ScheduledExecutorService scheduledExecutorService) {
+    return () -> scheduledExecutorService.shutdownNow();
   }
 
   @Override

xds/src/test/java/io/grpc/xds/internal/TlsXdsCredentialsProviderTest.java

@@ -29,11 +29,13 @@
 import com.google.common.collect.Iterables;
 import io.grpc.ChannelCredentials;
 import io.grpc.InternalServiceProviders;
+import io.grpc.ResourceAllocatingChannelCredentials;
 import io.grpc.TlsChannelCredentials;
 import io.grpc.internal.testing.TestUtils;
 import io.grpc.util.AdvancedTlsX509KeyManager;
 import io.grpc.util.AdvancedTlsX509TrustManager;
 import io.grpc.xds.XdsCredentialsProvider;
+import java.io.Closeable;
 import java.util.Map;
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.TrustManager;
@@ -111,13 +113,20 @@ public void channelCredentialsWhenValidConfig() throws Exception {
         "refresh_interval", "440s");
 
     ChannelCredentials creds = provider.newChannelCredentials(jsonConfig);
-    assertSame(TlsChannelCredentials.class, creds.getClass());
-    TlsChannelCredentials tlsChannelCredentials = (TlsChannelCredentials) creds;
+    assertSame(ResourceAllocatingChannelCredentials.class, creds.getClass());
+    ResourceAllocatingChannelCredentials resourceAllocatingChannelCredentials =
+        (ResourceAllocatingChannelCredentials) creds;
+    TlsChannelCredentials tlsChannelCredentials =
+        (TlsChannelCredentials) resourceAllocatingChannelCredentials.getChannelCredentials();
     assertThat(tlsChannelCredentials.getKeyManagers()).hasSize(1);
     KeyManager keyManager = Iterables.getOnlyElement(tlsChannelCredentials.getKeyManagers());
     assertThat(keyManager).isInstanceOf(AdvancedTlsX509KeyManager.class);
     assertThat(tlsChannelCredentials.getTrustManagers()).hasSize(1);
     TrustManager trustManager = Iterables.getOnlyElement(tlsChannelCredentials.getTrustManagers());
     assertThat(trustManager).isInstanceOf(AdvancedTlsX509TrustManager.class);
+    assertThat(resourceAllocatingChannelCredentials.getAllocatedResources()).hasSize(3);
+    for (Closeable resource : resourceAllocatingChannelCredentials.getAllocatedResources()) {
+      resource.close();
+    }
   }
 }