diff --git a/src/consumer/Consumer.java b/src/consumer/Consumer.java index c1b3a99..dfff32c 100644 --- a/src/consumer/Consumer.java +++ b/src/consumer/Consumer.java @@ -44,7 +44,7 @@ public class Consumer extends AbstractService { /** Version of the client library. Should be kept in sync with Go client version. */ - public static final String CLIENT_VERSION = "1.2.7"; + public static final String CLIENT_VERSION = "1.2.8"; private static final Logger LOGGER = LoggerFactory.getLogger(Consumer.class); @@ -59,8 +59,13 @@ public class Consumer extends AbstractService { private final BooleanSupplier verbose; private final Function, JoinMessage> registerFactory; + // Guards session start against shutdown, so shutdown never misses a newly started session. + private final Object sessionLock = new Object(); private @Nullable volatile Session session = null; - private ClusterMap cluster; + private boolean stopping = false; // guarded by sessionLock + + // Each update swaps in a new map; volatile lets readers see it fully built without locking. + private volatile ClusterMap cluster; private final Object lock = new Object(); /** Dials plaintext channels with the standard gRPC builder unless the application overrides it. */ @@ -204,40 +209,43 @@ protected void doStart() { State.RUNNING); private void connectLoop() { + var started = false; while (CONNECTABLE_STATES.contains(state())) { connectLimiter.acquire(); - var session = - Session.startAsync( - clock, - instance.toProto().getInstance(), - internalPool, - joinClient::join, - Messages::unwrapSessionMessage, - Messages::wrapSessionMessage, - this.verbose, - this::handleConsumerMessage); + Session session; + synchronized (sessionLock) { + if (stopping) { + return; + } + session = + Session.startAsync( + clock, + instance.toProto().getInstance(), + internalPool, + joinClient::join, + Messages::unwrapSessionMessage, + Messages::wrapSessionMessage, + this.verbose, + this::handleConsumerMessage); + // Queue Register before publishing, so a shutdown Deregister can only follow it. + session.sendAsync(registerFactory.apply(workPool.getStaleGrants())); + this.session = session; + } LOGGER.info("Started session {}", session); - var grants = workPool.getStaleGrants(); - session.sendAsync(registerFactory.apply(grants)); - session.awaitEstablishedOrTerminated(); if (session.isTerminated()) { LOGGER.info("Session {} is terminated before establishing. Reconnecting", session); + this.session = null; continue; } - this.session = session; - - // This loop can reconnect, but we only want to notify that we started once. We should only be - // in the STARTING state once. If state == STOPPING, we've been told to stop while we were - // starting. In this case, we'll announce that we've started, which should trigger - // termination after start finishes. - var state = state(); - if (state == State.STARTING || state == State.STOPPING) { + // Notify once: Guava defers a stop during startup until here, but not one during reconnect. + if (!started) { + started = true; notifyStarted(); } @@ -326,10 +334,15 @@ private void handleClusterMessage(ClusterMessage msg) { */ @Override protected void doStop() { - internalPool.submit(() -> terminate(session)); + internalPool.submit(this::terminate); } - private void terminate(@Nullable Session session) { + private void terminate() { + Session session; + synchronized (sessionLock) { + stopping = true; + session = this.session; + } if (session != null) { sendDeregister(session); } @@ -349,7 +362,8 @@ private void terminate(@Nullable Session session) { private void trySend(JoinMessage msg) { var session = this.session; - if (session == null) { + // The session is published before Established, drop it until it's established. + if (session == null || !session.isEstablished()) { LOGGER.warn("No active session. Ignoring message: {}", msg); return; } diff --git a/src/consumer/ConsumerTest.java b/src/consumer/ConsumerTest.java index 748b3a9..64bba13 100644 --- a/src/consumer/ConsumerTest.java +++ b/src/consumer/ConsumerTest.java @@ -169,6 +169,28 @@ public void reconnect(TestInfo testInfo) throws Exception { deps.stop(); } + @Test + public void stopWhileReconnecting(TestInfo testInfo) { + var deps = setupDeps(testInfo); + + deps.consumer.initialize(EXP); + + var g1 = Grant.create("g1", SHARD_D1_UNIT_0e, ALLOCATED, EXP, TS); + deps.assign(g1); + + deps.consumer.terminateConnection(); + deps.consumer.assertSentEstablish(); + + // Shutdown picks a session (revoking grants) before the replacement session is established. + deps.consumer.stopAsync(); + deps.grants.get("g1").revoked.await(); + + // The replacement session is still deregistered. + deps.consumer.sendEstablished(); + deps.consumer.assertSentRegister(g1); + deps.waitForOrderlyStop(); + } + @Test public void terminatedBeforeEstablished(TestInfo testInfo) { var deps = setupDeps(testInfo);