fix(offline): close the ordering gaps the terminal lock still left open

src/client.rs

@@ -1812,6 +1812,13 @@ pub struct Client {
     /// worker inside the flush and drive a concurrent generation change.
     #[cfg(test)]
     pub(crate) signal_flush_test_block: AtomicBool,
+    /// Set by `cleanup_connection_state` immediately before it takes
+    /// `offline_terminal_lock`. Lets a test prove the teardown reached the
+    /// lock rather than inferring it from elapsed scheduler turns: everything
+    /// the transition writes is on the far side of this point, so a test
+    /// holding the lock knows nothing has been reset yet when it fires.
+    #[cfg(test)]
+    pub(crate) offline_terminal_gate_reached: AtomicBool,
     /// Counts entries into the coalesced flush attempt, so a test can wait
     /// until a worker is actually inside the (blocked) flush.
     #[cfg(test)]

src/client/lifecycle.rs

@@ -543,6 +543,8 @@ impl Client {
             #[cfg(test)]
             signal_flush_test_block: AtomicBool::new(false),
             #[cfg(test)]
+            offline_terminal_gate_reached: AtomicBool::new(false),
+            #[cfg(test)]
             signal_flush_test_in_attempt: AtomicU32::new(0),
             #[cfg(test)]
             app_state_key_share_prepare_test_failures: AtomicU32::new(0),
@@ -1053,11 +1055,17 @@ impl Client {
         // refuse the login. `run` already clears it before each attempt; a
         // caller driving connections itself has nowhere else to learn of it.
         self.expected_disconnect.store(false, Ordering::Relaxed);
+        // Under the same lock as the teardown's resets: the previous
+        // connection's finisher is detached and can still be publishing, and
+        // its writes must not interleave with the state this attempt is
+        // establishing.
+        //
         // Runs before the flags it reads are cleared below. Teardown normally
         // reported the interruption already and this is a no-op; it covers the
         // paths that reach a new attempt without one. The current generation
         // is the right stamp here: this connection has not logged in yet, so
         // its own drain arms at a higher one.
+        let terminal_gate = self.offline_terminal_lock.lock().await;
         self.abandon_offline_sync_if_interrupted(
             self.connection_generation.load(Ordering::Acquire),
         );
@@ -1079,6 +1087,7 @@ impl Client {
             self.signal_cache.clear().await;
         }
         self.offline_batch.reset();
+        drop(terminal_gate);
         self.outbound_flush.reopen();
 
         // WA Web: both MQTT and DGW transports use a 20s connect timeout.
@@ -1842,20 +1851,26 @@ impl Client {
             );
             self.signal_cache.clear().await;
         }
+        // Everything the drain-to-live transition also writes goes under the
+        // lock the finisher publishes beneath, starting with the permit count:
+        // whichever of the two wins the stamp, its writes land wholly before
+        // or wholly after the other's. Outside it, a finisher that had already
+        // claimed its stamp could widen the semaphore back to 64 after this
+        // reset, and the next connection would drain its backlog concurrently
+        // with no permit serializing the Signal state.
+        #[cfg(test)]
+        self.offline_terminal_gate_reached
+            .store(true, Ordering::Release);
+        let terminal_gate = self.offline_terminal_lock.lock().await;
         // Reset semaphore to 1 permit for next offline sync.
         self.swap_message_semaphore(1);
         // Reset dead-socket timestamps so stale values from the previous
         // connection don't trigger an immediate reconnect on the next one.
         self.stats.reset_connection_activity();
         self.pending_device_sync.clear();
-        // Reset offline sync state for next connection, under the lock the
-        // finisher publishes beneath: whichever of the two wins the stamp, its
-        // writes land wholly before or wholly after the other's, so a widened
-        // semaphore can never survive into the next connection's drain.
         // The report is stamped with the generation being retired, not the one
         // just installed, so it silences that drain's own stale finisher
         // without claiming the slot the next drain will need.
-        let terminal_gate = self.offline_terminal_lock.lock().await;
         self.abandon_offline_sync_if_interrupted(closed_generation);
         self.offline_sync_completed.store(false, Ordering::Relaxed);
         self.clear_offline_receipt_buffer();

src/client/sessions.rs

@@ -417,6 +417,17 @@ impl Client {
     /// widens the semaphore then. `None` (upgrade-failure fallback) leaves
     /// the buffer alone for the connection-state reset to clear.
     fn publish_offline_sync_live_state(&self, count: i32, durable: Option<bool>, generation: u64) {
+        // Re-read under `offline_terminal_lock`, which every caller holds: the
+        // check each of them made before taking it could have been overtaken
+        // by the teardown waiting for that same lock.
+        if self.connection_generation.load(Ordering::Acquire) != generation {
+            log::debug!(
+                target: "Client/OfflineSync",
+                "Generation {} was retired before its drain could publish; leaving live state alone",
+                generation,
+            );
+            return;
+        }
         // The claim is the serialization point for the whole transition, not
         // just its event. Losing it means a teardown already reported this
         // drain's end, or a newer drain overtook it: either way that teardown

src/client/tests.rs

@@ -7585,6 +7585,104 @@ async fn wait_for_startup_sync_reports_a_teardown_without_waiting_out_its_timeou
     );
 }
 
+/// `async_lock::Semaphore` does not report its count, so ask it the question
+/// the drain actually asks: can two stanzas be in flight at once?
+fn concurrent_permits(client: &Arc<Client>) -> bool {
+    let semaphore = client.read_message_semaphore().1;
+    let first = semaphore.try_acquire_arc();
+    let second = semaphore.try_acquire_arc();
+    first.is_some() && second.is_some()
+}
+
+/// Exactly one, not merely "not two": a drain that could acquire nothing at
+/// all would satisfy the negation of [`concurrent_permits`] while being just
+/// as broken.
+fn exactly_one_permit(client: &Arc<Client>) -> bool {
+    let semaphore = client.read_message_semaphore().1;
+    let first = semaphore.try_acquire_arc();
+    let second = semaphore.try_acquire_arc();
+    first.is_some() && second.is_none()
+}
+
+/// A finisher that arrives after its teardown cannot widen the semaphore the
+/// next drain needs narrow.
+///
+/// This is the sequential half of the invariant: the finisher's own late
+/// publication is refused. The contended half, that the reset itself sits
+/// inside the lock, is [`the_permit_reset_is_inside_the_terminal_lock`].
+#[tokio::test]
+async fn a_late_finisher_cannot_widen_the_next_drains_semaphore() {
+    let client = offline_resume_test_client().await;
+
+    arm_offline_drain(&client, 711, 700).await;
+    let stale_generation = client.connection_generation.load(Ordering::Acquire);
+    client.enter_live_mode_for_tests();
+    assert!(
+        concurrent_permits(&client),
+        "live mode is the wide semaphore"
+    );
+
+    client.cleanup_connection_state().await;
+
+    // The finisher of the retired drain runs late and finds its slot taken.
+    client
+        .complete_offline_sync_for_generation(711, stale_generation)
+        .await;
+
+    assert!(
+        exactly_one_permit(&client),
+        "the next drain starts on exactly one permit, whatever the old finisher does"
+    );
+}
+
+/// The permit reset happens under `offline_terminal_lock`, not before it.
+///
+/// Holding the lock is the only way to observe that from outside: while it is
+/// held, a teardown must not have narrowed the semaphore, because the reset is
+/// on the far side of the same lock a publishing finisher holds. With the
+/// reset moved back out, the teardown reaches it without waiting and the
+/// assertion below fails.
+///
+/// Not vacuous, because it waits on `offline_terminal_gate_reached` rather
+/// than on elapsed scheduler turns: that flag is set on the line above the
+/// acquisition, so the teardown has provably arrived at the lock by the time
+/// the assertion runs.
+#[tokio::test]
+async fn the_permit_reset_is_inside_the_terminal_lock() {
+    let client = offline_resume_test_client().await;
+
+    arm_offline_drain(&client, 711, 700).await;
+    client.enter_live_mode_for_tests();
+
+    let terminal_gate = client.offline_terminal_lock.lock().await;
+
+    let teardown = tokio::spawn({
+        let client = Arc::clone(&client);
+        async move { client.cleanup_connection_state().await }
+    });
+
+    // Wait for the teardown to reach the lock itself, not for a guess at how
+    // many scheduler turns that takes: the flag is set on the line above the
+    // acquisition, so when it fires the reset is still ahead of the lock the
+    // test holds. Move the reset back out and it has already run by then.
+    crate::test_utils::poll_until("the teardown to reach the terminal lock", || {
+        client.offline_terminal_gate_reached.load(Ordering::Acquire)
+    })
+    .await;
+
+    assert!(
+        concurrent_permits(&client),
+        "the teardown narrowed the semaphore without the lock a finisher publishes under"
+    );
+
+    drop(terminal_gate);
+    teardown.await.expect("teardown must not panic");
+    assert!(
+        exactly_one_permit(&client),
+        "and once it has the lock, it does narrow it to exactly one"
+    );
+}
+
 /// A completion belongs to the connection whose drain it is.
 ///
 /// The inactivity watchdog and the offline-delivery waiter both check the