feat: add registry-level client interceptor support (#42)

* feat: add registry-level client interceptor support

* refactor: add in process support

* style: spotless

Author: aaron-steinfeld

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

@@ -8,6 +8,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
@@ -16,16 +17,27 @@
 public class GrpcChannelRegistry {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcChannelRegistry.class);
   private final Map<String, ManagedChannel> channelMap;
+  private final GrpcRegistryConfig registryConfig;
   private volatile boolean isShutdown = false;
 
   public GrpcChannelRegistry(GrpcChannelRegistry sourceRegistry) {
     // Copy constructor
+    this(sourceRegistry.registryConfig, new ConcurrentHashMap<>(sourceRegistry.channelMap));
     this.isShutdown = sourceRegistry.isShutdown();
-    this.channelMap = new ConcurrentHashMap<>(sourceRegistry.channelMap);
   }
 
   public GrpcChannelRegistry() {
-    this.channelMap = new ConcurrentHashMap<>();
+    this(GrpcRegistryConfig.builder().build());
+  }
+
+  public GrpcChannelRegistry(GrpcRegistryConfig registryConfig) {
+    this(registryConfig, new ConcurrentHashMap<>());
+  }
+
+  private GrpcChannelRegistry(
+      GrpcRegistryConfig registryConfig, ConcurrentMap<String, ManagedChannel> channelMap) {
+    this.registryConfig = registryConfig;
+    this.channelMap = channelMap;
   }
 
   /**
@@ -78,6 +90,7 @@ protected ManagedChannel configureAndBuildChannel(
     if (config.getMaxInboundMessageSize() != null) {
       builder.maxInboundMessageSize(config.getMaxInboundMessageSize());
     }
+    this.registryConfig.getDefaultInterceptors().forEach(builder::intercept);
     return builder.intercept(config.getClientInterceptors()).build();
   }
 

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

@@ -0,0 +1,14 @@
+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(toBuilder = true)
+public class GrpcRegistryConfig {
+
+  @Singular List<ClientInterceptor> defaultInterceptors;
+}

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

@@ -16,6 +16,16 @@ public InProcessGrpcChannelRegistry() {
   }
 
   public InProcessGrpcChannelRegistry(Map<String, String> authorityToInProcessNamedOverride) {
+    this(authorityToInProcessNamedOverride, GrpcRegistryConfig.builder().build());
+  }
+
+  public InProcessGrpcChannelRegistry(GrpcRegistryConfig registryConfig) {
+    this(Collections.emptyMap(), registryConfig);
+  }
+
+  public InProcessGrpcChannelRegistry(
+      Map<String, String> authorityToInProcessNamedOverride, GrpcRegistryConfig registryConfig) {
+    super(registryConfig);
     this.authorityToInProcessNamedOverride = authorityToInProcessNamedOverride;
   }
 

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

@@ -5,14 +5,19 @@
 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.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 
 import io.grpc.Channel;
+import io.grpc.ClientInterceptor;
 import io.grpc.Deadline;
 import io.grpc.Deadline.Ticker;
 import io.grpc.ManagedChannel;
@@ -139,4 +144,36 @@ void copyConstructorReusesExistingChannels() {
 
     assertSame(firstChannel, new GrpcChannelRegistry(firstRegistry).forSecureAddress("foo", 1000));
   }
+
+  @Test
+  void registersRegistryInterceptors() {
+    try (MockedStatic<ManagedChannelBuilder> mockedBuilderStatic =
+        Mockito.mockStatic(ManagedChannelBuilder.class)) {
+      ManagedChannelBuilder<?> mockBuilder = mock(ManagedChannelBuilder.class);
+      ManagedChannel mockChannel = mock(ManagedChannel.class);
+
+      mockedBuilderStatic
+          .when(() -> ManagedChannelBuilder.forAddress("foo", 1000))
+          .thenAnswer(
+              invocation -> {
+                when(mockBuilder.intercept(anyList())).then(RETURNS_SELF);
+                when(mockBuilder.intercept(any(ClientInterceptor.class))).then(RETURNS_SELF);
+                when(mockBuilder.build()).thenReturn(mockChannel);
+                return mockBuilder;
+              });
+
+      ClientInterceptor mockInterceptor1 = mock(ClientInterceptor.class);
+      ClientInterceptor mockInterceptor2 = mock(ClientInterceptor.class);
+      this.channelRegistry =
+          new GrpcChannelRegistry(
+              GrpcRegistryConfig.builder()
+                  .defaultInterceptor(mockInterceptor1)
+                  .defaultInterceptor(mockInterceptor2)
+                  .build());
+
+      assertSame(mockChannel, this.channelRegistry.forPlaintextAddress("foo", 1000));
+      verify(mockBuilder, times(1)).intercept(mockInterceptor1);
+      verify(mockBuilder, times(1)).intercept(mockInterceptor2);
+    }
+  }
 }

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

@@ -3,11 +3,23 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.mockito.Answers.RETURNS_SELF;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import io.grpc.Channel;
+import io.grpc.ClientInterceptor;
+import io.grpc.ManagedChannel;
+import io.grpc.inprocess.InProcessChannelBuilder;
 import java.util.Map;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 class InProcessGrpcChannelRegistryTest {
   InProcessGrpcChannelRegistry channelRegistry;
@@ -74,4 +86,38 @@ void copyConstructorReusesExistingChannelsAndOverrides() {
         firstChannel,
         new InProcessGrpcChannelRegistry(firstRegistry).forSecureAddress("foo", 1000));
   }
+
+  @Test
+  void registersRegistryInterceptors() {
+    try (MockedStatic<InProcessChannelBuilder> mockedBuilderStatic =
+        Mockito.mockStatic(InProcessChannelBuilder.class)) {
+      InProcessChannelBuilder mockBuilder = mock(InProcessChannelBuilder.class);
+      ManagedChannel mockChannel = mock(ManagedChannel.class);
+
+      mockedBuilderStatic
+          .when(() -> InProcessChannelBuilder.forName("test"))
+          .thenAnswer(
+              invocation -> {
+                when(mockBuilder.intercept(anyList())).then(RETURNS_SELF);
+                when(mockBuilder.intercept(any(ClientInterceptor.class))).then(RETURNS_SELF);
+                when(mockBuilder.build()).thenReturn(mockChannel);
+                return mockBuilder;
+              });
+
+      ClientInterceptor mockInterceptor1 = mock(ClientInterceptor.class);
+      ClientInterceptor mockInterceptor2 = mock(ClientInterceptor.class);
+      this.channelRegistry =
+          new InProcessGrpcChannelRegistry(
+              Map.of("host:1000", "test"),
+              GrpcRegistryConfig.builder()
+                  .defaultInterceptor(mockInterceptor1)
+                  .defaultInterceptor(mockInterceptor2)
+                  .build());
+
+      assertSame(mockChannel, this.channelRegistry.forName("test"));
+      assertSame(mockChannel, this.channelRegistry.forPlaintextAddress("host", 1000));
+      verify(mockBuilder, times(1)).intercept(mockInterceptor1);
+      verify(mockBuilder, times(1)).intercept(mockInterceptor2);
+    }
+  }
 }