diff --git a/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java index 72de4ba..073070e 100644 --- a/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java +++ b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistry.java @@ -12,19 +12,41 @@ public class GrpcChannelRegistry { private final Map 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() { diff --git a/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistryTest.java b/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistryTest.java index 26d2a74..de0f0ce 100644 --- a/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/GrpcChannelRegistryTest.java +++ b/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)); } }