From ed5b6f1aec1a44d2fb9781c8220ab708cf68b4d4 Mon Sep 17 00:00:00 2001 From: "pallav.garg" Date: Wed, 21 Oct 2026 11:41:53 -0700 Subject: [PATCH] Close Java peer channels when consumers leave --- src/consumer/Consumer.java | 52 +++++++++++++ src/consumer/ConsumerBuilder.java | 20 ++++- src/consumer/ConsumerTest.java | 66 +++++++++++++++++ src/consumer/PeerChannels.java | 58 +++++++++++++++ src/consumer/PeerChannelsTest.java | 65 +++++++++++++++++ .../processor/GrpcRequestForwarder.java | 46 ++++++++---- .../processor/GrpcRequestForwarderTest.java | 73 +++++++++++++++++++ src/consumer/processor/GrpcStubsCache.java | 7 +- src/consumer/testing/TestConsumer.java | 17 +++++ 9 files changed, 388 insertions(+), 16 deletions(-) create mode 100644 src/consumer/PeerChannels.java create mode 100644 src/consumer/PeerChannelsTest.java diff --git a/src/consumer/Consumer.java b/src/consumer/Consumer.java index b47cfb5..ff135ef 100644 --- a/src/consumer/Consumer.java +++ b/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 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, 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, 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 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 session) { } catch (InterruptedException e) { throw new RuntimeException(e); } finally { + peerChannels.close(); notifyStopped(); } } diff --git a/src/consumer/ConsumerBuilder.java b/src/consumer/ConsumerBuilder.java index 494608b..61111e3 100644 --- a/src/consumer/ConsumerBuilder.java +++ b/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, JoinMessage> registerFactory = (grants) -> { @@ -203,7 +220,8 @@ public Consumer build() { handlerPool, scheduler, verbose, - registerFactory); + registerFactory, + peerChannels); } static Instance defaultInstance(int port) { diff --git a/src/consumer/ConsumerTest.java b/src/consumer/ConsumerTest.java index 44a8152..748b3a9 100644 --- a/src/consumer/ConsumerTest.java +++ b/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(); + 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); diff --git a/src/consumer/PeerChannels.java b/src/consumer/PeerChannels.java new file mode 100644 index 0000000..86916ef --- /dev/null +++ b/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 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 dial) { + this.channels = + Caffeine.newBuilder() + .recordStats() + .maximumSize(100) + .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 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); + } + } +} diff --git a/src/consumer/PeerChannelsTest.java b/src/consumer/PeerChannelsTest.java new file mode 100644 index 0000000..b08cf8c --- /dev/null +++ b/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()); + } +} diff --git a/src/consumer/processor/GrpcRequestForwarder.java b/src/consumer/processor/GrpcRequestForwarder.java index 532523f..8e8e970 100644 --- a/src/consumer/processor/GrpcRequestForwarder.java +++ b/src/consumer/processor/GrpcRequestForwarder.java @@ -4,6 +4,7 @@ import co.atoms.splitter.consumer.processor.exceptions.GrpcExceptions; import co.atoms.splitter.internal.Metrics; import co.atoms.splitter.model.GrantState; +import co.atoms.splitter.model.Instance; import co.atoms.splitter.model.Location; import co.atoms.splitter.model.QualifiedDomainKey; import io.grpc.ManagedChannel; @@ -37,7 +38,7 @@ public interface RemoteHandler grantManager; - private final GrpcStubsCache stubs; + private final Function stubProvider; private final RemoteHandler remoteHandler; private final OwnershipRetryer retryer; private final Duration timeout; @@ -45,6 +46,9 @@ public interface RemoteHandlerThe supplied cache retains its own channel factory and expiration policy; + * its channels are not managed by the consumer's peer channel pool. + * * @param consumer Splitter consumer * @param grantManager Grant manager, usually an instance of a subclass of {@link * BaseWorkProcessor} @@ -60,12 +64,7 @@ public GrpcRequestForwarder( RemoteHandler remoteHandler, OwnershipRetryer retryer, Duration timeout) { - this.consumer = consumer; - this.grantManager = grantManager; - this.remoteHandler = remoteHandler; - this.retryer = retryer; - this.stubs = stubsCache; - this.timeout = timeout; + this(consumer, grantManager, remoteHandler, retryer, timeout, stubsCache::getStub); } /** @@ -74,8 +73,8 @@ public GrpcRequestForwarder( * @param consumer Splitter consumer * @param grantManager Grant manager, usually an instance of a subclass of {@link * BaseWorkProcessor} - * @param stubFactory Creates a new gRPC stub. The created stubs are cached using {@link - * GrpcStubsCache}. + * @param stubFactory Creates a new gRPC stub over a channel from the consumer's {@link + * co.atoms.splitter.consumer.PeerChannels}. * @param remoteHandler Invokes gRPC request * @param retryer Retryer to use when forwarding requests * @param timeout Timeout for one invocation of the request executed to a remote instance. @@ -87,7 +86,28 @@ public GrpcRequestForwarder( RemoteHandler remoteHandler, OwnershipRetryer retryer, Duration timeout) { - this(consumer, grantManager, new GrpcStubsCache(stubFactory), remoteHandler, retryer, timeout); + this( + consumer, + grantManager, + remoteHandler, + retryer, + timeout, + instance -> stubFactory.apply(consumer.getPeerChannels().channel(instance))); + } + + private GrpcRequestForwarder( + Consumer consumer, + GrantManager grantManager, + RemoteHandler remoteHandler, + OwnershipRetryer retryer, + Duration timeout, + Function stubProvider) { + this.consumer = consumer; + this.grantManager = grantManager; + this.remoteHandler = remoteHandler; + this.retryer = retryer; + this.stubProvider = stubProvider; + this.timeout = timeout; } /** @@ -96,8 +116,8 @@ public GrpcRequestForwarder( * @param consumer Splitter consumer * @param grantManager Grant manager, usually an instance of a subclass of {@link * BaseWorkProcessor} - * @param stubFactory Creates a new gRPC stub. The created stubs are cached using {@link - * GrpcStubsCache}. + * @param stubFactory Creates a new gRPC stub over a channel from the consumer's {@link + * co.atoms.splitter.consumer.PeerChannels}. * @param remoteHandler Invokes gRPC request * @param retryer Retryer to use when forwarding requests */ @@ -177,7 +197,7 @@ private void forwardOnce( return; } var consumer = owner.get().consumer(); - var stub = stubs.getStub(consumer); + var stub = stubProvider.apply(consumer); if (stub == null) { LOGGER.error("No stub found for instance {}", consumer); Metrics.recordForwardedRequest(key.domain(), "unknown", "owner_not_found", Location.UNKNOWN); diff --git a/src/consumer/processor/GrpcRequestForwarderTest.java b/src/consumer/processor/GrpcRequestForwarderTest.java index 31c54be..1c0cc38 100644 --- a/src/consumer/processor/GrpcRequestForwarderTest.java +++ b/src/consumer/processor/GrpcRequestForwarderTest.java @@ -13,6 +13,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import co.atoms.splitter.cluster.GrantInfo; @@ -28,10 +29,15 @@ import co.atoms.splitter.testing.TestStreamObserver; import io.grpc.CallOptions; import io.grpc.Channel; +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.StreamObserver; import java.time.Duration; import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -226,6 +232,73 @@ public void handleWithActiveLocal(TestInfo testInfo) { consumer.stopAsync(); } + @Test + public void handleRemoteWithExistingStubCache(TestInfo testInfo) { + var processor = new TestWorkProcessor((id, shard, ownership) -> new TestRange()); + var consumer = + TestConsumer.create( + Fixtures.createClock(), testInfo.getDisplayName(), DOMAIN1.service(), processor); + consumer.initialize(EXP); + + var instance = Fixtures.newInstance("remote"); + var grant = GrantInfo.create("g1", SHARD_D2_GLOBAL_5c, ALLOCATED_LOADED); + consumer.sendClusterSnapshot(MoreMessages.newAssignment(instance, grant)); + + var dialed = new ArrayList(); + var channel = ManagedChannelBuilder.forTarget(instance.endpoint()).usePlaintext().build(); + var cache = + new GrpcStubsCache<>( + managedChannel -> new TestStub(managedChannel, CallOptions.DEFAULT), + target -> { + dialed.add(target); + return channel; + }); + try { + var stub = cache.getStub(instance); + assertSame(stub, cache.getStub(instance)); + var forwarder = + new GrpcRequestForwarder<>( + consumer, + processor, + cache, + (remote, request) -> { + assertSame(channel, remote.getChannel()); + var deadline = remote.getCallOptions().getDeadline(); + assertTrue(deadline.timeRemaining(TimeUnit.MILLISECONDS) > 0); + assertTrue(deadline.timeRemaining(TimeUnit.MILLISECONDS) <= 17000); + return remote.handle(request); + }, + retrier, + Duration.ofSeconds(17)); + var key = Fixtures.newKey(DOMAIN2, "5"); + assertForwarded(o -> forwarder.forward(key, key, o), key); + assertForwarded(o -> forwarder.forward(key, key, o), key); + assertEquals(List.of(instance.endpoint()), dialed); + assertNull(stub.getCallOptions().getDeadline()); + + consumer.sendClusterRemove(1, instance.id()); + assertFalse(channel.isShutdown()); + consumer.stop(); + consumer.awaitTerminated(); + // Existing caches keep ownership of their channels, even after consumer shutdown. + assertFalse(channel.isShutdown()); + } finally { + channel.shutdownNow(); + } + } + + @Test + public void existingStubCacheDefaultConstructor() { + var cache = new GrpcStubsCache<>(channel -> new TestStub(channel, CallOptions.DEFAULT)); + var instance = Fixtures.newInstance("remote"); + var stub = cache.getStub(instance); + try { + assertSame(stub, cache.getStub(instance)); + } finally { + ((ManagedChannel) stub.getChannel()).shutdownNow(); + } + } + @Test public void handleRetry(TestInfo testInfo) { var r = new TestRange(); diff --git a/src/consumer/processor/GrpcStubsCache.java b/src/consumer/processor/GrpcStubsCache.java index aba99eb..1abb241 100644 --- a/src/consumer/processor/GrpcStubsCache.java +++ b/src/consumer/processor/GrpcStubsCache.java @@ -12,8 +12,11 @@ import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder; /** - * Cache for gRPC stubs to avoid creating new channels on each request. Stubs are created using the - * provided factory method and are cached for some time after last access. + * Cache for gRPC stubs to avoid creating new channels on each request. Entries expire ten minutes + * after last access; this cache does not track consumer membership. + * + *

Retained for compatibility. New callers should use the consumer's {@link + * co.atoms.splitter.consumer.PeerChannels} and create stubs over its shared channels. */ public class GrpcStubsCache { diff --git a/src/consumer/testing/TestConsumer.java b/src/consumer/testing/TestConsumer.java index b195a4c..032caa5 100644 --- a/src/consumer/testing/TestConsumer.java +++ b/src/consumer/testing/TestConsumer.java @@ -11,6 +11,8 @@ import co.atoms.splitter.model.Instance; import co.atoms.splitter.model.QualifiedServiceName; import co.atoms.splitter.model.Shard; +import co.atoms.splitter.proto.ClusterMessage; +import co.atoms.splitter.proto.ConsumerMessage; import co.atoms.splitter.proto.JoinMessage; import co.atoms.splitter.testing.Fixtures; import co.atoms.splitter.testing.MoreMessages; @@ -297,4 +299,19 @@ public void sendClusterAssign( int version, co.atoms.splitter.cluster.Messages.Assignment... assignments) { send(MoreMessages.newClusterAssign(SERVER.getId(), version, assignments)); } + + /** Sends a message to the consumer removing consumers from the cluster map */ + public void sendClusterRemove(int version, String... consumers) { + var remove = ClusterMessage.Remove.newBuilder().addAllConsumers(Arrays.asList(consumers)); + send( + JoinMessage.newBuilder() + .setConsumer( + ConsumerMessage.newBuilder() + .setCluster( + ClusterMessage.newBuilder() + .setId(SERVER.getId()) + .setVersion(version) + .setChange(ClusterMessage.Change.newBuilder().setRemove(remove)))) + .build()); + } }