Deregister reconnecting sessions on consumer shutdown

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<Stream<Grant>, 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<JoinMessage, JoinMessage> 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<JoinMessage, JoinMessage> 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<JoinMessage, JoinMessage> session) {
+  private void terminate() {
+    Session<JoinMessage, JoinMessage> session;
+    synchronized (sessionLock) {
+      stopping = true;
+      session = this.session;
+    }
     if (session != null) {
       sendDeregister(session);
     }
@@ -349,7 +362,8 @@ private void terminate(@Nullable Session<JoinMessage, JoinMessage> 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;
     }

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);