Close Java peer channels when consumers leave

src/consumer/Consumer.java

@@ -16,6 +16,8 @@
 import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.AbstractService;
 import com.google.common.util.concurrent.RateLimiter;
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
 import java.time.Clock;
 import java.util.List;
 import java.util.Set;
@@ -61,6 +63,12 @@ public class Consumer extends AbstractService {
   private ClusterMap cluster;
   private final Object lock = new Object();
 
+  /** Dials plaintext channels with the standard gRPC builder unless the application overrides it. */
+  static final Function<String, ManagedChannel> DEFAULT_PEER_CHANNEL_FACTORY =
+      target -> ManagedChannelBuilder.forTarget(target).usePlaintext().build();
+
+  private final PeerChannels peerChannels;
+
   /**
    * Creates a new consumer with default values. It will use the local IP address, pod region and
    * pod hostname to register with the coordinator. The grant handlers will be started in a new
@@ -99,6 +107,36 @@ protected Consumer(
       ScheduledExecutorService scheduler,
       BooleanSupplier verbose,
       Function<Stream<Grant>, JoinMessage> registerFactory) {
+    this(
+        clock,
+        service,
+        instance,
+        joinClient,
+        handler,
+        handlerPool,
+        scheduler,
+        verbose,
+        registerFactory,
+        new PeerChannels(DEFAULT_PEER_CHANNEL_FACTORY));
+  }
+
+  /**
+   * Creates a new consumer with the given peer channel pool.
+   *
+   * @param peerChannels Pool of channels to peer consumers. The consumer manages its lifecycle:
+   *     channels follow cluster membership and close on shutdown.
+   */
+  protected Consumer(
+      Clock clock,
+      QualifiedServiceName service,
+      Instance instance,
+      JoinClient joinClient,
+      WorkHandler handler,
+      ExecutorService handlerPool,
+      ScheduledExecutorService scheduler,
+      BooleanSupplier verbose,
+      Function<Stream<Grant>, JoinMessage> registerFactory,
+      PeerChannels peerChannels) {
     this.clock = clock;
     this.instance = instance;
     this.service = service;
@@ -117,6 +155,7 @@ protected Consumer(
     this.cluster = new ClusterMap(ClusterId.DEFAULT, List.of());
     this.verbose = verbose;
     this.registerFactory = registerFactory;
+    this.peerChannels = peerChannels;
   }
 
   /**
@@ -134,6 +173,13 @@ public Instance getInstance() {
     return instance;
   }
 
+  /**
+   * @return channels to other consumers in the cluster, closed as consumers leave.
+   */
+  public PeerChannels getPeerChannels() {
+    return peerChannels;
+  }
+
   /**
    * Start the consumer and initiate a new session to the service. The consumer will re-connect
    * automatically in case of session termination until {@link #stopAsync()} is called. No-op if the
@@ -259,14 +305,19 @@ private void handleClientMessage(ClientMessage msg) {
   }
 
   private void handleClusterMessage(ClusterMessage msg) {
+    List<Instance> consumers = null;
     synchronized (lock) {
       try {
         this.cluster =
             this.cluster.update(co.atoms.splitter.cluster.Messages.ClusterMessage.fromProto(msg));
+        consumers = this.cluster.getConsumers();
       } catch (Exception e) {
         LOGGER.error("Failed to update cluster with message {}. Ignoring the update.", msg, e);
       }
     }
+    if (consumers != null) {
+      peerChannels.retain(consumers);
+    }
   }
 
   /**
@@ -291,6 +342,7 @@ private void terminate(@Nullable Session<JoinMessage, JoinMessage> session) {
     } catch (InterruptedException e) {
       throw new RuntimeException(e);
     } finally {
+      peerChannels.close();
       notifyStopped();
     }
   }

src/consumer/ConsumerBuilder.java

@@ -42,6 +42,7 @@ public class ConsumerBuilder {
   private @Nullable DomainKeyName[] keyNames;
   private @Nullable BooleanSupplier verbose;
   private @Nullable String version;
+  private @Nullable PeerChannels peerChannels;
 
   /**
    * Creates a new builder for a {@link Consumer} with default values. It will use the local IP
@@ -145,6 +146,18 @@ public ConsumerBuilder withVersion(String version) {
     return this;
   }
 
+  /**
+   * Sets the pool of channels to peer consumers. By default, the pool uses the standard gRPC
+   * builder with plaintext transport. The consumer manages the pool's lifecycle: channels follow
+   * cluster membership and close on shutdown.
+   *
+   * @return the builder with the peer channels set
+   */
+  public ConsumerBuilder withPeerChannels(PeerChannels peerChannels) {
+    this.peerChannels = peerChannels;
+    return this;
+  }
+
   public Consumer build() {
     var clock = this.clock;
     if (clock == null) {
@@ -167,6 +180,10 @@ public Consumer build() {
       verbose = () -> false;
     }
     var version = this.version;
+    var peerChannels = this.peerChannels;
+    if (peerChannels == null) {
+      peerChannels = new PeerChannels(Consumer.DEFAULT_PEER_CHANNEL_FACTORY);
+    }
 
     Function<Stream<Grant>, JoinMessage> registerFactory =
         (grants) -> {
@@ -203,7 +220,8 @@ public Consumer build() {
         handlerPool,
         scheduler,
         verbose,
-        registerFactory);
+        registerFactory,
+        peerChannels);
   }
 
   static Instance defaultInstance(int port) {

src/consumer/ConsumerTest.java

@@ -11,7 +11,9 @@
 import static co.atoms.splitter.testing.Fixtures.SHARD_D2_GLOBAL_15;
 import static co.atoms.splitter.testing.Fixtures.TS;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Timeout.ThreadMode.SEPARATE_THREAD;
 
 import co.atoms.splitter.cluster.Cluster;
@@ -20,12 +22,14 @@
 import co.atoms.splitter.model.Grant;
 import co.atoms.splitter.model.GrantState;
 import co.atoms.splitter.model.Instance;
+import co.atoms.splitter.model.Location;
 import co.atoms.splitter.model.Shard;
 import co.atoms.splitter.testing.Fixtures;
 import co.atoms.splitter.testing.Gate;
 import co.atoms.splitter.testing.MoreMessages;
 import java.time.Clock;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
@@ -80,6 +84,68 @@ void clusterUpdates(TestInfo testInfo) {
     deps.stop();
   }
 
+  @Test
+  void consumerRemovalClosesPeerChannel(TestInfo testInfo) throws Exception {
+    var deps = setupDeps(testInfo);
+
+    deps.consumer.initialize(EXP);
+
+    var other = Instance.create("other-id", Location.LOCAL, TS, "other", "other-endpoint");
+    var g1 = GrantInfo.create("g1", SHARD_D1_UNIT_0e, GrantState.ALLOCATED_LOADED);
+    var g2 = GrantInfo.create("g2", SHARD_D2_GLOBAL_15, GrantState.ACTIVE);
+    deps.consumer.sendClusterSnapshot(
+        List.of(SHARD_D1_UNIT_0e, SHARD_D2_GLOBAL_15),
+        MoreMessages.newAssignment(deps.instance, g1),
+        MoreMessages.newAssignment(other, g2));
+    assertEquals(deps.consumer.getCluster().getConsumers().size(), 2);
+
+    var channel = deps.consumer.getPeerChannels().channel(other);
+    assertFalse(channel.isShutdown());
+
+    deps.consumer.sendClusterRemove(1, other.id());
+
+    assertEquals(deps.consumer.getCluster().getConsumers(), List.of(deps.instance));
+    // Caffeine runs removal listeners asynchronously.
+    awaitShutdown(channel);
+
+    // Stopping the consumer closes the remaining peer channels.
+    var retained = deps.consumer.getPeerChannels().channel(deps.instance);
+    deps.stop();
+    awaitShutdown(retained);
+  }
+
+  private static void awaitShutdown(io.grpc.ManagedChannel channel) throws InterruptedException {
+    for (int i = 0; i < 1000 && !channel.isShutdown(); i++) {
+      Thread.sleep(10);
+    }
+    assertTrue(channel.isShutdown());
+  }
+
+  @Test
+  void builderSharesInjectedPeerChannels() {
+    var dialed = new ArrayList<String>();
+    var peerChannels =
+        new PeerChannels(
+            target -> {
+              dialed.add(target);
+              return io.grpc.ManagedChannelBuilder.forTarget(target).usePlaintext().build();
+            });
+    var consumer =
+        new ConsumerBuilder(
+                DOMAIN1.service(),
+                responseObserver -> null,
+                (id, shard, ownership) -> {},
+                Fixtures.newInstance("factory-test"))
+            .withPeerChannels(peerChannels)
+            .build();
+
+    assertEquals(peerChannels, consumer.getPeerChannels());
+    var other = Instance.create("other-id", Location.LOCAL, TS, "other", "other-endpoint");
+    var channel = consumer.getPeerChannels().channel(other);
+    assertEquals(List.of("other-endpoint"), dialed);
+    channel.shutdown();
+  }
+
   @Test
   public void reconnect(TestInfo testInfo) throws Exception {
     var deps = setupDeps(testInfo);

src/consumer/PeerChannels.java

@@ -0,0 +1,58 @@
+package co.atoms.splitter.consumer;
+
+import co.atoms.splitter.model.Instance;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import io.grpc.ManagedChannel;
+import java.util.Collection;
+import java.util.Set;
+import java.util.function.Function;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Channels to other consumers in the cluster. Channels are dialed on first use, reused, and closed
+ * when their consumer leaves the cluster.
+ */
+public class PeerChannels {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(PeerChannels.class);
+
+  private final LoadingCache<Instance, ManagedChannel> channels;
+
+  /**
+   * @param dial Creates a channel for a consumer endpoint. Injected so applications choose the
+   *     channel builder (e.g. one with interceptors) and this class stays platform-agnostic.
+   */
+  public PeerChannels(Function<String, ManagedChannel> dial) {
+    this.channels =
+        Caffeine.newBuilder()
+            .recordStats()
+            .maximumSize(100)
+            .<Instance, ManagedChannel>removalListener((k, channel, cause) -> shutdown(channel))
+            .build(consumer -> dial.apply(consumer.endpoint()));
+  }
+
+  /** Returns a shared channel to the consumer, dialing on first use. */
+  public ManagedChannel channel(Instance consumer) {
+    return channels.get(consumer);
+  }
+
+  /** Closes channels to consumers no longer in the cluster. */
+  public void retain(Collection<Instance> consumers) {
+    channels.asMap().keySet().retainAll(Set.copyOf(consumers));
+  }
+
+  /** Closes all channels. */
+  public void close() {
+    channels.invalidateAll();
+  }
+
+  private static void shutdown(ManagedChannel channel) {
+    try {
+      channel.shutdown();
+    } catch (Exception error) {
+      LOGGER.error("Failed to shut down channel", error);
+    }
+  }
+}

src/consumer/PeerChannelsTest.java

@@ -0,0 +1,65 @@
+package co.atoms.splitter.consumer;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import co.atoms.splitter.model.Instance;
+import co.atoms.splitter.model.Location;
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import java.time.Instant;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+public class PeerChannelsTest {
+
+  private static Instance instance(String id) {
+    return Instance.create(id, Location.LOCAL, Instant.EPOCH, id, "localhost:1");
+  }
+
+  private static PeerChannels newPeerChannels() {
+    return new PeerChannels(
+        target -> ManagedChannelBuilder.forTarget(target).usePlaintext().build());
+  }
+
+  @Test
+  void retainClosesDepartedChannels() throws Exception {
+    var channels = newPeerChannels();
+
+    var kept = instance("kept");
+    var departed = instance("departed");
+    var keptChannel = channels.channel(kept);
+    var departedChannel = channels.channel(departed);
+    assertSame(keptChannel, channels.channel(kept));
+
+    channels.retain(List.of(kept));
+
+    // Caffeine runs removal listeners asynchronously.
+    awaitShutdown(departedChannel);
+    assertFalse(keptChannel.isShutdown());
+    assertSame(keptChannel, channels.channel(kept));
+
+    // A fresh channel is dialed when a departed consumer is accessed again.
+    var redialed = channels.channel(departed);
+    assertNotSame(departedChannel, redialed);
+
+    channels.retain(List.of());
+    awaitShutdown(keptChannel);
+    awaitShutdown(redialed);
+  }
+
+  @Test
+  void retainWithoutCachedChannelsIsNoop() {
+    var channels = newPeerChannels();
+    channels.retain(List.of(instance("unknown")));
+  }
+
+  private static void awaitShutdown(ManagedChannel channel) throws InterruptedException {
+    for (int i = 0; i < 1000 && !channel.isShutdown(); i++) {
+      Thread.sleep(10);
+    }
+    assertTrue(channel.isShutdown());
+  }
+}