fix: Detect clean-session reconnects in Subscription::next_message

src/client.rs

@@ -122,25 +122,6 @@ impl<'a, M: RawMutex> MqttClient<'a, M> {
         .await
     }
 
-    /// Cleans the session by waiting for the connection state to change.
-    async fn clean_session(&self) {
-        let mut prev_state = self.shared.lock().await.connected;
-
-        poll_fn(|cx| {
-            if let Ok(mut shared) = self.shared.try_lock() {
-                if let (None | Some(true), Some(false)) = (prev_state, shared.connected) {
-                    return Poll::Ready(());
-                }
-                prev_state = shared.connected;
-                shared.connection_waker.register(cx.waker());
-                return Poll::Pending;
-            }
-            cx.waker().wake_by_ref();
-            Poll::Pending
-        })
-        .await;
-    }
-
     /// Enqueues an MQTT packet into the transmit queue, to be sent by the [`MqttStack`].
     ///
     /// # Parameters
@@ -354,10 +335,13 @@ impl<'a, M: RawMutex> MqttClient<'a, M> {
             .map(|sub| TopicFilter::new(sub.topic_path))
             .collect::<Result<_, _>>()?;
 
+        let created_clean_session_count = self.shared.lock().await.clean_session_count;
+
         Ok(Subscription {
             subscriber,
             topic_filters,
             client: self,
+            created_clean_session_count,
         })
     }
 
@@ -477,6 +461,14 @@ pub struct Subscription<'a, 'b, M: RawMutex, const MAX_TOPICS: usize> {
     topic_filters: heapless::Vec<TopicFilter, MAX_TOPICS>,
     subscriber: FrameSubscriber<'a, M, SliceBufferProvider<'a>, MAX_SUBSCRIBERS>,
     client: &'b MqttClient<'a, M>,
+    /// Snapshot of `Shared::clean_session_count` at subscribe time.
+    /// Compared in `poll_next`; if the live counter has advanced, the broker
+    /// has reported `session_present=false` since this subscription was
+    /// created — meaning it has dropped our subscriptions — and the wait
+    /// returns `None`. Transient disconnects and session-resume reconnects
+    /// do not advance the counter, so quick reconnects keep subscriptions
+    /// alive.
+    created_clean_session_count: u8,
 }
 
 impl<'a, 'b, M: RawMutex, const MAX_TOPICS: usize> Subscription<'a, 'b, M, MAX_TOPICS> {
@@ -567,15 +559,24 @@ impl<'a, 'b, M: RawMutex, const MAX_TOPICS: usize> futures_util::Stream
             subscriber,
             topic_filters,
             client,
+            created_clean_session_count,
         } = self.get_mut();
 
-        // Check for clean session first - if clean session occurred, return None
-        // to signal the subscription is no longer valid
-        let clean_session_future = client.clean_session();
-        futures_util::pin_mut!(clean_session_future);
-
-        if clean_session_future.poll(cx).is_ready() {
-            return Poll::Ready(None);
+        // Return None if the broker has dropped our subs since this
+        // Subscription was created. `clean_session_count` is bumped by
+        // `set_connected` only on `session_present=false`; transient
+        // disconnects and session-resume reconnects leave subs valid.
+        // Register the waker before checking so a concurrent
+        // `set_connected` can't slip through.
+        if let Ok(mut shared) = client.shared.try_lock() {
+            shared.connection_state_change_waker.register(cx.waker());
+            if shared.clean_session_count != *created_clean_session_count {
+                return Poll::Ready(None);
+            }
+        } else {
+            // Lock contended (likely `set_connected` running) — retry.
+            cx.waker().wake_by_ref();
+            return Poll::Pending;
         }
 
         match subscriber.read_with_context(Some(cx)) {

src/state.rs

@@ -31,6 +31,17 @@ pub struct Shared {
     /// - `None` if disconnected
     pub(crate) connected: Option<bool>,
 
+    /// Monotonic counter incremented only when a reconnect lands with
+    /// `session_present=false` (clean session). `Subscription` snapshots
+    /// this on creation so `poll_next` can detect that the broker has
+    /// dropped our subscriptions and signal `None` to wake parked waiters.
+    /// Transient disconnects and session-resume reconnects do not bump
+    /// this counter — both broker and local subscriber state survive, so
+    /// the subscription stays valid. This survives `Subscription`'s
+    /// lifetime (it lives in the long-lived `Shared`), unlike a future-local
+    /// `prev_state`, which `poll_next` can't preserve across re-creation.
+    pub(crate) clean_session_count: u8,
+
     /// Outgoing QoS 1, 2 publishes which aren't acked yet
     pub(crate) inflight_pub: FnvIndexSet<u16, MAX_INFLIGHT>,
     pub(crate) ack_status: FnvIndexMap<u16, AckStatus, { MAX_SUBSCRIBERS }>,
@@ -55,6 +66,7 @@ impl Shared {
             connection_waker: WakerRegistration::new(),
             connection_state_change_waker: WakerRegistration::new(),
             connected: None,
+            clean_session_count: 0,
 
             inflight_pub: heapless::IndexSet::new(),
             ack_status: IndexMap::new(),
@@ -101,6 +113,16 @@ impl Shared {
 
         // Wake state change waker if connection state actually changed
         if previous != connected {
+            // Only count clean-session reconnects. A transient disconnect
+            // (`None`) or a session-resume reconnect (`Some(true)`) leaves
+            // both the broker's and our local subscriber state intact, so
+            // active `Subscription`s stay valid. Bumping the counter only
+            // on `Some(false)` (broker reports session_present=false)
+            // means quick reconnects no longer spuriously invalidate
+            // subscriptions.
+            if connected == Some(false) {
+                self.clean_session_count = self.clean_session_count.wrapping_add(1);
+            }
             self.connection_state_change_waker.wake();
         }
     }