refactor: add support for client interceptors to registry (#34)

Author: aaron-steinfeld

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

@@ -1,10 +1,15 @@
 package org.hypertrace.core.grpcutils.client;
 
+import io.grpc.ClientInterceptor;
+import java.util.List;
 import lombok.Builder;
+import lombok.Singular;
 import lombok.Value;
 
 @Value
-@Builder
+@Builder(toBuilder = true)
 public class GrpcChannelConfig {
   Integer maxInboundMessageSize;
+
+  @Singular List<ClientInterceptor> clientInterceptors;
 }

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

@@ -69,19 +69,20 @@ private ManagedChannel buildNewChannel(
     return this.configureAndBuildChannel(builder, config);
   }
 
-  ManagedChannelBuilder<?> getBuilderForAddress(String host, int port) {
+  protected ManagedChannelBuilder<?> getBuilderForAddress(String host, int port) {
     return ManagedChannelBuilder.forAddress(host, port);
   }
 
-  ManagedChannel configureAndBuildChannel(
+  protected ManagedChannel configureAndBuildChannel(
       ManagedChannelBuilder<?> builder, GrpcChannelConfig config) {
     if (config.getMaxInboundMessageSize() != null) {
       builder.maxInboundMessageSize(config.getMaxInboundMessageSize());
     }
-    return builder.build();
+    return builder.intercept(config.getClientInterceptors()).build();
   }
 
-  String getChannelId(String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
+  protected String getChannelId(
+      String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
     String securePrefix = isPlaintext ? "plaintext" : "secure";
     return securePrefix + ":" + host + ":" + port + ":" + Objects.hash(config);
   }

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

@@ -40,14 +40,15 @@ public ManagedChannel forName(String name, GrpcChannelConfig config) {
         this.getInProcessChannelId(name, config), () -> this.buildInProcessChannel(name, config));
   }
 
-  String getChannelId(String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
+  protected String getChannelId(
+      String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
     return this.getInProcessOverrideNameForHostAndPort(host, port)
         .map(name -> this.getInProcessChannelId(name, config))
         .orElseGet(() -> super.getChannelId(host, port, isPlaintext, config));
   }
 
   @Override
-  ManagedChannelBuilder<?> getBuilderForAddress(String host, int port) {
+  protected ManagedChannelBuilder<?> getBuilderForAddress(String host, int port) {
     return this.getInProcessOverrideNameForHostAndPort(host, port)
         .<ManagedChannelBuilder<?>>map(InProcessChannelBuilder::forName)
         .orElseGet(() -> super.getBuilderForAddress(host, port));

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

@@ -4,6 +4,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Answers.RETURNS_SELF;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.inOrder;
@@ -16,6 +17,7 @@
 import io.grpc.Deadline.Ticker;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -79,6 +81,7 @@ void shutdownAllChannelsOnShutdown() throws InterruptedException {
           .thenAnswer(
               invocation -> {
                 ManagedChannelBuilder<?> mockBuilder = mock(ManagedChannelBuilder.class);
+                when(mockBuilder.intercept(List.of())).then(RETURNS_SELF);
                 when(mockBuilder.build()).thenReturn(mock(ManagedChannel.class));
                 return mockBuilder;
               });