feat: add in process channel support (#31)

* feat: add in process channel support

* test: add tests for inprocess and copy constructor

Author: aaron-steinfeld

grpc-client-utils/build.gradle.kts

@@ -11,15 +11,16 @@ dependencies {
       because("https://snyk.io/vuln/SNYK-JAVA-IONETTY-2812456")
     }
   }
-  api(platform("io.grpc:grpc-bom:1.45.1"))
+  api(platform("io.grpc:grpc-bom:1.47.0"))
   api("io.grpc:grpc-context")
   api("io.grpc:grpc-api")
 
   implementation(project(":grpc-context-utils"))
   implementation("org.slf4j:slf4j-api:1.7.36")
+  implementation("io.grpc:grpc-core")
 
-  annotationProcessor("org.projectlombok:lombok:1.18.22")
-  compileOnly("org.projectlombok:lombok:1.18.22")
+  annotationProcessor("org.projectlombok:lombok:1.18.24")
+  compileOnly("org.projectlombok:lombok:1.18.24")
 
   testImplementation("org.junit.jupiter:junit-jupiter:5.8.2")
   testImplementation("org.mockito:mockito-core:4.4.0")

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

@@ -9,14 +9,25 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class GrpcChannelRegistry {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcChannelRegistry.class);
-  private final Map<String, ManagedChannel> channelMap = new ConcurrentHashMap<>();
+  private final Map<String, ManagedChannel> channelMap;
   private volatile boolean isShutdown = false;
 
+  public GrpcChannelRegistry(GrpcChannelRegistry sourceRegistry) {
+    // Copy constructor
+    this.isShutdown = sourceRegistry.isShutdown();
+    this.channelMap = new ConcurrentHashMap<>(sourceRegistry.channelMap);
+  }
+
+  public GrpcChannelRegistry() {
+    this.channelMap = new ConcurrentHashMap<>();
+  }
+
   /**
    * Use either {@link #forSecureAddress(String, int)} or {@link #forPlaintextAddress(String, int)}
    */
@@ -30,44 +41,56 @@ public ManagedChannel forSecureAddress(String host, int port) {
   }
 
   public ManagedChannel forSecureAddress(String host, int port, GrpcChannelConfig config) {
-    assert !this.isShutdown;
+    assert !this.isShutdown();
     String channelId = this.getChannelId(host, port, false, config);
-    return this.channelMap.computeIfAbsent(
-        channelId, unused -> this.buildNewChannel(host, port, false, config));
+    return this.getOrComputeChannel(
+        channelId, () -> this.buildNewChannel(host, port, false, config));
   }
 
   public ManagedChannel forPlaintextAddress(String host, int port) {
     return forPlaintextAddress(host, port, GrpcChannelConfig.builder().build());
   }
 
   public ManagedChannel forPlaintextAddress(String host, int port, GrpcChannelConfig config) {
-    assert !this.isShutdown;
+    assert !this.isShutdown();
     String channelId = this.getChannelId(host, port, true, config);
-    return this.channelMap.computeIfAbsent(
-        channelId, unused -> this.buildNewChannel(host, port, true, config));
+    return this.getOrComputeChannel(
+        channelId, () -> this.buildNewChannel(host, port, true, config));
   }
 
   private ManagedChannel buildNewChannel(
       String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
     LOG.info("Creating new channel {}", this.getChannelId(host, port, isPlaintext, config));
 
-    ManagedChannelBuilder<?> builder = ManagedChannelBuilder.forAddress(host, port);
+    ManagedChannelBuilder<?> builder = this.getBuilderForAddress(host, port);
     if (isPlaintext) {
       builder.usePlaintext();
     }
+    return this.configureAndBuildChannel(builder, config);
+  }
+
+  ManagedChannelBuilder<?> getBuilderForAddress(String host, int port) {
+    return ManagedChannelBuilder.forAddress(host, port);
+  }
 
+  ManagedChannel configureAndBuildChannel(
+      ManagedChannelBuilder<?> builder, GrpcChannelConfig config) {
     if (config.getMaxInboundMessageSize() != null) {
       builder.maxInboundMessageSize(config.getMaxInboundMessageSize());
     }
     return builder.build();
   }
 
-  private String getChannelId(
-      String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
+  String getChannelId(String host, int port, boolean isPlaintext, GrpcChannelConfig config) {
     String securePrefix = isPlaintext ? "plaintext" : "secure";
     return securePrefix + ":" + host + ":" + port + ":" + Objects.hash(config);
   }
 
+  final ManagedChannel getOrComputeChannel(
+      String channelId, Supplier<ManagedChannel> channelSupplier) {
+    return this.channelMap.computeIfAbsent(channelId, unused -> channelSupplier.get());
+  }
+
   /**
    * Shuts down channels using a default deadline of 1 minute.
    *
@@ -109,6 +132,10 @@ public void shutdown(Deadline deadline) {
     this.channelMap.clear();
   }
 
+  public boolean isShutdown() {
+    return this.isShutdown;
+  }
+
   private void initiateChannelShutdown(String channelId, ManagedChannel managedChannel) {
     LOG.info("Starting shutdown for channel [{}]", channelId);
     managedChannel.shutdown();

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

@@ -0,0 +1,72 @@
+package org.hypertrace.core.grpcutils.client;
+
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+
+public class InProcessGrpcChannelRegistry extends GrpcChannelRegistry {
+  private final Map<String, String> authorityToInProcessNamedOverride;
+
+  public InProcessGrpcChannelRegistry() {
+    this(Collections.emptyMap());
+  }
+
+  public InProcessGrpcChannelRegistry(Map<String, String> authorityToInProcessNamedOverride) {
+    this.authorityToInProcessNamedOverride = authorityToInProcessNamedOverride;
+  }
+
+  public InProcessGrpcChannelRegistry(GrpcChannelRegistry sourceRegistry) {
+    super(sourceRegistry);
+    if (sourceRegistry instanceof InProcessGrpcChannelRegistry) {
+      this.authorityToInProcessNamedOverride =
+          Map.copyOf(
+              ((InProcessGrpcChannelRegistry) sourceRegistry).authorityToInProcessNamedOverride);
+    } else {
+      this.authorityToInProcessNamedOverride = Collections.emptyMap();
+    }
+  }
+
+  public ManagedChannel forName(String name) {
+    return this.forName(name, GrpcChannelConfig.builder().build());
+  }
+
+  public ManagedChannel forName(String name, GrpcChannelConfig config) {
+    assert !this.isShutdown();
+    return this.getOrComputeChannel(
+        this.getInProcessChannelId(name, config), () -> this.buildInProcessChannel(name, config));
+  }
+
+  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) {
+    return this.getInProcessOverrideNameForHostAndPort(host, port)
+        .<ManagedChannelBuilder<?>>map(InProcessChannelBuilder::forName)
+        .orElseGet(() -> super.getBuilderForAddress(host, port));
+  }
+
+  ManagedChannel buildInProcessChannel(String name, GrpcChannelConfig config) {
+    return this.configureAndBuildChannel(InProcessChannelBuilder.forName(name), config);
+  }
+
+  String getInProcessChannelId(String name, GrpcChannelConfig config) {
+    return "inprocess:" + name + ":" + Objects.hash(config);
+  }
+
+  private Optional<String> getInProcessOverrideNameForHostAndPort(String host, int port) {
+    return Optional.ofNullable(
+        this.authorityToInProcessNamedOverride.get(this.toAuthority(host, port)));
+  }
+
+  private String toAuthority(String host, int port) {
+    return host + ":" + port;
+  }
+}

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

@@ -16,6 +16,7 @@
 import io.grpc.Deadline.Ticker;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -127,4 +128,13 @@ void throwsIfNewChannelRequestedAfterShutdown() {
     assertThrows(AssertionError.class, () -> this.channelRegistry.forPlaintextAddress("foo", 1000));
     assertThrows(AssertionError.class, () -> this.channelRegistry.forSecureAddress("foo", 1000));
   }
+
+  @Test
+  void copyConstructorReusesExistingChannels() {
+    GrpcChannelRegistry firstRegistry = new GrpcChannelRegistry();
+
+    Channel firstChannel = firstRegistry.forSecureAddress("foo", 1000);
+
+    assertSame(firstChannel, new GrpcChannelRegistry(firstRegistry).forSecureAddress("foo", 1000));
+  }
 }

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

@@ -0,0 +1,77 @@
+package org.hypertrace.core.grpcutils.client;
+
+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 io.grpc.Channel;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class InProcessGrpcChannelRegistryTest {
+  InProcessGrpcChannelRegistry channelRegistry;
+
+  @BeforeEach
+  void beforeEach() {
+    this.channelRegistry = new InProcessGrpcChannelRegistry();
+  }
+
+  @Test
+  void createsNewChannelsAsRequested() {
+    // Regular channel
+    assertEquals("foo:1000", this.channelRegistry.forPlaintextAddress("foo", 1000).authority());
+    assertEquals("foo:1000", this.channelRegistry.forSecureAddress("foo", 1000).authority());
+    // In process runs on localhost
+    assertEquals("localhost", this.channelRegistry.forName("inprocess-name").authority());
+  }
+
+  @Test
+  void reusesInProcessChannels() {
+    assertSame(
+        this.channelRegistry.forName("inprocess-name"),
+        this.channelRegistry.forName("inprocess-name"));
+
+    // But not if their config differs
+    assertNotSame(
+        this.channelRegistry.forName("inprocess-name"),
+        this.channelRegistry.forName(
+            "inprocess-name", GrpcChannelConfig.builder().maxInboundMessageSize(100).build()));
+  }
+
+  @Test
+  void overridesAuthorityByConfig() {
+    this.channelRegistry = new InProcessGrpcChannelRegistry(Map.of("foo:1000", "inprocess-name"));
+    assertSame(
+        this.channelRegistry.forSecureAddress("foo", 1000),
+        this.channelRegistry.forName("inprocess-name"));
+    assertSame(
+        this.channelRegistry.forPlaintextAddress("foo", 1000),
+        this.channelRegistry.forName("inprocess-name"));
+    // Also works with custom config
+
+    assertSame(
+        this.channelRegistry.forSecureAddress(
+            "foo", 1000, GrpcChannelConfig.builder().maxInboundMessageSize(100).build()),
+        this.channelRegistry.forName(
+            "inprocess-name", GrpcChannelConfig.builder().maxInboundMessageSize(100).build()));
+
+    // but custom config shouldn't match the default config ones
+    assertNotSame(
+        this.channelRegistry.forSecureAddress("foo", 1000),
+        this.channelRegistry.forName(
+            "inprocess-name", GrpcChannelConfig.builder().maxInboundMessageSize(100).build()));
+  }
+
+  @Test
+  void copyConstructorReusesExistingChannelsAndOverrides() {
+    InProcessGrpcChannelRegistry firstRegistry =
+        new InProcessGrpcChannelRegistry(Map.of("foo:1000", "inprocess-name"));
+
+    Channel firstChannel = firstRegistry.forName("inprocess-name");
+
+    assertSame(
+        firstChannel,
+        new InProcessGrpcChannelRegistry(firstRegistry).forSecureAddress("foo", 1000));
+  }
+}