refactor: add registry support for secure channels (#18)

* refactor: add registry support for secure channels

* docs: finish writing javadoc comment

Author: aaron-steinfeld

grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java

@@ -12,19 +12,41 @@ public class GrpcChannelRegistry {
   private final Map<String, ManagedChannel> channelMap = new ConcurrentHashMap<>();
   private volatile boolean isShutdown = false;
 
+  /**
+   * Use either {@link #forSecureAddress(String, int)} or {@link #forPlaintextAddress(String, int)}
+   */
+  @Deprecated
   public ManagedChannel forAddress(String host, int port) {
+    return this.forPlaintextAddress(host, port);
+  }
+
+  public ManagedChannel forSecureAddress(String host, int port) {
+    assert !this.isShutdown;
+    String channelId = this.getChannelId(host, port, false);
+    return this.channelMap.computeIfAbsent(
+        channelId, unused -> this.buildNewChannel(host, port, false));
+  }
+
+  public ManagedChannel forPlaintextAddress(String host, int port) {
     assert !this.isShutdown;
-    String channelId = this.getChannelId(host, port);
-    return this.channelMap.computeIfAbsent(channelId, unused -> this.buildNewChannel(host, port));
+    String channelId = this.getChannelId(host, port, true);
+    return this.channelMap.computeIfAbsent(
+        channelId, unused -> this.buildNewChannel(host, port, true));
   }
 
-  private ManagedChannel buildNewChannel(String host, int port) {
-    LOG.info("Creating new channel for {}:{}", host, port);
-    return ManagedChannelBuilder.forAddress(host, port).usePlaintext().build();
+  private ManagedChannel buildNewChannel(String host, int port, boolean isPlaintext) {
+    LOG.info("Creating new channel {}", this.getChannelId(host, port, isPlaintext));
+
+    ManagedChannelBuilder<?> builder = ManagedChannelBuilder.forAddress(host, port);
+    if (isPlaintext) {
+      builder.usePlaintext();
+    }
+    return builder.build();
   }
 
-  private String getChannelId(String host, int port) {
-    return host + ":" + port;
+  private String getChannelId(String host, int port, boolean isPlaintext) {
+    String securePrefix = isPlaintext ? "plaintext" : "secure";
+    return securePrefix + ":" + host + ":" + port;
   }
 
   public void shutdown() {

grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistryTest.java

@@ -23,21 +23,26 @@ void beforeEach() {
 
   @Test
   void createsNewChannelsAsRequested() {
-    assertNotNull(this.channelRegistry.forAddress("foo", 1000));
+    assertNotNull(this.channelRegistry.forPlaintextAddress("foo", 1000));
   }
 
   @Test
   void reusesChannelsForDuplicateRequests() {
-    Channel firstChannel = this.channelRegistry.forAddress("foo", 1000);
-    assertSame(firstChannel, this.channelRegistry.forAddress("foo", 1000));
-    assertNotSame(firstChannel, this.channelRegistry.forAddress("foo", 1001));
-    assertNotSame(firstChannel, this.channelRegistry.forAddress("bar", 1000));
+    Channel firstChannel = this.channelRegistry.forPlaintextAddress("foo", 1000);
+    assertSame(firstChannel, this.channelRegistry.forPlaintextAddress("foo", 1000));
+    Channel firstChannelSecure = this.channelRegistry.forSecureAddress("foo", 1000);
+    assertSame(firstChannelSecure, this.channelRegistry.forSecureAddress("foo", 1000));
+    assertNotSame(firstChannel, firstChannelSecure);
+    assertNotSame(firstChannel, this.channelRegistry.forPlaintextAddress("foo", 1001));
+    assertNotSame(firstChannel, this.channelRegistry.forPlaintextAddress("foo", 1001));
+    assertNotSame(firstChannelSecure, this.channelRegistry.forSecureAddress("bar", 1000));
+    assertNotSame(firstChannelSecure, this.channelRegistry.forSecureAddress("bar", 1000));
   }
 
   @Test
   void shutdownAllChannelsOnShutdown() {
-    ManagedChannel firstChannel = this.channelRegistry.forAddress("foo", 1000);
-    ManagedChannel secondChannel = this.channelRegistry.forAddress("foo", 1002);
+    ManagedChannel firstChannel = this.channelRegistry.forPlaintextAddress("foo", 1000);
+    ManagedChannel secondChannel = this.channelRegistry.forSecureAddress("foo", 1002);
     assertFalse(firstChannel.isShutdown());
     assertFalse(secondChannel.isShutdown());
     this.channelRegistry.shutdown();
@@ -48,6 +53,7 @@ void shutdownAllChannelsOnShutdown() {
   @Test
   void throwsIfNewChannelRequestedAfterShutdown() {
     this.channelRegistry.shutdown();
-    assertThrows(AssertionError.class, () -> this.channelRegistry.forAddress("foo", 1000));
+    assertThrows(AssertionError.class, () -> this.channelRegistry.forPlaintextAddress("foo", 1000));
+    assertThrows(AssertionError.class, () -> this.channelRegistry.forSecureAddress("foo", 1000));
   }
 }