fix(fpss): evict in-flight subscribe correlation on unsubscribe so a superseded reject cannot drop a live sub (#883)

A per-contract or full-stream subscribe records a pending correlation keyed by its req_id so a rejecting REQ_RESPONSE can untrack exactly that subscription. The correlation is only an authority to untrack while the tracked entry it created is still live. An unsubscribe removed the tracked entry but left the correlation resident, so a subscribe -> unsubscribe -> re-subscribe of the same identity, all before the first subscribe's REQ_RESPONSE landed, left the first request's correlation pointing at a slot now occupied by a new live entry. A late rejection of that obsolete req_id then untracked the re-subscribed entry by value and dropped it from the reconnect-replay set.

Evict the in-flight correlation for an identity at the unsubscribe boundary, alongside removing the tracked entry. This restores the invariant that at most one correlation per (kind, contract) or (kind, sec_type) identity is resident and it always names the current live entry, so apply_req_response can untrack purely by req_id lookup and a rejection of a superseded request is a no-op rather than a value match against a newer entry.

Co-authored-by: preview <noreply@anthropic.com>

Author: userFRM

crates/thetadatadx/src/fpss/io_loop/mod.rs

@@ -140,6 +140,28 @@ pub(in crate::fpss) fn evict_stale_pending(map: &mut HashMap<i32, PendingSubEntr
         );
     }
 }
+
+/// Drop the in-flight subscribe correlation(s) for one tracked identity.
+///
+/// A pending correlation is only an authority to untrack while the tracked
+/// entry it created is still live. Once that entry leaves the tracked set — an
+/// unsubscribe removes it — the correlation points at nothing, so a later
+/// rejection of its (now superseded) `req_id` must not act on the set: the
+/// `(kind, contract)` slot may since have been re-subscribed into a new live
+/// entry that a value match would wrongly drop. Removing the correlation at the
+/// unsubscribe boundary keeps the invariant that at most one resident
+/// correlation per identity exists and it always names the current live entry,
+/// so an `apply_req_response` rejection can untrack purely by `req_id` lookup.
+///
+/// The map is keyed by `req_id`, so identity removal is a single retain pass.
+/// The caller holds the `pending_subs` lock; this touches only in-memory map
+/// state and performs no I/O.
+pub(in crate::fpss) fn evict_pending_for_identity(
+    map: &mut HashMap<i32, PendingSubEntry>,
+    identity: &super::protocol::PendingSub,
+) {
+    map.retain(|_, entry| &entry.sub != identity);
+}
 type ActiveFullSubs = Arc<
     Mutex<
         Vec<(
@@ -280,6 +302,16 @@ fn host_index(hosts: &[(String, u16)], addr: &str) -> Option<usize> {
 /// id from a span longer than the 31-bit counter cycle) leaves the tracked set
 /// untouched: with no correlation, untracking would risk dropping a healthy
 /// subscription, so the conservative choice is to keep it.
+///
+/// The `req_id` lookup is a safe untrack because the pending registry holds the
+/// invariant that at most one correlation per `(kind, contract)` (or `(kind,
+/// sec_type)`) identity is resident, and it always names the current live
+/// entry. A duplicate subscribe shares the live entry and registers no second
+/// correlation; an unsubscribe removes the entry and evicts its correlation
+/// (see [`evict_pending_for_identity`]). So a subscribe that is superseded by an
+/// unsubscribe + re-subscribe of the same identity has no resident correlation
+/// for its old `req_id`, and a late rejection of that id is a no-op rather than
+/// a value match that would drop the re-subscribed live entry.
 fn apply_req_response(
     req_id: i32,
     result: StreamResponseType,

crates/thetadatadx/src/fpss/mod.rs

@@ -2236,15 +2236,32 @@ impl StreamingClient {
             unsubscribe,
             "sent full-stream subscription frame"
         );
-        // Track / untrack for reconnection.
+        // Untrack-on-unsubscribe is terminal: remove the tracked entry, then
+        // evict the in-flight correlation for this identity for the same reason
+        // as the per-contract unsubscribe — the removed entry is no longer live,
+        // so a late rejection of its subscribe must not survive to untrack a
+        // future re-subscribe of the same `(kind, sec_type)` by value.
+        if unsubscribe {
+            self.active_full_subs
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .retain(|(k, s)| !(*k == kind_for_track && *s == sec_type));
+            io_loop::evict_pending_for_identity(
+                &mut self
+                    .pending_subs
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner),
+                &protocol::PendingSub::Full(kind_for_track, sec_type),
+            );
+            return Ok(());
+        }
+
+        // Track for reconnection.
         let mut subs = self
             .active_full_subs
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner);
-        let newly_tracked = if unsubscribe {
-            subs.retain(|(k, s)| !(*k == kind_for_track && *s == sec_type));
-            false
-        } else if subs
+        let newly_tracked = if subs
             .iter()
             .any(|(k, s)| *k == kind_for_track && *s == sec_type)
         {
@@ -2261,10 +2278,9 @@ impl StreamingClient {
         // Record the pending full-stream subscribe by `req_id` so a server
         // rejection drops exactly this entry from the replay set. Only the
         // subscribe that actually added the tracked entry may carry an
-        // untrack-capable correlation: an unsubscribe removed its entry above
-        // and has nothing to untrack, and a duplicate subscribe shares the one
-        // live entry, so letting its rejection untrack by value would drop the
-        // live subscription.
+        // untrack-capable correlation: an unsubscribe is handled terminally
+        // above and a duplicate subscribe shares the one live entry, so letting
+        // its rejection untrack by value would drop the live subscription.
         if newly_tracked {
             let mut pending = self
                 .pending_subs
@@ -2372,6 +2388,21 @@ impl StreamingClient {
             subs.retain(|(k, c)| !(k == &kind && c == contract));
         }
 
+        // Evict any in-flight correlation for this identity. The removed entry
+        // is no longer live, so a late rejection of the subscribe that created
+        // it must not survive to untrack a future re-subscribe of the same
+        // `(kind, contract)` by value.
+        {
+            let mut pending = self
+                .pending_subs
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            io_loop::evict_pending_for_identity(
+                &mut pending,
+                &protocol::PendingSub::Contract(kind, contract.clone()),
+            );
+        }
+
         tracing::debug!(
             req_id,
             kind = ?kind,
@@ -3425,6 +3456,140 @@ mod full_stream_guard_tests {
         client.shutdown();
     }
 
+    /// Subscribe, unsubscribe, then re-subscribe the same contract, all before
+    /// the first subscribe's `REQ_RESPONSE` lands; a late rejection of that
+    /// superseded first request must not drop the re-subscribed live entry.
+    ///
+    /// The first subscribe owns the correlation for its `req_id`; the
+    /// unsubscribe both removes the tracked entry and evicts that correlation,
+    /// so when the re-subscribe re-adds the entry it owns a fresh correlation
+    /// under a new `req_id`. A server rejection of the obsolete first `req_id`
+    /// then finds no resident correlation and is a no-op, leaving the live
+    /// re-subscribed entry tracked and in the reconnect-replay set (which is a
+    /// clone of `active_subs`).
+    #[test]
+    fn unsub_resub_then_reject_superseded_keeps_live_sub() {
+        use super::io_loop::apply_req_response_for_test;
+        use crate::tdbe::types::enums::StreamResponseType;
+
+        let client = StreamingClient::for_self_join_test(
+            0,
+            64,
+            HarnessPublishMode::BlockingPublish,
+            None,
+            |_event| {},
+        );
+
+        // req_id counter starts at 1: subscribe -> 1, unsubscribe -> 2,
+        // re-subscribe -> 3.
+        let contract = Contract::stock("AAPL");
+        client
+            .subscribe(contract.clone().trade())
+            .expect("first subscribe");
+        client
+            .unsubscribe(contract.clone().trade())
+            .expect("unsubscribe");
+        client
+            .subscribe(contract.clone().trade())
+            .expect("re-subscribe");
+
+        // Exactly one resident correlation: the unsubscribe evicted the first
+        // subscribe's (req_id 1) correlation, and the re-subscribe registered a
+        // fresh one (req_id 3).
+        assert_eq!(
+            client
+                .pending_subs
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .len(),
+            1,
+            "unsubscribe must evict the superseded correlation; only the \
+             re-subscribe's correlation may remain"
+        );
+
+        // The server rejects the obsolete first req_id (1). With its
+        // correlation already evicted this is a no-op, not a value match that
+        // would drop the live re-subscribed entry.
+        apply_req_response_for_test(
+            &client.pending_subs,
+            &client.active_subs,
+            &client.active_full_subs,
+            1,
+            StreamResponseType::MaxStreamsReached,
+        );
+
+        let tracked = client.active_subscriptions();
+        assert_eq!(
+            tracked.len(),
+            1,
+            "rejecting a superseded subscribe must leave the re-subscribed entry \
+             tracked, got {tracked:?}"
+        );
+        assert!(
+            tracked.iter().any(|(_, c)| *c == contract),
+            "the re-subscribed contract must remain in active_subscriptions() so \
+             it is replayed on reconnect, got {tracked:?}"
+        );
+
+        client.shutdown();
+    }
+
+    /// The full-stream analogue: subscribe, unsubscribe, re-subscribe a
+    /// full-stream `(kind, sec_type)`, then reject the superseded first
+    /// request — the re-subscribed full-stream entry must stay tracked.
+    #[test]
+    fn full_unsub_resub_then_reject_superseded_keeps_live_sub() {
+        use super::io_loop::apply_req_response_for_test;
+        use crate::tdbe::types::enums::StreamResponseType;
+
+        let client = StreamingClient::for_self_join_test(
+            0,
+            64,
+            HarnessPublishMode::BlockingPublish,
+            None,
+            |_event| {},
+        );
+
+        // subscribe -> 1, unsubscribe -> 2, re-subscribe -> 3.
+        client
+            .subscribe(SecType::Stock.full_trades())
+            .expect("first full subscribe");
+        client
+            .unsubscribe(SecType::Stock.full_trades())
+            .expect("full unsubscribe");
+        client
+            .subscribe(SecType::Stock.full_trades())
+            .expect("re-subscribe full");
+
+        assert_eq!(
+            client
+                .pending_subs
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .len(),
+            1,
+            "full-stream unsubscribe must evict the superseded correlation"
+        );
+
+        apply_req_response_for_test(
+            &client.pending_subs,
+            &client.active_subs,
+            &client.active_full_subs,
+            1,
+            StreamResponseType::MaxStreamsReached,
+        );
+
+        let tracked = client.active_full_subscriptions();
+        assert_eq!(
+            tracked.len(),
+            1,
+            "rejecting a superseded full-stream subscribe must leave the \
+             re-subscribed entry tracked, got {tracked:?}"
+        );
+
+        client.shutdown();
+    }
+
     /// The command channel is bounded: once it is saturated, a further
     /// `try_send` reports `Full` rather than growing without limit. This
     /// pins the backpressure contract `send_cmd` relies on to surface a