Skip to main content

nautilus_network/websocket/
subscription.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Adapter-managed subscription intent and acknowledgment tracking.
17//!
18//! [`SubscriptionState`] keeps confirmed, `pending_subscribe`, and `pending_unsubscribe` topics
19//! separate. Server acknowledgments move topics between those states; late subscribe
20//! acknowledgments do not revive cancelled intent, and stale unsubscribe acknowledgments do not
21//! remove a later resubscription. [`SubscriptionState::all_topics`] returns confirmed and pending
22//! subscribe intent for recovery, excluding pending unsubscriptions.
23//!
24//! Reference counts are independent of acknowledgment state. The first reference tells the caller
25//! to send a subscribe request, and removing the last tells it to send an unsubscribe request. The
26//! tracker records state but never sends protocol messages.
27//!
28//! # Topic format
29//!
30//! Topics use `channel{delimiter}symbol`; the first delimiter occurrence separates the channel
31//! from the optional symbol. A topic without a delimiter represents the whole channel.
32
33use std::{
34    num::NonZeroUsize,
35    ops::Deref,
36    sync::{Arc, LazyLock},
37};
38
39use ahash::{AHashMap, AHashSet};
40use dashmap::DashMap;
41use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
42use ustr::Ustr;
43
44/// Marker for channel-level subscriptions (no specific symbol).
45///
46/// An empty string in the symbol set indicates a channel-level subscription
47/// that applies to all symbols for that channel.
48pub(crate) static CHANNEL_LEVEL_MARKER: LazyLock<Ustr> = LazyLock::new(|| Ustr::from(""));
49
50/// Read-only snapshot of subscription topics grouped by channel.
51///
52/// The snapshot supports immutable map operations through [`Deref`], but cannot mutate the
53/// underlying [`SubscriptionState`]. Use the subscription lifecycle methods to change state.
54///
55/// ```compile_fail
56/// use nautilus_network::websocket::SubscriptionState;
57///
58/// let state = SubscriptionState::new('.');
59/// let mut confirmed = state.confirmed();
60/// confirmed.clear();
61/// ```
62#[derive(Clone, Debug, Default, Eq, PartialEq)]
63pub struct SubscriptionSnapshot(AHashMap<Ustr, AHashSet<Ustr>>);
64
65impl Deref for SubscriptionSnapshot {
66    type Target = AHashMap<Ustr, AHashSet<Ustr>>;
67
68    fn deref(&self) -> &Self::Target {
69        &self.0
70    }
71}
72
73/// Tracks subscription intent and acknowledgment state for WebSocket connections.
74///
75/// # State management
76///
77/// The tracker maintains desired subscription intent and three acknowledgment states:
78///
79/// - **Desired**: Topics that should be active, independent of response ordering or connection.
80/// - **Confirmed**: Subscriptions acknowledged by the server and expected to stream data.
81/// - **Pending subscribe**: Subscribe requests awaiting server acknowledgment.
82/// - **Pending unsubscribe**: Unsubscribe requests awaiting server acknowledgment.
83///
84/// Late subscribe acknowledgments do not revive cancelled intent, and stale unsubscribe
85/// acknowledgments do not remove a later resubscription.
86///
87/// # Reference counting
88///
89/// Reference counts remain independent of acknowledgment state. The first consumer tells the
90/// caller to send a subscribe request, while removing the last tells it to send an unsubscribe
91/// request. The tracker records these transitions but does not send protocol messages.
92///
93/// # Topic format
94///
95/// Topics use `channel{delimiter}symbol`, with delimiters such as `.` or `:`. A topic without the
96/// delimiter represents a channel-level subscription.
97///
98/// # Thread safety
99///
100/// Clones share all state. Operations are thread-safe and can run concurrently from multiple
101/// tasks.
102#[derive(Clone, Debug)]
103pub struct SubscriptionState {
104    confirmed: Arc<DashMap<Ustr, AHashSet<Ustr>>>,
105    pending_subscribe: Arc<DashMap<Ustr, AHashSet<Ustr>>>,
106    pending_unsubscribe: Arc<DashMap<Ustr, AHashSet<Ustr>>>,
107    desired: Arc<DashMap<Ustr, AHashSet<Ustr>>>,
108    reference_counts: Arc<DashMap<Ustr, NonZeroUsize>>,
109    state_lock: Arc<RwLock<()>>,
110    delimiter: char,
111}
112
113impl SubscriptionState {
114    /// Creates a new subscription state tracker with the specified topic delimiter.
115    #[must_use]
116    pub fn new(delimiter: char) -> Self {
117        Self {
118            confirmed: Arc::new(DashMap::new()),
119            pending_subscribe: Arc::new(DashMap::new()),
120            pending_unsubscribe: Arc::new(DashMap::new()),
121            desired: Arc::new(DashMap::new()),
122            reference_counts: Arc::new(DashMap::new()),
123            state_lock: Arc::new(RwLock::new(())),
124            delimiter,
125        }
126    }
127
128    /// Returns the delimiter character used for topic splitting.
129    #[must_use]
130    pub fn delimiter(&self) -> char {
131        self.delimiter
132    }
133
134    /// Returns a read-only snapshot of confirmed subscriptions.
135    #[must_use]
136    pub fn confirmed(&self) -> SubscriptionSnapshot {
137        let _guard = self.lock_state_read();
138        snapshot(&self.confirmed)
139    }
140
141    /// Returns a read-only snapshot of pending subscriptions.
142    #[must_use]
143    pub fn pending_subscribe(&self) -> SubscriptionSnapshot {
144        let _guard = self.lock_state_read();
145        snapshot(&self.pending_subscribe)
146    }
147
148    /// Returns a read-only snapshot of pending unsubscriptions.
149    #[must_use]
150    pub fn pending_unsubscribe(&self) -> SubscriptionSnapshot {
151        let _guard = self.lock_state_read();
152        snapshot(&self.pending_unsubscribe)
153    }
154
155    /// Returns the number of confirmed subscriptions.
156    ///
157    /// Counts both channel-level and symbol-level subscriptions.
158    #[must_use]
159    pub fn len(&self) -> usize {
160        let _guard = self.lock_state_read();
161        self.confirmed.iter().map(|entry| entry.value().len()).sum()
162    }
163
164    /// Returns true if there are no subscriptions (confirmed or pending).
165    #[must_use]
166    pub fn is_empty(&self) -> bool {
167        let _guard = self.lock_state_read();
168        self.confirmed.is_empty()
169            && self.pending_subscribe.is_empty()
170            && self.pending_unsubscribe.is_empty()
171            && self.desired.is_empty()
172    }
173
174    /// Returns true if a channel:symbol pair is subscribed (confirmed or pending subscribe).
175    #[must_use]
176    pub fn is_subscribed(&self, channel: &Ustr, symbol: &Ustr) -> bool {
177        let _guard = self.lock_state_read();
178
179        if let Some(symbols) = self.confirmed.get(channel)
180            && symbols.contains(symbol)
181        {
182            return true;
183        }
184
185        if let Some(symbols) = self.pending_subscribe.get(channel)
186            && symbols.contains(symbol)
187        {
188            return true;
189        }
190        false
191    }
192
193    /// Returns all pending subscribe topics as strings.
194    #[must_use]
195    pub fn pending_subscribe_topics(&self) -> Vec<String> {
196        let _guard = self.lock_state_read();
197        self.topics_from_map(&self.pending_subscribe)
198    }
199
200    /// Returns all pending unsubscribe topics as strings.
201    #[must_use]
202    pub fn pending_unsubscribe_topics(&self) -> Vec<String> {
203        let _guard = self.lock_state_read();
204        self.topics_from_map(&self.pending_unsubscribe)
205    }
206
207    /// Returns all topics that should be active after reconnect recovery.
208    ///
209    /// The result includes confirmed and pending subscribe topics, but excludes pending
210    /// unsubscribe topics.
211    #[must_use]
212    pub fn all_topics(&self) -> Vec<String> {
213        let _guard = self.lock_state_read();
214        let mut topics = self.topics_from_map(&self.confirmed);
215        topics.extend(self.topics_from_map(&self.pending_subscribe));
216        topics
217    }
218
219    /// Marks a topic as pending subscription.
220    ///
221    /// Call this after sending a subscribe request. This operation is idempotent for a confirmed
222    /// topic and cancels any pending unsubscription for the same topic.
223    pub fn mark_subscribe(&self, topic: &str) {
224        let _guard = self.lock_state_write();
225        let (channel, symbol) = split_topic(topic, self.delimiter);
226        track_topic(&self.desired, channel, symbol);
227
228        // If already confirmed, don't re-add to pending (idempotent)
229        if is_tracked(&self.confirmed, channel, symbol) {
230            return;
231        }
232
233        // Remove from pending_unsubscribe if present
234        untrack_topic(&self.pending_unsubscribe, channel, symbol);
235
236        track_topic(&self.pending_subscribe, channel, symbol);
237    }
238
239    /// Atomically tries to mark a topic as pending subscription.
240    ///
241    /// Returns `true` if the topic was newly marked as pending (should send subscribe).
242    /// Returns `false` if the topic was already confirmed or pending (skip sending).
243    ///
244    /// The check and state transition are atomic across concurrent subscribe calls.
245    #[must_use]
246    pub fn try_mark_subscribe(&self, topic: &str) -> bool {
247        let _guard = self.lock_state_write();
248        let (channel, symbol) = split_topic(topic, self.delimiter);
249
250        // If already desired, no action needed
251        if !track_topic(&self.desired, channel, symbol) {
252            return false;
253        }
254
255        let inserted = track_topic(&self.pending_subscribe, channel, symbol);
256        if inserted {
257            untrack_topic(&self.pending_unsubscribe, channel, symbol);
258        }
259
260        inserted
261    }
262
263    /// Confirms a subscription by moving it from pending to confirmed.
264    ///
265    /// Call this when the server acknowledges a subscribe request. A late confirmation cannot
266    /// restore a topic that is no longer desired.
267    pub fn confirm_subscribe(&self, topic: &str) {
268        let _guard = self.lock_state_write();
269        let (channel, symbol) = split_topic(topic, self.delimiter);
270
271        if !is_tracked(&self.desired, channel, symbol)
272            || is_tracked(&self.pending_unsubscribe, channel, symbol)
273        {
274            return;
275        }
276
277        untrack_topic(&self.pending_subscribe, channel, symbol);
278        track_topic(&self.confirmed, channel, symbol);
279    }
280
281    /// Marks a topic as pending unsubscription.
282    ///
283    /// Removes the topic from confirmed and `pending_subscribe` state before adding it to
284    /// `pending_unsubscribe`. This also handles unsubscription before initial confirmation.
285    pub fn mark_unsubscribe(&self, topic: &str) {
286        let _guard = self.lock_state_write();
287        let (channel, symbol) = split_topic(topic, self.delimiter);
288        untrack_topic(&self.desired, channel, symbol);
289        track_topic(&self.pending_unsubscribe, channel, symbol);
290        untrack_topic(&self.confirmed, channel, symbol);
291        untrack_topic(&self.pending_subscribe, channel, symbol);
292    }
293
294    /// Confirms an unsubscription by removing it from pending and confirmed state.
295    ///
296    /// Call this when the server acknowledges an unsubscribe request. A stale acknowledgment is
297    /// ignored if the topic is no longer pending unsubscription. `pending_subscribe` remains intact
298    /// so an immediate resubscription survives a late unsubscribe acknowledgment.
299    pub fn confirm_unsubscribe(&self, topic: &str) {
300        let _guard = self.lock_state_write();
301        let (channel, symbol) = split_topic(topic, self.delimiter);
302
303        // Only process if topic is actually pending unsubscription
304        // This ignores stale unsubscribe ACKs after user has re-subscribed
305        if !is_tracked(&self.pending_unsubscribe, channel, symbol) {
306            return; // Stale ACK, ignore
307        }
308
309        untrack_topic(&self.pending_unsubscribe, channel, symbol);
310        untrack_topic(&self.confirmed, channel, symbol);
311        // Don't clear pending_subscribe - it's a valid re-subscribe request
312    }
313
314    /// Marks a subscription as failed, moving it from confirmed back to pending.
315    ///
316    /// This keeps failed subscriptions available for retry after reconnect. A topic pending
317    /// unsubscription is unchanged because its subscription was cancelled.
318    pub fn mark_failure(&self, topic: &str) {
319        let _guard = self.lock_state_write();
320        let (channel, symbol) = split_topic(topic, self.delimiter);
321
322        if !is_tracked(&self.desired, channel, symbol)
323            || is_tracked(&self.pending_unsubscribe, channel, symbol)
324        {
325            return;
326        }
327
328        untrack_topic(&self.confirmed, channel, symbol);
329        track_topic(&self.pending_subscribe, channel, symbol);
330    }
331
332    /// Resets acknowledgment state for a replacement connection.
333    ///
334    /// Returns the desired topics to replay. Confirmed topics become pending subscriptions,
335    /// pending unsubscriptions are completed by the closed connection, and reference counts are
336    /// preserved.
337    #[allow(
338        clippy::must_use_candidate,
339        reason = "some adapters replay from separate subscription registries"
340    )]
341    pub fn reset_after_reconnect(&self) -> Vec<String> {
342        let _guard = self.lock_state_write();
343        let mut topics = self.topics_from_map(&self.confirmed);
344        topics.extend(self.topics_from_map(&self.pending_subscribe));
345
346        self.confirmed.clear();
347        self.pending_subscribe.clear();
348        self.pending_unsubscribe.clear();
349
350        for topic in &topics {
351            let (channel, symbol) = split_topic(topic, self.delimiter);
352            track_topic(&self.pending_subscribe, channel, symbol);
353        }
354
355        topics
356    }
357
358    /// Increments the reference count for a topic.
359    ///
360    /// Returns `true` if this is the first subscription (caller should send subscribe
361    /// message to server).
362    ///
363    /// # Panics
364    ///
365    /// Panics if the reference count exceeds `usize::MAX` subscriptions for a single topic.
366    #[allow(
367        clippy::must_use_candidate,
368        reason = "callers use this for side effects"
369    )]
370    pub fn add_reference(&self, topic: &str) -> bool {
371        let mut should_subscribe = false;
372        let topic_ustr = Ustr::from(topic);
373
374        self.reference_counts
375            .entry(topic_ustr)
376            .and_modify(|count| {
377                *count = NonZeroUsize::new(count.get() + 1).expect("reference count overflow");
378            })
379            .or_insert_with(|| {
380                should_subscribe = true;
381                NonZeroUsize::MIN
382            });
383
384        should_subscribe
385    }
386
387    /// Decrements the reference count for a topic.
388    ///
389    /// Returns `true` if this was the last subscription (caller should send unsubscribe
390    /// message to server).
391    ///
392    /// # Panics
393    ///
394    /// Panics if the internal reference count state becomes inconsistent (should never happen
395    /// if the API is used correctly).
396    #[allow(
397        clippy::must_use_candidate,
398        reason = "callers use this for side effects"
399    )]
400    pub fn remove_reference(&self, topic: &str) -> bool {
401        let topic_ustr = Ustr::from(topic);
402
403        // Use entry API to atomically decrement and remove if zero
404        // This prevents race where another thread adds a reference between the check and remove
405        if let dashmap::mapref::entry::Entry::Occupied(mut entry) =
406            self.reference_counts.entry(topic_ustr)
407        {
408            let current = entry.get().get();
409
410            if current == 1 {
411                entry.remove();
412                return true;
413            }
414
415            *entry.get_mut() = NonZeroUsize::new(current - 1)
416                .expect("reference count should never reach zero here");
417        }
418
419        false
420    }
421
422    /// Returns the current reference count for a topic.
423    ///
424    /// Returns 0 if the topic has no references.
425    #[must_use]
426    pub fn get_reference_count(&self, topic: &str) -> usize {
427        let topic_ustr = Ustr::from(topic);
428        self.reference_counts
429            .get(&topic_ustr)
430            .map_or(0, |count| count.get())
431    }
432
433    /// Clears all subscription state.
434    ///
435    /// This resets desired intent, acknowledgment state, and reference counts.
436    pub fn clear(&self) {
437        let _guard = self.lock_state_write();
438        self.confirmed.clear();
439        self.pending_subscribe.clear();
440        self.pending_unsubscribe.clear();
441        self.desired.clear();
442        self.reference_counts.clear();
443    }
444
445    // Converts a subscription map to sorted topic strings
446    fn topics_from_map(&self, map: &DashMap<Ustr, AHashSet<Ustr>>) -> Vec<String> {
447        let mut topics = Vec::new();
448        let marker = *CHANNEL_LEVEL_MARKER;
449
450        for entry in map {
451            let channel = entry.key();
452            let symbols = entry.value();
453
454            // Check for channel-level subscription marker
455            if symbols.contains(&marker) {
456                topics.push(channel.to_string());
457            }
458
459            // Add symbol-level subscriptions (skip marker)
460            for symbol in symbols {
461                if *symbol != marker {
462                    topics.push(format!("{channel}{}{symbol}", self.delimiter));
463                }
464            }
465        }
466
467        // Sort so resubscription after a reconnect replays topics in the same sequence
468        // across runs; both the outer DashMap and the inner symbol sets are unordered.
469        topics.sort();
470        topics
471    }
472
473    fn lock_state_read(&self) -> RwLockReadGuard<'_, ()> {
474        self.state_lock.read()
475    }
476
477    fn lock_state_write(&self) -> RwLockWriteGuard<'_, ()> {
478        self.state_lock.write()
479    }
480}
481
482/// Splits a topic into channel and optional symbol using the specified delimiter.
483#[must_use]
484pub fn split_topic(topic: &str, delimiter: char) -> (&str, Option<&str>) {
485    topic
486        .split_once(delimiter)
487        .map_or((topic, None), |(channel, symbol)| (channel, Some(symbol)))
488}
489
490fn snapshot(map: &DashMap<Ustr, AHashSet<Ustr>>) -> SubscriptionSnapshot {
491    SubscriptionSnapshot(
492        map.iter()
493            .map(|entry| (*entry.key(), entry.value().clone()))
494            .collect(),
495    )
496}
497
498/// Tracks a topic in the given map by adding it to the channel's symbol set.
499///
500/// Channel-level subscriptions are stored using an empty string marker,
501/// allowing both channel-level and symbol-level subscriptions to coexist.
502fn track_topic(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) -> bool {
503    map.entry(Ustr::from(channel))
504        .or_default()
505        .insert(topic_symbol(symbol))
506}
507
508/// Removes a topic from the given map by removing it from the channel's symbol set.
509///
510/// Removes the entire channel entry if no subscriptions remain after removal.
511fn untrack_topic(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) {
512    let symbol = topic_symbol(symbol);
513
514    // Use entry API to atomically remove symbol and check if empty
515    // This prevents race conditions where another thread adds a symbol between operations
516    if let dashmap::mapref::entry::Entry::Occupied(mut entry) = map.entry(Ustr::from(channel)) {
517        entry.get_mut().remove(&symbol);
518        if entry.get().is_empty() {
519            entry.remove();
520        }
521    }
522}
523
524/// Checks if a topic exists in the given map.
525fn is_tracked(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) -> bool {
526    let symbol = topic_symbol(symbol);
527    map.get(&Ustr::from(channel))
528        .is_some_and(|entry| entry.contains(&symbol))
529}
530
531fn topic_symbol(symbol: Option<&str>) -> Ustr {
532    symbol.map_or(*CHANNEL_LEVEL_MARKER, Ustr::from)
533}
534
535#[cfg(test)]
536mod tests {
537    use std::sync::{
538        Barrier,
539        atomic::{AtomicUsize, Ordering},
540    };
541
542    use rstest::rstest;
543
544    use super::*;
545
546    #[rstest]
547    fn test_split_topic_with_symbol() {
548        let (channel, symbol) = split_topic("tickers.BTCUSDT", '.');
549        assert_eq!(channel, "tickers");
550        assert_eq!(symbol, Some("BTCUSDT"));
551
552        let (channel, symbol) = split_topic("orderBookL2:XBTUSD", ':');
553        assert_eq!(channel, "orderBookL2");
554        assert_eq!(symbol, Some("XBTUSD"));
555    }
556
557    #[rstest]
558    fn test_split_topic_without_symbol() {
559        let (channel, symbol) = split_topic("orderbook", '.');
560        assert_eq!(channel, "orderbook");
561        assert_eq!(symbol, None);
562    }
563
564    #[rstest]
565    fn test_topic_without_symbol() {
566        let state = SubscriptionState::new('.');
567        state.mark_subscribe("orderbook");
568        state.confirm_subscribe("orderbook");
569
570        assert_eq!(state.len(), 1);
571        assert_eq!(state.all_topics(), vec!["orderbook"]);
572    }
573
574    #[rstest]
575    fn test_different_delimiters() {
576        let state_dot = SubscriptionState::new('.');
577        state_dot.mark_subscribe("tickers.BTCUSDT");
578        assert_eq!(
579            state_dot.pending_subscribe_topics(),
580            vec!["tickers.BTCUSDT"]
581        );
582
583        let state_colon = SubscriptionState::new(':');
584        state_colon.mark_subscribe("orderBookL2:XBTUSD");
585        assert_eq!(
586            state_colon.pending_subscribe_topics(),
587            vec!["orderBookL2:XBTUSD"]
588        );
589    }
590
591    #[rstest]
592    fn test_different_delimiter_does_not_affect_storage() {
593        // Verify delimiter is only used for parsing, not storage
594        let state_dot = SubscriptionState::new('.');
595        let state_colon = SubscriptionState::new(':');
596
597        // Add same logical subscription with different delimiters
598        state_dot.mark_subscribe("channel.SYMBOL");
599        state_colon.mark_subscribe("channel:SYMBOL");
600
601        // Both should work correctly
602        assert_eq!(state_dot.pending_subscribe_topics(), vec!["channel.SYMBOL"]);
603        assert_eq!(
604            state_colon.pending_subscribe_topics(),
605            vec!["channel:SYMBOL"]
606        );
607    }
608
609    #[rstest]
610    fn test_multiple_symbols_same_channel() {
611        let state = SubscriptionState::new('.');
612        state.mark_subscribe("tickers.BTCUSDT");
613        state.mark_subscribe("tickers.ETHUSDT");
614        state.confirm_subscribe("tickers.BTCUSDT");
615        state.confirm_subscribe("tickers.ETHUSDT");
616
617        assert_eq!(state.len(), 2);
618        let topics = state.all_topics();
619        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
620        assert!(topics.contains(&"tickers.ETHUSDT".to_string()));
621    }
622
623    #[rstest]
624    fn test_mixed_channel_and_symbol_subscriptions() {
625        let state = SubscriptionState::new('.');
626
627        // Subscribe to channel-level first
628        state.mark_subscribe("tickers");
629        state.confirm_subscribe("tickers");
630        assert_eq!(state.len(), 1);
631        assert_eq!(state.all_topics(), vec!["tickers"]);
632
633        // Add symbol-level subscription to same channel
634        state.mark_subscribe("tickers.BTCUSDT");
635        state.confirm_subscribe("tickers.BTCUSDT");
636        assert_eq!(state.len(), 2);
637
638        // Both should be present
639        let topics = state.all_topics();
640        assert_eq!(topics.len(), 2);
641        assert!(topics.contains(&"tickers".to_string()));
642        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
643
644        // Add another symbol
645        state.mark_subscribe("tickers.ETHUSDT");
646        state.confirm_subscribe("tickers.ETHUSDT");
647        assert_eq!(state.len(), 3);
648
649        let topics = state.all_topics();
650        assert_eq!(topics.len(), 3);
651        assert!(topics.contains(&"tickers".to_string()));
652        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
653        assert!(topics.contains(&"tickers.ETHUSDT".to_string()));
654
655        // Unsubscribe from channel-level only
656        state.mark_unsubscribe("tickers");
657        state.confirm_unsubscribe("tickers");
658        assert_eq!(state.len(), 2);
659
660        let topics = state.all_topics();
661        assert_eq!(topics.len(), 2);
662        assert!(!topics.contains(&"tickers".to_string()));
663        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
664        assert!(topics.contains(&"tickers.ETHUSDT".to_string()));
665    }
666
667    #[rstest]
668    fn test_symbol_subscription_before_channel() {
669        let state = SubscriptionState::new('.');
670
671        // Subscribe to symbol first
672        state.mark_subscribe("tickers.BTCUSDT");
673        state.confirm_subscribe("tickers.BTCUSDT");
674        assert_eq!(state.len(), 1);
675
676        // Then add channel-level
677        state.mark_subscribe("tickers");
678        state.confirm_subscribe("tickers");
679        assert_eq!(state.len(), 2);
680
681        // Both should be present after reconnect
682        let topics = state.all_topics();
683        assert_eq!(topics.len(), 2);
684        assert!(topics.contains(&"tickers".to_string()));
685        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
686    }
687
688    #[rstest]
689    fn test_edge_case_empty_channel_name() {
690        let state = SubscriptionState::new('.');
691
692        // Edge case: empty string as topic
693        state.mark_subscribe("");
694        state.confirm_subscribe("");
695
696        assert_eq!(state.len(), 1);
697        assert_eq!(state.all_topics(), vec![""]);
698    }
699
700    #[rstest]
701    fn test_special_characters_in_topics() {
702        let state = SubscriptionState::new('.');
703
704        // Topics with special characters
705        let special_topics = vec![
706            "channel.symbol-with-dash",
707            "channel.SYMBOL_WITH_UNDERSCORE",
708            "channel.symbol123",
709            "channel.symbol@special",
710        ];
711
712        for topic in &special_topics {
713            state.mark_subscribe(topic);
714            state.confirm_subscribe(topic);
715        }
716
717        assert_eq!(state.len(), special_topics.len());
718
719        let all_topics = state.all_topics();
720
721        for topic in &special_topics {
722            assert!(
723                all_topics.contains(&(*topic).to_string()),
724                "Missing topic: {topic}"
725            );
726        }
727    }
728
729    #[rstest]
730    fn test_edge_case_malformed_topics() {
731        let state = SubscriptionState::new('.');
732
733        // Topics with multiple delimiters (splits on first delimiter)
734        state.mark_subscribe("channel.symbol.extra");
735        state.confirm_subscribe("channel.symbol.extra");
736        let topics = state.all_topics();
737        assert!(topics.contains(&"channel.symbol.extra".to_string()));
738
739        // Topic with leading delimiter (empty channel, symbol is "channel")
740        state.mark_subscribe(".channel");
741        state.confirm_subscribe(".channel");
742        assert_eq!(state.len(), 2);
743
744        // Topic with trailing delimiter - treated as channel-level (empty symbol = marker)
745        // "channel." splits to ("channel", Some("")), and empty string is the channel marker
746        state.mark_subscribe("channel.");
747        state.confirm_subscribe("channel.");
748        assert_eq!(state.len(), 3);
749
750        // Topic without delimiter - explicitly channel-level
751        state.mark_subscribe("tickers");
752        state.confirm_subscribe("tickers");
753        assert_eq!(state.len(), 4);
754
755        // Verify all are retrievable (note: "channel." becomes "channel")
756        let all = state.all_topics();
757        assert_eq!(all.len(), 4);
758        assert!(all.contains(&"channel.symbol.extra".to_string()));
759        assert!(all.contains(&".channel".to_string()));
760        assert!(all.contains(&"channel".to_string())); // "channel." treated as channel-level
761        assert!(all.contains(&"tickers".to_string()));
762    }
763
764    #[rstest]
765    fn test_new_state_is_empty() {
766        let state = SubscriptionState::new('.');
767        assert!(state.is_empty());
768        assert_eq!(state.len(), 0);
769    }
770
771    #[rstest]
772    fn test_subscription_map_accessors_return_isolated_snapshots() {
773        let state = SubscriptionState::new('.');
774        state.mark_subscribe("tickers.BTCUSDT");
775        state.confirm_subscribe("tickers.BTCUSDT");
776
777        let confirmed = state.confirmed();
778        let pending_subscribe = state.pending_subscribe();
779        let pending_unsubscribe = state.pending_unsubscribe();
780
781        state.mark_unsubscribe("tickers.BTCUSDT");
782
783        assert_eq!(
784            confirmed.get(&Ustr::from("tickers")),
785            Some(&AHashSet::from_iter([Ustr::from("BTCUSDT")]))
786        );
787        assert!(pending_subscribe.is_empty());
788        assert!(pending_unsubscribe.is_empty());
789        assert!(state.confirmed().is_empty());
790        assert_eq!(
791            state.pending_unsubscribe().get(&Ustr::from("tickers")),
792            Some(&AHashSet::from_iter([Ustr::from("BTCUSDT")]))
793        );
794    }
795
796    #[rstest]
797    fn test_is_subscribed_empty_state() {
798        let state = SubscriptionState::new('.');
799        let channel = Ustr::from("tickers");
800        let symbol = Ustr::from("BTCUSDT");
801
802        assert!(!state.is_subscribed(&channel, &symbol));
803    }
804
805    #[rstest]
806    fn test_is_subscribed_pending() {
807        let state = SubscriptionState::new('.');
808        let channel = Ustr::from("tickers");
809        let symbol = Ustr::from("BTCUSDT");
810
811        state.mark_subscribe("tickers.BTCUSDT");
812
813        assert!(state.is_subscribed(&channel, &symbol));
814    }
815
816    #[rstest]
817    fn test_is_subscribed_confirmed() {
818        let state = SubscriptionState::new('.');
819        let channel = Ustr::from("tickers");
820        let symbol = Ustr::from("BTCUSDT");
821
822        state.mark_subscribe("tickers.BTCUSDT");
823        state.confirm_subscribe("tickers.BTCUSDT");
824
825        assert!(state.is_subscribed(&channel, &symbol));
826    }
827
828    #[rstest]
829    fn test_is_subscribed_after_unsubscribe() {
830        let state = SubscriptionState::new('.');
831        let channel = Ustr::from("tickers");
832        let symbol = Ustr::from("BTCUSDT");
833
834        state.mark_subscribe("tickers.BTCUSDT");
835        state.confirm_subscribe("tickers.BTCUSDT");
836        state.mark_unsubscribe("tickers.BTCUSDT");
837
838        // Pending unsubscribe should not count as subscribed
839        assert!(!state.is_subscribed(&channel, &symbol));
840    }
841
842    #[rstest]
843    fn test_is_subscribed_after_confirm_unsubscribe() {
844        let state = SubscriptionState::new('.');
845        let channel = Ustr::from("tickers");
846        let symbol = Ustr::from("BTCUSDT");
847
848        state.mark_subscribe("tickers.BTCUSDT");
849        state.confirm_subscribe("tickers.BTCUSDT");
850        state.mark_unsubscribe("tickers.BTCUSDT");
851        state.confirm_unsubscribe("tickers.BTCUSDT");
852
853        assert!(!state.is_subscribed(&channel, &symbol));
854    }
855
856    #[rstest]
857    fn test_all_topics_includes_confirmed_and_pending_subscribe() {
858        let state = SubscriptionState::new('.');
859        state.mark_subscribe("tickers.BTCUSDT");
860        state.confirm_subscribe("tickers.BTCUSDT");
861        state.mark_subscribe("tickers.ETHUSDT");
862
863        let topics = state.all_topics();
864        assert_eq!(topics.len(), 2);
865        assert!(topics.contains(&"tickers.BTCUSDT".to_string()));
866        assert!(topics.contains(&"tickers.ETHUSDT".to_string()));
867    }
868
869    #[rstest]
870    fn test_all_topics_excludes_pending_unsubscribe() {
871        let state = SubscriptionState::new('.');
872        state.mark_subscribe("tickers.BTCUSDT");
873        state.confirm_subscribe("tickers.BTCUSDT");
874        state.mark_unsubscribe("tickers.BTCUSDT");
875
876        let topics = state.all_topics();
877        assert!(topics.is_empty());
878    }
879
880    #[rstest]
881    fn test_all_topics_is_sorted_within_each_group() {
882        let state = SubscriptionState::new('.');
883
884        // Insert scrambled across channels and symbols so hash order cannot pass.
885        for topic in [
886            "trades.SOLUSDT",
887            "tickers.ETHUSDT",
888            "trades.BTCUSDT",
889            "tickers.BTCUSDT",
890        ] {
891            state.mark_subscribe(topic);
892            state.confirm_subscribe(topic);
893        }
894
895        state.mark_subscribe("orders.XRPUSDT");
896        state.mark_subscribe("orders.ADAUSDT");
897
898        // Confirmed topics sort among themselves, then pending ones do the same.
899        assert_eq!(
900            state.all_topics(),
901            vec![
902                "tickers.BTCUSDT",
903                "tickers.ETHUSDT",
904                "trades.BTCUSDT",
905                "trades.SOLUSDT",
906                "orders.ADAUSDT",
907                "orders.XRPUSDT",
908            ]
909        );
910    }
911
912    #[rstest]
913    fn test_pending_subscribe_excludes_pending_unsubscribe() {
914        let state = SubscriptionState::new('.');
915
916        // Subscribe and confirm
917        state.mark_subscribe("tickers.BTCUSDT");
918        state.confirm_subscribe("tickers.BTCUSDT");
919
920        // Mark for unsubscribe
921        state.mark_unsubscribe("tickers.BTCUSDT");
922
923        // Should be in pending_unsubscribe but NOT in all_topics
924        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
925        assert!(state.all_topics().is_empty());
926        assert_eq!(state.len(), 0);
927    }
928
929    #[rstest]
930    fn test_mark_subscribe() {
931        let state = SubscriptionState::new('.');
932        state.mark_subscribe("tickers.BTCUSDT");
933
934        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
935        assert_eq!(state.len(), 0); // Not confirmed yet
936    }
937
938    #[rstest]
939    fn test_try_mark_subscribe_returns_true_once_across_concurrent_calls() {
940        const CALLERS: usize = 32;
941
942        let state = Arc::new(SubscriptionState::new('.'));
943        let start = Arc::new(Barrier::new(CALLERS));
944        let send_count = Arc::new(AtomicUsize::new(0));
945
946        std::thread::scope(|scope| {
947            for _ in 0..CALLERS {
948                let state = Arc::clone(&state);
949                let start = Arc::clone(&start);
950                let send_count = Arc::clone(&send_count);
951
952                scope.spawn(move || {
953                    start.wait();
954
955                    if state.try_mark_subscribe("tickers.BTCUSDT") {
956                        send_count.fetch_add(1, Ordering::SeqCst);
957                    }
958                });
959            }
960        });
961
962        assert_eq!(send_count.load(Ordering::SeqCst), 1);
963        assert_eq!(state.pending_subscribe_topics(), ["tickers.BTCUSDT"]);
964        assert!(state.pending_unsubscribe_topics().is_empty());
965    }
966
967    #[rstest]
968    fn test_try_mark_subscribe_respects_lifecycle_state() {
969        let state = SubscriptionState::new('.');
970        let topic = "tickers.BTCUSDT";
971
972        assert!(state.try_mark_subscribe(topic));
973        assert!(!state.try_mark_subscribe(topic));
974
975        state.confirm_subscribe(topic);
976        assert!(!state.try_mark_subscribe(topic));
977
978        state.mark_unsubscribe(topic);
979        assert!(state.try_mark_subscribe(topic));
980        assert_eq!(state.pending_subscribe_topics(), [topic]);
981        assert!(state.pending_unsubscribe_topics().is_empty());
982    }
983
984    #[rstest]
985    fn test_confirm_subscribe() {
986        let state = SubscriptionState::new('.');
987        state.mark_subscribe("tickers.BTCUSDT");
988        state.confirm_subscribe("tickers.BTCUSDT");
989
990        assert!(state.pending_subscribe_topics().is_empty());
991        assert_eq!(state.len(), 1);
992    }
993
994    #[rstest]
995    fn test_mark_unsubscribe() {
996        let state = SubscriptionState::new('.');
997        state.mark_subscribe("tickers.BTCUSDT");
998        state.confirm_subscribe("tickers.BTCUSDT");
999        state.mark_unsubscribe("tickers.BTCUSDT");
1000
1001        assert_eq!(state.len(), 0); // Removed from confirmed
1002        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1003    }
1004
1005    #[rstest]
1006    fn test_confirm_unsubscribe() {
1007        let state = SubscriptionState::new('.');
1008        state.mark_subscribe("tickers.BTCUSDT");
1009        state.confirm_subscribe("tickers.BTCUSDT");
1010        state.mark_unsubscribe("tickers.BTCUSDT");
1011        state.confirm_unsubscribe("tickers.BTCUSDT");
1012
1013        assert!(state.is_empty());
1014    }
1015
1016    #[rstest]
1017    fn test_unsubscribe_before_subscribe_confirmed() {
1018        let state = SubscriptionState::new('.');
1019
1020        // User subscribes
1021        state.mark_subscribe("tickers.BTCUSDT");
1022        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1023
1024        // User immediately changes mind before server confirms
1025        state.mark_unsubscribe("tickers.BTCUSDT");
1026
1027        // Should be removed from pending_subscribe and added to pending_unsubscribe
1028        assert!(state.pending_subscribe_topics().is_empty());
1029        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1030
1031        // Confirm the unsubscribe
1032        state.confirm_unsubscribe("tickers.BTCUSDT");
1033
1034        // Should be completely gone
1035        assert!(state.is_empty());
1036        assert!(state.all_topics().is_empty());
1037        assert_eq!(state.len(), 0);
1038    }
1039
1040    #[rstest]
1041    fn test_late_subscribe_confirmation_after_unsubscribe() {
1042        let state = SubscriptionState::new('.');
1043
1044        // User subscribes
1045        state.mark_subscribe("tickers.BTCUSDT");
1046
1047        // User immediately unsubscribes
1048        state.mark_unsubscribe("tickers.BTCUSDT");
1049
1050        // Late subscribe confirmation arrives from server
1051        state.confirm_subscribe("tickers.BTCUSDT");
1052
1053        // Should NOT be added to confirmed (unsubscribe takes precedence)
1054        assert_eq!(state.len(), 0);
1055        assert!(state.pending_subscribe_topics().is_empty());
1056
1057        // Confirm the unsubscribe
1058        state.confirm_unsubscribe("tickers.BTCUSDT");
1059
1060        // Should still be empty
1061        assert!(state.is_empty());
1062        assert!(state.all_topics().is_empty());
1063    }
1064
1065    #[rstest]
1066    fn test_late_subscribe_ack_after_unsubscribe_ack_does_not_restore_topic() {
1067        let state = SubscriptionState::new('.');
1068        state.mark_subscribe("tickers.BTCUSDT");
1069        state.mark_unsubscribe("tickers.BTCUSDT");
1070
1071        state.confirm_unsubscribe("tickers.BTCUSDT");
1072        state.confirm_subscribe("tickers.BTCUSDT");
1073
1074        assert!(state.is_empty());
1075        assert!(state.all_topics().is_empty());
1076        assert!(state.pending_subscribe_topics().is_empty());
1077        assert!(state.pending_unsubscribe_topics().is_empty());
1078    }
1079
1080    #[rstest]
1081    fn test_resubscribe_before_unsubscribe_ack() {
1082        // Regression test for race condition:
1083        // User unsubscribes, then immediately resubscribes before the unsubscribe ACK arrives.
1084        // The unsubscribe ACK should NOT clear the pending_subscribe entry.
1085        let state = SubscriptionState::new('.');
1086
1087        state.mark_subscribe("tickers.BTCUSDT");
1088        state.confirm_subscribe("tickers.BTCUSDT");
1089        assert_eq!(state.len(), 1);
1090
1091        state.mark_unsubscribe("tickers.BTCUSDT");
1092        assert_eq!(state.len(), 0);
1093        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1094
1095        // User immediately resubscribes (before unsubscribe ACK)
1096        state.mark_subscribe("tickers.BTCUSDT");
1097        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1098
1099        // Stale unsubscribe ACK arrives - should be ignored (pending_unsubscribe already cleared)
1100        state.confirm_unsubscribe("tickers.BTCUSDT");
1101        assert!(state.pending_unsubscribe_topics().is_empty());
1102        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]); // Must still be pending
1103
1104        // Subscribe ACK confirms successfully
1105        state.confirm_subscribe("tickers.BTCUSDT");
1106        assert_eq!(state.len(), 1);
1107        assert!(state.pending_subscribe_topics().is_empty());
1108
1109        // Topic available for reconnect
1110        let all = state.all_topics();
1111        assert_eq!(all.len(), 1);
1112        assert!(all.contains(&"tickers.BTCUSDT".to_string()));
1113    }
1114
1115    #[rstest]
1116    fn test_stale_unsubscribe_ack_after_resubscribe_confirmed() {
1117        // Regression test for P1 bug: Stale unsubscribe ACK removing confirmed topic.
1118        // Scenario: User unsubscribes, immediately resubscribes, subscribe ACK arrives
1119        // FIRST (out of order), then stale unsubscribe ACK arrives.
1120        // The stale ACK must NOT remove the topic from confirmed state.
1121        let state = SubscriptionState::new('.');
1122
1123        // Initial subscription
1124        state.mark_subscribe("tickers.BTCUSDT");
1125        state.confirm_subscribe("tickers.BTCUSDT");
1126        assert_eq!(state.len(), 1);
1127
1128        // User unsubscribes
1129        state.mark_unsubscribe("tickers.BTCUSDT");
1130        assert_eq!(state.len(), 0);
1131        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1132
1133        // User immediately resubscribes (before unsubscribe ACK)
1134        state.mark_subscribe("tickers.BTCUSDT");
1135        assert!(state.pending_unsubscribe_topics().is_empty()); // Cleared by mark_subscribe
1136        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1137
1138        // Subscribe ACK arrives FIRST (out of order!)
1139        state.confirm_subscribe("tickers.BTCUSDT");
1140        assert_eq!(state.len(), 1); // Back in confirmed
1141        assert!(state.pending_subscribe_topics().is_empty());
1142
1143        // NOW the stale unsubscribe ACK arrives
1144        // This must be ignored because topic is no longer in pending_unsubscribe
1145        state.confirm_unsubscribe("tickers.BTCUSDT");
1146
1147        // Topic should STILL be confirmed (not removed by stale ACK)
1148        assert_eq!(state.len(), 1); // Must remain confirmed
1149        assert!(state.pending_unsubscribe_topics().is_empty());
1150        assert!(state.pending_subscribe_topics().is_empty());
1151
1152        // Topic should be in all_topics (for reconnect)
1153        let all = state.all_topics();
1154        assert_eq!(all.len(), 1);
1155        assert!(all.contains(&"tickers.BTCUSDT".to_string()));
1156    }
1157
1158    #[rstest]
1159    fn test_unsubscribe_clears_all_states() {
1160        let state = SubscriptionState::new('.');
1161
1162        // Subscribe and confirm
1163        state.mark_subscribe("tickers.BTCUSDT");
1164        state.confirm_subscribe("tickers.BTCUSDT");
1165        assert_eq!(state.len(), 1);
1166
1167        // Unsubscribe
1168        state.mark_unsubscribe("tickers.BTCUSDT");
1169
1170        // Should be removed from confirmed
1171        assert_eq!(state.len(), 0);
1172        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1173
1174        // Late subscribe confirmation somehow arrives (race condition)
1175        state.confirm_subscribe("tickers.BTCUSDT");
1176
1177        // confirm_unsubscribe should clean everything
1178        state.confirm_unsubscribe("tickers.BTCUSDT");
1179
1180        // Completely empty
1181        assert!(state.is_empty());
1182        assert_eq!(state.len(), 0);
1183        assert!(state.pending_subscribe_topics().is_empty());
1184        assert!(state.pending_unsubscribe_topics().is_empty());
1185        assert!(state.all_topics().is_empty());
1186    }
1187
1188    #[rstest]
1189    fn test_concurrent_subscribe_confirmation_and_unsubscribe_preserve_intent() {
1190        const ITERATIONS: usize = 10_000;
1191
1192        let state = Arc::new(SubscriptionState::new('.'));
1193        let start = Arc::new(Barrier::new(3));
1194        let finish = Arc::new(Barrier::new(3));
1195        let failed_iteration = Arc::new(AtomicUsize::new(usize::MAX));
1196
1197        std::thread::scope(|scope| {
1198            let confirming_state = Arc::clone(&state);
1199            let confirming_start = Arc::clone(&start);
1200            let confirming_finish = Arc::clone(&finish);
1201
1202            scope.spawn(move || {
1203                for _ in 0..ITERATIONS {
1204                    confirming_start.wait();
1205                    confirming_state.confirm_subscribe("tickers.BTCUSDT");
1206                    confirming_finish.wait();
1207                }
1208            });
1209
1210            let unsubscribing_state = Arc::clone(&state);
1211            let unsubscribing_start = Arc::clone(&start);
1212            let unsubscribing_finish = Arc::clone(&finish);
1213
1214            scope.spawn(move || {
1215                for _ in 0..ITERATIONS {
1216                    unsubscribing_start.wait();
1217                    unsubscribing_state.mark_unsubscribe("tickers.BTCUSDT");
1218                    unsubscribing_finish.wait();
1219                }
1220            });
1221
1222            for iteration in 0..ITERATIONS {
1223                state.clear();
1224                state.mark_subscribe("tickers.BTCUSDT");
1225                start.wait();
1226                finish.wait();
1227
1228                if !state.all_topics().is_empty()
1229                    || state.pending_unsubscribe_topics() != ["tickers.BTCUSDT"]
1230                {
1231                    _ = failed_iteration.compare_exchange(
1232                        usize::MAX,
1233                        iteration,
1234                        Ordering::SeqCst,
1235                        Ordering::SeqCst,
1236                    );
1237                }
1238            }
1239        });
1240
1241        assert_eq!(
1242            failed_iteration.load(Ordering::SeqCst),
1243            usize::MAX,
1244            "late subscribe confirmation restored the unsubscribe intent"
1245        );
1246    }
1247
1248    #[rstest]
1249    fn test_state_machine_invalid_transitions() {
1250        let state = SubscriptionState::new('.');
1251
1252        // Confirm subscribe without matching intent - should be ignored
1253        state.confirm_subscribe("tickers.BTCUSDT");
1254        assert_eq!(state.len(), 0);
1255
1256        // Confirm unsubscribe without marking first - should not crash
1257        state.confirm_unsubscribe("tickers.ETHUSDT");
1258        assert_eq!(state.len(), 0); // Nothing changes
1259
1260        // Double confirm subscribe
1261        state.mark_subscribe("orderbook");
1262        state.confirm_subscribe("orderbook");
1263        state.confirm_subscribe("orderbook"); // Second confirm is idempotent
1264        assert_eq!(state.len(), 1);
1265
1266        // Unsubscribe something that was never subscribed
1267        state.mark_unsubscribe("nonexistent");
1268        state.confirm_unsubscribe("nonexistent");
1269        assert_eq!(state.len(), 1); // Still 1
1270    }
1271
1272    #[rstest]
1273    fn test_mark_failure() {
1274        let state = SubscriptionState::new('.');
1275        state.mark_subscribe("tickers.BTCUSDT");
1276        state.confirm_subscribe("tickers.BTCUSDT");
1277        state.mark_failure("tickers.BTCUSDT");
1278
1279        assert_eq!(state.len(), 0);
1280        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1281    }
1282
1283    #[rstest]
1284    fn test_mark_failure_moves_to_pending() {
1285        let state = SubscriptionState::new('.');
1286
1287        // Subscribe and confirm
1288        state.mark_subscribe("tickers.BTCUSDT");
1289        state.confirm_subscribe("tickers.BTCUSDT");
1290        assert_eq!(state.len(), 1);
1291        assert!(state.pending_subscribe_topics().is_empty());
1292
1293        // Mark as failed
1294        state.mark_failure("tickers.BTCUSDT");
1295
1296        // Should be removed from confirmed and back in pending
1297        assert_eq!(state.len(), 0);
1298        assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1299
1300        // all_topics should still include it for reconnection
1301        assert_eq!(state.all_topics(), vec!["tickers.BTCUSDT"]);
1302    }
1303
1304    #[rstest]
1305    fn test_mark_failure_respects_pending_unsubscribe() {
1306        let state = SubscriptionState::new('.');
1307
1308        // Subscribe and confirm
1309        state.mark_subscribe("tickers.BTCUSDT");
1310        state.confirm_subscribe("tickers.BTCUSDT");
1311        assert_eq!(state.len(), 1);
1312
1313        // User unsubscribes
1314        state.mark_unsubscribe("tickers.BTCUSDT");
1315        assert_eq!(state.len(), 0);
1316        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1317
1318        // Meanwhile, a network error triggers mark_failure
1319        state.mark_failure("tickers.BTCUSDT");
1320
1321        // Should NOT be added to pending_subscribe (user wanted to unsubscribe)
1322        assert!(state.pending_subscribe_topics().is_empty());
1323        assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1324
1325        // all_topics should NOT include it
1326        assert!(state.all_topics().is_empty());
1327
1328        // Confirm unsubscribe
1329        state.confirm_unsubscribe("tickers.BTCUSDT");
1330        assert!(state.is_empty());
1331    }
1332
1333    #[rstest]
1334    fn test_reconnection_scenario() {
1335        let state = SubscriptionState::new('.');
1336
1337        // Initial subscriptions
1338        state.add_reference("tickers.BTCUSDT");
1339        state.mark_subscribe("tickers.BTCUSDT");
1340        state.confirm_subscribe("tickers.BTCUSDT");
1341
1342        state.add_reference("tickers.ETHUSDT");
1343        state.mark_subscribe("tickers.ETHUSDT");
1344        state.confirm_subscribe("tickers.ETHUSDT");
1345
1346        state.add_reference("orderbook");
1347        state.mark_subscribe("orderbook");
1348        state.confirm_subscribe("orderbook");
1349
1350        assert_eq!(state.len(), 3);
1351
1352        // Simulate disconnect - topics should be available for resubscription
1353        let topics_to_resubscribe = state.all_topics();
1354        assert_eq!(topics_to_resubscribe.len(), 3);
1355        assert!(topics_to_resubscribe.contains(&"tickers.BTCUSDT".to_string()));
1356        assert!(topics_to_resubscribe.contains(&"tickers.ETHUSDT".to_string()));
1357        assert!(topics_to_resubscribe.contains(&"orderbook".to_string()));
1358
1359        // On reconnect, mark all as pending again
1360        for topic in &topics_to_resubscribe {
1361            state.mark_subscribe(topic);
1362        }
1363
1364        // Simulate server confirmations
1365        for topic in &topics_to_resubscribe {
1366            state.confirm_subscribe(topic);
1367        }
1368
1369        // Should still have all 3 subscriptions
1370        assert_eq!(state.len(), 3);
1371        assert_eq!(state.all_topics().len(), 3);
1372    }
1373
1374    #[rstest]
1375    fn test_reconnection_with_partial_state() {
1376        let state = SubscriptionState::new('.');
1377
1378        // Setup: Some confirmed, some pending subscribe, some pending unsubscribe
1379        // Confirmed
1380        state.add_reference("confirmed.BTCUSDT");
1381        state.mark_subscribe("confirmed.BTCUSDT");
1382        state.confirm_subscribe("confirmed.BTCUSDT");
1383
1384        // Pending subscribe (not yet confirmed)
1385        state.add_reference("pending.ETHUSDT");
1386        state.mark_subscribe("pending.ETHUSDT");
1387
1388        // Pending unsubscribe (user cancelled)
1389        state.mark_subscribe("cancelled.XRPUSDT");
1390        state.confirm_subscribe("cancelled.XRPUSDT");
1391        state.mark_unsubscribe("cancelled.XRPUSDT");
1392
1393        // Verify state before reconnect
1394        assert_eq!(state.len(), 1); // Only confirmed.BTCUSDT
1395        let all = state.all_topics();
1396        assert_eq!(all.len(), 2); // confirmed + pending_subscribe (not pending_unsubscribe)
1397        assert!(all.contains(&"confirmed.BTCUSDT".to_string()));
1398        assert!(all.contains(&"pending.ETHUSDT".to_string()));
1399        assert!(!all.contains(&"cancelled.XRPUSDT".to_string())); // Should NOT be included
1400
1401        // Simulate disconnect and reconnect
1402        let topics_to_resubscribe = state.reset_after_reconnect();
1403        assert_eq!(
1404            topics_to_resubscribe,
1405            [
1406                "confirmed.BTCUSDT".to_string(),
1407                "pending.ETHUSDT".to_string()
1408            ]
1409        );
1410        assert_eq!(state.reset_after_reconnect(), topics_to_resubscribe);
1411        assert_eq!(state.len(), 0);
1412        assert_eq!(
1413            state.pending_subscribe_topics(),
1414            [
1415                "confirmed.BTCUSDT".to_string(),
1416                "pending.ETHUSDT".to_string()
1417            ]
1418        );
1419        assert!(state.pending_unsubscribe_topics().is_empty());
1420        assert_eq!(state.get_reference_count("confirmed.BTCUSDT"), 1);
1421        assert_eq!(state.get_reference_count("pending.ETHUSDT"), 1);
1422
1423        // Server confirms both
1424        for topic in &topics_to_resubscribe {
1425            state.confirm_subscribe(topic);
1426        }
1427
1428        // Verify final state
1429        assert_eq!(state.len(), 2); // Both confirmed
1430        let final_topics = state.all_topics();
1431        assert_eq!(final_topics.len(), 2);
1432        assert!(final_topics.contains(&"confirmed.BTCUSDT".to_string()));
1433        assert!(final_topics.contains(&"pending.ETHUSDT".to_string()));
1434        assert!(!final_topics.contains(&"cancelled.XRPUSDT".to_string()));
1435    }
1436
1437    #[rstest]
1438    fn test_reference_counting_single_topic() {
1439        let state = SubscriptionState::new('.');
1440
1441        assert!(state.add_reference("tickers.BTCUSDT"));
1442        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 1);
1443
1444        assert!(!state.add_reference("tickers.BTCUSDT"));
1445        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 2);
1446
1447        assert!(!state.remove_reference("tickers.BTCUSDT"));
1448        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 1);
1449
1450        assert!(state.remove_reference("tickers.BTCUSDT"));
1451        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 0);
1452    }
1453
1454    #[rstest]
1455    fn test_reference_counting_multiple_topics() {
1456        let state = SubscriptionState::new('.');
1457
1458        assert!(state.add_reference("tickers.BTCUSDT"));
1459        assert!(state.add_reference("tickers.ETHUSDT"));
1460
1461        assert!(!state.add_reference("tickers.BTCUSDT"));
1462        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 2);
1463        assert_eq!(state.get_reference_count("tickers.ETHUSDT"), 1);
1464
1465        assert!(!state.remove_reference("tickers.BTCUSDT"));
1466        assert!(state.remove_reference("tickers.ETHUSDT"));
1467    }
1468
1469    #[rstest]
1470    fn test_remove_reference_nonexistent_topic() {
1471        let state = SubscriptionState::new('.');
1472
1473        // Removing reference to topic that was never added
1474        let should_unsubscribe = state.remove_reference("nonexistent");
1475
1476        // Should return false and not crash
1477        assert!(!should_unsubscribe);
1478        assert_eq!(state.get_reference_count("nonexistent"), 0);
1479    }
1480
1481    #[rstest]
1482    fn test_reference_count_underflow_safety() {
1483        let state = SubscriptionState::new('.');
1484
1485        // Remove without ever adding
1486        assert!(!state.remove_reference("never.added"));
1487        assert_eq!(state.get_reference_count("never.added"), 0);
1488
1489        // Add one, remove multiple times
1490        state.add_reference("once.added");
1491        assert_eq!(state.get_reference_count("once.added"), 1);
1492
1493        assert!(state.remove_reference("once.added")); // Should return true (last ref)
1494        assert_eq!(state.get_reference_count("once.added"), 0);
1495
1496        assert!(!state.remove_reference("once.added")); // Should not crash, returns false
1497        assert!(!state.remove_reference("once.added")); // Multiple times
1498        assert_eq!(state.get_reference_count("once.added"), 0);
1499
1500        // Verify we can add again after underflow attempts
1501        assert!(state.add_reference("once.added"));
1502        assert_eq!(state.get_reference_count("once.added"), 1);
1503    }
1504
1505    #[rstest]
1506    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1507    async fn test_concurrent_reference_counting_same_topic() {
1508        let state = Arc::new(SubscriptionState::new('.'));
1509        let topic = "tickers.BTCUSDT";
1510        let mut handles = vec![];
1511
1512        // Spawn 10 tasks all adding 10 references to the same topic
1513        for _ in 0..10 {
1514            let state_clone = Arc::clone(&state);
1515
1516            let handle = tokio::spawn(async move {
1517                for _ in 0..10 {
1518                    state_clone.add_reference(topic);
1519                }
1520            });
1521            handles.push(handle);
1522        }
1523
1524        for handle in handles {
1525            handle.await.unwrap();
1526        }
1527
1528        // Should have exactly 100 references (10 tasks * 10 refs each)
1529        assert_eq!(state.get_reference_count(topic), 100);
1530
1531        // Now remove 50 references sequentially
1532        for _ in 0..50 {
1533            state.remove_reference(topic);
1534        }
1535
1536        // Should have exactly 50 references remaining
1537        assert_eq!(state.get_reference_count(topic), 50);
1538    }
1539
1540    #[rstest]
1541    fn test_clear() {
1542        let state = SubscriptionState::new('.');
1543        state.mark_subscribe("tickers.BTCUSDT");
1544        state.confirm_subscribe("tickers.BTCUSDT");
1545        state.add_reference("tickers.BTCUSDT");
1546
1547        state.clear();
1548
1549        assert!(state.is_empty());
1550        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 0);
1551    }
1552
1553    #[rstest]
1554    fn test_clear_resets_all_state() {
1555        let state = SubscriptionState::new('.');
1556
1557        // Add multiple subscriptions and references
1558        for i in 0..10 {
1559            let topic = format!("channel{i}.SYMBOL");
1560            state.add_reference(&topic);
1561            state.add_reference(&topic); // Add twice
1562            state.mark_subscribe(&topic);
1563            state.confirm_subscribe(&topic);
1564        }
1565
1566        assert_eq!(state.len(), 10);
1567        assert!(!state.is_empty());
1568
1569        // Clear everything
1570        state.clear();
1571
1572        // Verify complete reset
1573        assert_eq!(state.len(), 0);
1574        assert!(state.is_empty());
1575        assert!(state.all_topics().is_empty());
1576        assert!(state.pending_subscribe_topics().is_empty());
1577        assert!(state.pending_unsubscribe_topics().is_empty());
1578
1579        // Verify reference counts are cleared
1580        for i in 0..10 {
1581            let topic = format!("channel{i}.SYMBOL");
1582            assert_eq!(state.get_reference_count(&topic), 0);
1583        }
1584    }
1585
1586    #[rstest]
1587    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1588    async fn test_concurrent_subscribe_same_topic() {
1589        let state = Arc::new(SubscriptionState::new('.'));
1590        let mut handles = vec![];
1591
1592        // Spawn 10 tasks all subscribing to the same topic
1593        for _ in 0..10 {
1594            let state_clone = Arc::clone(&state);
1595            let handle = tokio::spawn(async move {
1596                state_clone.add_reference("tickers.BTCUSDT");
1597                state_clone.mark_subscribe("tickers.BTCUSDT");
1598                state_clone.confirm_subscribe("tickers.BTCUSDT");
1599            });
1600            handles.push(handle);
1601        }
1602
1603        for handle in handles {
1604            handle.await.unwrap();
1605        }
1606
1607        // Reference count should be exactly 10
1608        assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 10);
1609        assert_eq!(state.len(), 1);
1610    }
1611
1612    #[rstest]
1613    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1614    async fn test_concurrent_subscribe_unsubscribe() {
1615        let state = Arc::new(SubscriptionState::new('.'));
1616        let mut handles = vec![];
1617
1618        // Spawn 20 tasks, each adding 2 references to their own unique topic
1619        // This ensures deterministic behavior - we know exactly what the final state should be
1620        for i in 0..20 {
1621            let state_clone = Arc::clone(&state);
1622
1623            let handle = tokio::spawn(async move {
1624                let topic = format!("tickers.SYMBOL{i}");
1625                // Add 2 references
1626                state_clone.add_reference(&topic);
1627                state_clone.add_reference(&topic);
1628                state_clone.mark_subscribe(&topic);
1629                state_clone.confirm_subscribe(&topic);
1630
1631                // Remove 1 reference (should still have 1 remaining)
1632                state_clone.remove_reference(&topic);
1633            });
1634            handles.push(handle);
1635        }
1636
1637        for handle in handles {
1638            handle.await.unwrap();
1639        }
1640
1641        // Each of the 20 topics should still have 1 reference
1642        for i in 0..20 {
1643            let topic = format!("tickers.SYMBOL{i}");
1644            assert_eq!(state.get_reference_count(&topic), 1);
1645        }
1646
1647        // Should have exactly 20 confirmed subscriptions
1648        assert_eq!(state.len(), 20);
1649        assert!(!state.is_empty());
1650    }
1651
1652    #[rstest]
1653    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1654    async fn test_concurrent_stress_mixed_operations() {
1655        let state = Arc::new(SubscriptionState::new('.'));
1656        let mut handles = vec![];
1657
1658        // Spawn 50 tasks doing random interleaved operations
1659        for i in 0..50 {
1660            let state_clone = Arc::clone(&state);
1661
1662            let handle = tokio::spawn(async move {
1663                let topic1 = format!("channel.SYMBOL{i}");
1664                let topic2 = format!("channel.SYMBOL{}", i + 100);
1665
1666                // Add references
1667                state_clone.add_reference(&topic1);
1668                state_clone.add_reference(&topic2);
1669
1670                // Mark and confirm subscriptions
1671                state_clone.mark_subscribe(&topic1);
1672                state_clone.confirm_subscribe(&topic1);
1673                state_clone.mark_subscribe(&topic2);
1674
1675                // Interleave some unsubscribes
1676                if i % 3 == 0 {
1677                    state_clone.mark_unsubscribe(&topic1);
1678                    state_clone.confirm_unsubscribe(&topic1);
1679                }
1680
1681                // More reference operations
1682                state_clone.add_reference(&topic2);
1683                state_clone.remove_reference(&topic2);
1684
1685                // Confirm topic2
1686                state_clone.confirm_subscribe(&topic2);
1687            });
1688            handles.push(handle);
1689        }
1690
1691        for handle in handles {
1692            handle.await.unwrap();
1693        }
1694
1695        let actual = state.all_topics().into_iter().collect::<AHashSet<_>>();
1696        let expected = (0..50)
1697            .flat_map(|i| {
1698                let topic2 = format!("channel.SYMBOL{}", i + 100);
1699                (i % 3 != 0)
1700                    .then(|| format!("channel.SYMBOL{i}"))
1701                    .into_iter()
1702                    .chain(std::iter::once(topic2))
1703            })
1704            .collect::<AHashSet<_>>();
1705
1706        assert_eq!(actual, expected);
1707        assert_eq!(state.len(), 83);
1708        assert!(state.pending_subscribe_topics().is_empty());
1709        assert!(state.pending_unsubscribe_topics().is_empty());
1710        assert_eq!(state.reference_counts.len(), 100);
1711    }
1712
1713    #[rstest]
1714    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1715    async fn test_stress_rapid_resubscribe_pattern() {
1716        // Stress test the race condition we fixed: rapid unsubscribe -> resubscribe
1717        let state = Arc::new(SubscriptionState::new('.'));
1718        let mut handles = vec![];
1719
1720        for i in 0..100 {
1721            let state_clone = Arc::clone(&state);
1722
1723            let handle = tokio::spawn(async move {
1724                let topic = format!("rapid.SYMBOL{}", i % 10); // 10 unique topics, lots of contention
1725
1726                // Initial subscribe
1727                state_clone.mark_subscribe(&topic);
1728                state_clone.confirm_subscribe(&topic);
1729
1730                // Rapid unsubscribe -> resubscribe (race condition scenario)
1731                state_clone.mark_unsubscribe(&topic);
1732                // Immediately resubscribe before unsubscribe ACK
1733                state_clone.mark_subscribe(&topic);
1734                // Now unsubscribe ACK arrives
1735                state_clone.confirm_unsubscribe(&topic);
1736                // Subscribe ACK arrives
1737                state_clone.confirm_subscribe(&topic);
1738            });
1739            handles.push(handle);
1740        }
1741
1742        for handle in handles {
1743            handle.await.unwrap();
1744        }
1745
1746        check_invariants(&state, "After rapid resubscribe stress test");
1747    }
1748
1749    #[rstest]
1750    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1751    async fn test_stress_failure_recovery_loop() {
1752        // Stress test failure -> recovery loops
1753        // Each task gets its own unique topic to avoid race conditions in the test itself
1754        let state = Arc::new(SubscriptionState::new('.'));
1755        let mut handles = vec![];
1756
1757        for i in 0..30 {
1758            let state_clone = Arc::clone(&state);
1759
1760            let handle = tokio::spawn(async move {
1761                let topic = format!("failure.SYMBOL{i}"); // Unique topic per task
1762
1763                // Subscribe and confirm
1764                state_clone.mark_subscribe(&topic);
1765                state_clone.confirm_subscribe(&topic);
1766
1767                // Simulate multiple failures and recoveries
1768                for _ in 0..5 {
1769                    state_clone.mark_failure(&topic);
1770                    state_clone.confirm_subscribe(&topic); // Re-confirm after retry
1771                }
1772            });
1773            handles.push(handle);
1774        }
1775
1776        for handle in handles {
1777            handle.await.unwrap();
1778        }
1779
1780        check_invariants(&state, "After failure recovery loops");
1781
1782        // All should eventually be confirmed (30 unique topics)
1783        assert_eq!(state.len(), 30);
1784    }
1785
1786    #[rstest]
1787    fn test_exhaustive_two_step_transitions() {
1788        let operations = [
1789            "mark_subscribe",
1790            "confirm_subscribe",
1791            "mark_unsubscribe",
1792            "confirm_unsubscribe",
1793            "mark_failure",
1794        ];
1795
1796        for &op1 in &operations {
1797            for &op2 in &operations {
1798                let state = SubscriptionState::new('.');
1799                let topic = "test.TOPIC";
1800
1801                // Apply two operations
1802                apply_op(&state, op1, topic);
1803                apply_op(&state, op2, topic);
1804
1805                // Verify invariants hold
1806                check_invariants(&state, &format!("{op1} -> {op2}"));
1807                check_topic_exclusivity(&state, topic, &format!("{op1} -> {op2}"));
1808            }
1809        }
1810    }
1811
1812    fn apply_op(state: &SubscriptionState, op: &str, topic: &str) {
1813        match op {
1814            "mark_subscribe" => state.mark_subscribe(topic),
1815            "confirm_subscribe" => state.confirm_subscribe(topic),
1816            "mark_unsubscribe" => state.mark_unsubscribe(topic),
1817            "confirm_unsubscribe" => state.confirm_unsubscribe(topic),
1818            "mark_failure" => state.mark_failure(topic),
1819            _ => panic!("Unknown operation: {op}"),
1820        }
1821    }
1822
1823    /// Verifies all invariants of the subscription state.
1824    ///
1825    /// # Invariants
1826    ///
1827    /// 1. **Mutual exclusivity**: A topic cannot exist in multiple states simultaneously
1828    ///    (one of: confirmed, `pending_subscribe`, `pending_unsubscribe`, or none).
1829    /// 2. **`all_topics` consistency**: `all_topics()` must equal `confirmed ∪ pending_subscribe`
1830    /// 3. **len consistency**: `len()` must equal total count of symbols in confirmed map
1831    /// 4. **`is_empty` consistency**: `is_empty()` true iff all maps are empty
1832    /// 5. **Reference count non-negative**: All reference counts >= 0
1833    fn check_invariants(state: &SubscriptionState, label: &str) {
1834        // Collect all topics from each state
1835        let confirmed_topics: AHashSet<String> = state
1836            .topics_from_map(&state.confirmed)
1837            .into_iter()
1838            .collect();
1839        let pending_sub_topics: AHashSet<String> =
1840            state.pending_subscribe_topics().into_iter().collect();
1841        let pending_unsub_topics: AHashSet<String> =
1842            state.pending_unsubscribe_topics().into_iter().collect();
1843
1844        // INVARIANT 1: Mutual exclusivity - no topic in multiple states
1845        let confirmed_and_pending_sub: Vec<_> =
1846            confirmed_topics.intersection(&pending_sub_topics).collect();
1847        assert!(
1848            confirmed_and_pending_sub.is_empty(),
1849            "{label}: Topic in both confirmed and pending_subscribe: {confirmed_and_pending_sub:?}"
1850        );
1851
1852        let confirmed_and_pending_unsub: Vec<_> = confirmed_topics
1853            .intersection(&pending_unsub_topics)
1854            .collect();
1855        assert!(
1856            confirmed_and_pending_unsub.is_empty(),
1857            "{label}: Topic in both confirmed and pending_unsubscribe: {confirmed_and_pending_unsub:?}"
1858        );
1859
1860        let pending_sub_and_unsub: Vec<_> = pending_sub_topics
1861            .intersection(&pending_unsub_topics)
1862            .collect();
1863        assert!(
1864            pending_sub_and_unsub.is_empty(),
1865            "{label}: Topic in both pending_subscribe and pending_unsubscribe: {pending_sub_and_unsub:?}"
1866        );
1867
1868        // INVARIANT 2: all_topics() == confirmed ∪ pending_subscribe
1869        let all_topics: AHashSet<String> = state.all_topics().into_iter().collect();
1870        let expected_all: AHashSet<String> = confirmed_topics
1871            .union(&pending_sub_topics)
1872            .cloned()
1873            .collect();
1874        assert_eq!(
1875            all_topics, expected_all,
1876            "{label}: all_topics() doesn't match confirmed ∪ pending_subscribe"
1877        );
1878
1879        // Ensure pending_unsubscribe is NOT in all_topics
1880        for topic in &pending_unsub_topics {
1881            assert!(
1882                !all_topics.contains(topic),
1883                "{label}: pending_unsubscribe topic {topic} incorrectly in all_topics()"
1884            );
1885        }
1886
1887        // INVARIANT 3: len() == sum of confirmed symbol counts
1888        let expected_len: usize = state
1889            .confirmed
1890            .iter()
1891            .map(|entry| entry.value().len())
1892            .sum();
1893        assert_eq!(
1894            state.len(),
1895            expected_len,
1896            "{label}: len() mismatch. Expected {expected_len}, was {}",
1897            state.len()
1898        );
1899
1900        // INVARIANT 4: is_empty() consistency
1901        let should_be_empty = state.confirmed.is_empty()
1902            && pending_sub_topics.is_empty()
1903            && pending_unsub_topics.is_empty();
1904        assert_eq!(
1905            state.is_empty(),
1906            should_be_empty,
1907            "{label}: is_empty() inconsistent. Maps empty: {should_be_empty}, is_empty(): {}",
1908            state.is_empty()
1909        );
1910
1911        // INVARIANT 5: Reference counts non-negative (NonZeroUsize enforces > 0, absence = 0)
1912        for entry in state.reference_counts.iter() {
1913            let count = entry.value().get();
1914            assert!(
1915                count > 0,
1916                "{label}: Reference count should be NonZeroUsize (> 0), was {count} for {:?}",
1917                entry.key()
1918            );
1919        }
1920    }
1921
1922    /// Checks that a topic exists in exactly one of the three states or none.
1923    fn check_topic_exclusivity(state: &SubscriptionState, topic: &str, label: &str) {
1924        let (channel, symbol) = split_topic(topic, state.delimiter);
1925
1926        let in_confirmed = is_tracked(&state.confirmed, channel, symbol);
1927        let in_pending_sub = is_tracked(&state.pending_subscribe, channel, symbol);
1928        let in_pending_unsub = is_tracked(&state.pending_unsubscribe, channel, symbol);
1929
1930        let count = [in_confirmed, in_pending_sub, in_pending_unsub]
1931            .iter()
1932            .filter(|&&x| x)
1933            .count();
1934
1935        assert!(
1936            count <= 1,
1937            "{label}: Topic {topic} in {count} states (should be 0 or 1). \
1938             confirmed: {in_confirmed}, pending_sub: {in_pending_sub}, pending_unsub: {in_pending_unsub}"
1939        );
1940    }
1941
1942    #[cfg(test)]
1943    mod property_tests {
1944        use ahash::AHashMap;
1945        use proptest::prelude::*;
1946
1947        use super::*;
1948
1949        #[derive(Debug, Clone)]
1950        enum Operation {
1951            MarkSubscribe(String),
1952            ConfirmSubscribe(String),
1953            MarkUnsubscribe(String),
1954            ConfirmUnsubscribe(String),
1955            MarkFailure(String),
1956            AddReference(String),
1957            RemoveReference(String),
1958            Clear,
1959        }
1960
1961        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
1962        enum ModelState {
1963            Confirmed,
1964            PendingSubscribe,
1965            PendingUnsubscribe,
1966        }
1967
1968        // Strategy for generating valid topics
1969        fn topic_strategy() -> impl Strategy<Value = String> {
1970            prop_oneof![
1971                // Symbol-level topics
1972                (any::<u8>(), any::<u8>())
1973                    .prop_map(|(ch, sym)| { format!("channel{}.SYMBOL{}", ch % 5, sym % 10) }),
1974                // Channel-level topics (no symbol)
1975                any::<u8>().prop_map(|ch| format!("channel{}", ch % 5)),
1976            ]
1977        }
1978
1979        // Strategy for generating random operations
1980        fn operation_strategy() -> impl Strategy<Value = Operation> {
1981            topic_strategy().prop_flat_map(|topic| {
1982                prop_oneof![
1983                    Just(Operation::MarkSubscribe(topic.clone())),
1984                    Just(Operation::ConfirmSubscribe(topic.clone())),
1985                    Just(Operation::MarkUnsubscribe(topic.clone())),
1986                    Just(Operation::ConfirmUnsubscribe(topic.clone())),
1987                    Just(Operation::MarkFailure(topic.clone())),
1988                    Just(Operation::AddReference(topic.clone())),
1989                    Just(Operation::RemoveReference(topic)),
1990                    Just(Operation::Clear),
1991                ]
1992            })
1993        }
1994
1995        // Apply an operation to the state
1996        fn apply_operation(state: &SubscriptionState, op: &Operation) {
1997            match op {
1998                Operation::MarkSubscribe(topic) => state.mark_subscribe(topic),
1999                Operation::ConfirmSubscribe(topic) => state.confirm_subscribe(topic),
2000                Operation::MarkUnsubscribe(topic) => state.mark_unsubscribe(topic),
2001                Operation::ConfirmUnsubscribe(topic) => state.confirm_unsubscribe(topic),
2002                Operation::MarkFailure(topic) => state.mark_failure(topic),
2003                Operation::AddReference(topic) => {
2004                    state.add_reference(topic);
2005                }
2006                Operation::RemoveReference(topic) => {
2007                    state.remove_reference(topic);
2008                }
2009                Operation::Clear => state.clear(),
2010            }
2011        }
2012
2013        fn apply_model_operation(model: &mut AHashMap<String, ModelState>, op: &Operation) {
2014            match op {
2015                Operation::MarkSubscribe(topic) => {
2016                    if model.get(topic) != Some(&ModelState::Confirmed) {
2017                        model.insert(topic.clone(), ModelState::PendingSubscribe);
2018                    }
2019                }
2020                Operation::ConfirmSubscribe(topic) => {
2021                    if matches!(
2022                        model.get(topic),
2023                        Some(ModelState::PendingSubscribe | ModelState::Confirmed)
2024                    ) {
2025                        model.insert(topic.clone(), ModelState::Confirmed);
2026                    }
2027                }
2028                Operation::MarkUnsubscribe(topic) => {
2029                    model.insert(topic.clone(), ModelState::PendingUnsubscribe);
2030                }
2031                Operation::ConfirmUnsubscribe(topic) => {
2032                    if model.get(topic) == Some(&ModelState::PendingUnsubscribe) {
2033                        model.remove(topic);
2034                    }
2035                }
2036                Operation::MarkFailure(topic) => {
2037                    if matches!(
2038                        model.get(topic),
2039                        Some(ModelState::PendingSubscribe | ModelState::Confirmed)
2040                    ) {
2041                        model.insert(topic.clone(), ModelState::PendingSubscribe);
2042                    }
2043                }
2044                Operation::AddReference(_) | Operation::RemoveReference(_) => {}
2045                Operation::Clear => model.clear(),
2046            }
2047        }
2048
2049        fn assert_state_matches_model(
2050            state: &SubscriptionState,
2051            model: &AHashMap<String, ModelState>,
2052        ) {
2053            let topics_for = |expected_state| {
2054                model
2055                    .iter()
2056                    .filter(|&(_topic, state)| *state == expected_state)
2057                    .map(|(topic, _state)| topic.clone())
2058                    .collect::<AHashSet<_>>()
2059            };
2060            let confirmed = state
2061                .topics_from_map(&state.confirmed)
2062                .into_iter()
2063                .collect::<AHashSet<_>>();
2064            let pending_subscribe = state
2065                .pending_subscribe_topics()
2066                .into_iter()
2067                .collect::<AHashSet<_>>();
2068            let pending_unsubscribe = state
2069                .pending_unsubscribe_topics()
2070                .into_iter()
2071                .collect::<AHashSet<_>>();
2072            let expected_confirmed = topics_for(ModelState::Confirmed);
2073            let expected_pending_subscribe = topics_for(ModelState::PendingSubscribe);
2074            let expected_pending_unsubscribe = topics_for(ModelState::PendingUnsubscribe);
2075            let expected_all = expected_confirmed
2076                .union(&expected_pending_subscribe)
2077                .cloned()
2078                .collect::<AHashSet<_>>();
2079            let all = state.all_topics().into_iter().collect::<AHashSet<_>>();
2080
2081            assert_eq!(confirmed, expected_confirmed);
2082            assert_eq!(pending_subscribe, expected_pending_subscribe);
2083            assert_eq!(pending_unsubscribe, expected_pending_unsubscribe);
2084            assert_eq!(all, expected_all);
2085            assert_eq!(state.len(), confirmed.len());
2086            assert_eq!(state.is_empty(), model.is_empty());
2087        }
2088
2089        proptest! {
2090            #![proptest_config(ProptestConfig::with_cases(500))]
2091
2092            /// Property: Invariants hold after any sequence of operations.
2093            #[rstest]
2094            fn prop_invariants_hold_after_operations(
2095                operations in prop::collection::vec(operation_strategy(), 1..50)
2096            ) {
2097                let state = SubscriptionState::new('.');
2098                let mut model = AHashMap::new();
2099
2100                for (i, op) in operations.iter().enumerate() {
2101                    apply_operation(&state, op);
2102                    apply_model_operation(&mut model, op);
2103
2104                    check_invariants(&state, &format!("After op {i}: {op:?}"));
2105                    assert_state_matches_model(&state, &model);
2106                }
2107
2108                check_invariants(&state, "Final state");
2109                assert_state_matches_model(&state, &model);
2110            }
2111
2112            /// Reference-count operations match an independent count model.
2113            #[rstest]
2114            fn prop_reference_counting_matches_reference(
2115                ops in prop::collection::vec(
2116                    topic_strategy().prop_flat_map(|t| {
2117                        prop_oneof![
2118                            Just(Operation::AddReference(t.clone())),
2119                            Just(Operation::RemoveReference(t)),
2120                        ]
2121                    }),
2122                    1..100
2123                )
2124            ) {
2125                let state = SubscriptionState::new('.');
2126                let mut expected = AHashMap::new();
2127
2128                for op in &ops {
2129                    match op {
2130                        Operation::AddReference(topic) => {
2131                            let count = expected.entry(topic.clone()).or_insert(0usize);
2132                            let should_subscribe = *count == 0;
2133                            *count += 1;
2134                            prop_assert_eq!(state.add_reference(topic), should_subscribe);
2135                        }
2136                        Operation::RemoveReference(topic) => {
2137                            let count = expected.get(topic).copied().unwrap_or(0);
2138                            let should_unsubscribe = count == 1;
2139                            if should_unsubscribe {
2140                                expected.remove(topic);
2141                            } else if count > 1 {
2142                                *expected.get_mut(topic).unwrap() -= 1;
2143                            }
2144                            prop_assert_eq!(state.remove_reference(topic), should_unsubscribe);
2145                        }
2146                        _ => unreachable!("reference-count strategy only generates reference operations"),
2147                    }
2148
2149                    prop_assert_eq!(state.reference_counts.len(), expected.len());
2150                    for (topic, count) in &expected {
2151                        prop_assert_eq!(state.get_reference_count(topic), *count);
2152                    }
2153                }
2154            }
2155
2156            /// Property: all_topics() always equals confirmed ∪ pending_subscribe.
2157            #[rstest]
2158            fn prop_all_topics_is_union(
2159                operations in prop::collection::vec(operation_strategy(), 1..50)
2160            ) {
2161                let state = SubscriptionState::new('.');
2162
2163                for op in &operations {
2164                    apply_operation(&state, op);
2165
2166                    // Verify all_topics() == confirmed ∪ pending_subscribe
2167                    let all_topics: AHashSet<String> = state.all_topics().into_iter().collect();
2168                    let confirmed: AHashSet<String> = state.topics_from_map(&state.confirmed).into_iter().collect();
2169                    let pending_sub: AHashSet<String> = state.pending_subscribe_topics().into_iter().collect();
2170                    let expected: AHashSet<String> = confirmed.union(&pending_sub).cloned().collect();
2171
2172                    assert_eq!(all_topics, expected);
2173
2174                    // Ensure pending_unsubscribe topics are NOT in all_topics
2175                    let pending_unsub: AHashSet<String> = state.pending_unsubscribe_topics().into_iter().collect();
2176                    for topic in pending_unsub {
2177                        assert!(!all_topics.contains(&topic));
2178                    }
2179                }
2180            }
2181
2182            /// Property: clear() resets to empty state.
2183            #[rstest]
2184            fn prop_clear_resets_completely(
2185                operations in prop::collection::vec(operation_strategy(), 1..30)
2186            ) {
2187                let state = SubscriptionState::new('.');
2188
2189                // Apply random operations
2190                for op in &operations {
2191                    apply_operation(&state, op);
2192                }
2193
2194                // Clear and verify complete reset
2195                state.clear();
2196
2197                assert!(state.is_empty());
2198                assert_eq!(state.len(), 0);
2199                assert!(state.all_topics().is_empty());
2200                assert!(state.pending_subscribe_topics().is_empty());
2201                assert!(state.pending_unsubscribe_topics().is_empty());
2202                assert!(state.confirmed.is_empty());
2203                assert!(state.pending_subscribe.is_empty());
2204                assert!(state.pending_unsubscribe.is_empty());
2205                assert!(state.reference_counts.is_empty());
2206            }
2207
2208            /// Property: Topics are mutually exclusive across states.
2209            #[rstest]
2210            fn prop_topic_mutual_exclusivity(
2211                operations in prop::collection::vec(operation_strategy(), 1..50),
2212                topic in topic_strategy()
2213            ) {
2214                let state = SubscriptionState::new('.');
2215
2216                for (i, op) in operations.iter().enumerate() {
2217                    apply_operation(&state, op);
2218                    check_topic_exclusivity(&state, &topic, &format!("After op {i}: {op:?}"));
2219                }
2220            }
2221        }
2222    }
2223}