1use 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
44pub(crate) static CHANNEL_LEVEL_MARKER: LazyLock<Ustr> = LazyLock::new(|| Ustr::from(""));
49
50#[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#[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 #[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 #[must_use]
130 pub fn delimiter(&self) -> char {
131 self.delimiter
132 }
133
134 #[must_use]
136 pub fn confirmed(&self) -> SubscriptionSnapshot {
137 let _guard = self.lock_state_read();
138 snapshot(&self.confirmed)
139 }
140
141 #[must_use]
143 pub fn pending_subscribe(&self) -> SubscriptionSnapshot {
144 let _guard = self.lock_state_read();
145 snapshot(&self.pending_subscribe)
146 }
147
148 #[must_use]
150 pub fn pending_unsubscribe(&self) -> SubscriptionSnapshot {
151 let _guard = self.lock_state_read();
152 snapshot(&self.pending_unsubscribe)
153 }
154
155 #[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 #[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 #[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 #[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 #[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 #[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 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 is_tracked(&self.confirmed, channel, symbol) {
230 return;
231 }
232
233 untrack_topic(&self.pending_unsubscribe, channel, symbol);
235
236 track_topic(&self.pending_subscribe, channel, symbol);
237 }
238
239 #[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 !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 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 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 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 if !is_tracked(&self.pending_unsubscribe, channel, symbol) {
306 return; }
308
309 untrack_topic(&self.pending_unsubscribe, channel, symbol);
310 untrack_topic(&self.confirmed, channel, symbol);
311 }
313
314 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 #[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 #[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 #[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 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 #[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 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 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 if symbols.contains(&marker) {
456 topics.push(channel.to_string());
457 }
458
459 for symbol in symbols {
461 if *symbol != marker {
462 topics.push(format!("{channel}{}{symbol}", self.delimiter));
463 }
464 }
465 }
466
467 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#[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
498fn 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
508fn untrack_topic(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) {
512 let symbol = topic_symbol(symbol);
513
514 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
524fn 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 let state_dot = SubscriptionState::new('.');
595 let state_colon = SubscriptionState::new(':');
596
597 state_dot.mark_subscribe("channel.SYMBOL");
599 state_colon.mark_subscribe("channel:SYMBOL");
600
601 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 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 state.mark_subscribe("tickers.BTCUSDT");
635 state.confirm_subscribe("tickers.BTCUSDT");
636 assert_eq!(state.len(), 2);
637
638 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 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 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 state.mark_subscribe("tickers.BTCUSDT");
673 state.confirm_subscribe("tickers.BTCUSDT");
674 assert_eq!(state.len(), 1);
675
676 state.mark_subscribe("tickers");
678 state.confirm_subscribe("tickers");
679 assert_eq!(state.len(), 2);
680
681 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 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 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 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 state.mark_subscribe(".channel");
741 state.confirm_subscribe(".channel");
742 assert_eq!(state.len(), 2);
743
744 state.mark_subscribe("channel.");
747 state.confirm_subscribe("channel.");
748 assert_eq!(state.len(), 3);
749
750 state.mark_subscribe("tickers");
752 state.confirm_subscribe("tickers");
753 assert_eq!(state.len(), 4);
754
755 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())); 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 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 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 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 state.mark_subscribe("tickers.BTCUSDT");
918 state.confirm_subscribe("tickers.BTCUSDT");
919
920 state.mark_unsubscribe("tickers.BTCUSDT");
922
923 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); }
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); 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 state.mark_subscribe("tickers.BTCUSDT");
1022 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1023
1024 state.mark_unsubscribe("tickers.BTCUSDT");
1026
1027 assert!(state.pending_subscribe_topics().is_empty());
1029 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1030
1031 state.confirm_unsubscribe("tickers.BTCUSDT");
1033
1034 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 state.mark_subscribe("tickers.BTCUSDT");
1046
1047 state.mark_unsubscribe("tickers.BTCUSDT");
1049
1050 state.confirm_subscribe("tickers.BTCUSDT");
1052
1053 assert_eq!(state.len(), 0);
1055 assert!(state.pending_subscribe_topics().is_empty());
1056
1057 state.confirm_unsubscribe("tickers.BTCUSDT");
1059
1060 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 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 state.mark_subscribe("tickers.BTCUSDT");
1097 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1098
1099 state.confirm_unsubscribe("tickers.BTCUSDT");
1101 assert!(state.pending_unsubscribe_topics().is_empty());
1102 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]); state.confirm_subscribe("tickers.BTCUSDT");
1106 assert_eq!(state.len(), 1);
1107 assert!(state.pending_subscribe_topics().is_empty());
1108
1109 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 let state = SubscriptionState::new('.');
1122
1123 state.mark_subscribe("tickers.BTCUSDT");
1125 state.confirm_subscribe("tickers.BTCUSDT");
1126 assert_eq!(state.len(), 1);
1127
1128 state.mark_unsubscribe("tickers.BTCUSDT");
1130 assert_eq!(state.len(), 0);
1131 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1132
1133 state.mark_subscribe("tickers.BTCUSDT");
1135 assert!(state.pending_unsubscribe_topics().is_empty()); assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1137
1138 state.confirm_subscribe("tickers.BTCUSDT");
1140 assert_eq!(state.len(), 1); assert!(state.pending_subscribe_topics().is_empty());
1142
1143 state.confirm_unsubscribe("tickers.BTCUSDT");
1146
1147 assert_eq!(state.len(), 1); assert!(state.pending_unsubscribe_topics().is_empty());
1150 assert!(state.pending_subscribe_topics().is_empty());
1151
1152 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 state.mark_subscribe("tickers.BTCUSDT");
1164 state.confirm_subscribe("tickers.BTCUSDT");
1165 assert_eq!(state.len(), 1);
1166
1167 state.mark_unsubscribe("tickers.BTCUSDT");
1169
1170 assert_eq!(state.len(), 0);
1172 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1173
1174 state.confirm_subscribe("tickers.BTCUSDT");
1176
1177 state.confirm_unsubscribe("tickers.BTCUSDT");
1179
1180 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 state.confirm_subscribe("tickers.BTCUSDT");
1254 assert_eq!(state.len(), 0);
1255
1256 state.confirm_unsubscribe("tickers.ETHUSDT");
1258 assert_eq!(state.len(), 0); state.mark_subscribe("orderbook");
1262 state.confirm_subscribe("orderbook");
1263 state.confirm_subscribe("orderbook"); assert_eq!(state.len(), 1);
1265
1266 state.mark_unsubscribe("nonexistent");
1268 state.confirm_unsubscribe("nonexistent");
1269 assert_eq!(state.len(), 1); }
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 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 state.mark_failure("tickers.BTCUSDT");
1295
1296 assert_eq!(state.len(), 0);
1298 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1299
1300 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 state.mark_subscribe("tickers.BTCUSDT");
1310 state.confirm_subscribe("tickers.BTCUSDT");
1311 assert_eq!(state.len(), 1);
1312
1313 state.mark_unsubscribe("tickers.BTCUSDT");
1315 assert_eq!(state.len(), 0);
1316 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1317
1318 state.mark_failure("tickers.BTCUSDT");
1320
1321 assert!(state.pending_subscribe_topics().is_empty());
1323 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1324
1325 assert!(state.all_topics().is_empty());
1327
1328 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 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 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 for topic in &topics_to_resubscribe {
1361 state.mark_subscribe(topic);
1362 }
1363
1364 for topic in &topics_to_resubscribe {
1366 state.confirm_subscribe(topic);
1367 }
1368
1369 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 state.add_reference("confirmed.BTCUSDT");
1381 state.mark_subscribe("confirmed.BTCUSDT");
1382 state.confirm_subscribe("confirmed.BTCUSDT");
1383
1384 state.add_reference("pending.ETHUSDT");
1386 state.mark_subscribe("pending.ETHUSDT");
1387
1388 state.mark_subscribe("cancelled.XRPUSDT");
1390 state.confirm_subscribe("cancelled.XRPUSDT");
1391 state.mark_unsubscribe("cancelled.XRPUSDT");
1392
1393 assert_eq!(state.len(), 1); let all = state.all_topics();
1396 assert_eq!(all.len(), 2); assert!(all.contains(&"confirmed.BTCUSDT".to_string()));
1398 assert!(all.contains(&"pending.ETHUSDT".to_string()));
1399 assert!(!all.contains(&"cancelled.XRPUSDT".to_string())); 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 for topic in &topics_to_resubscribe {
1425 state.confirm_subscribe(topic);
1426 }
1427
1428 assert_eq!(state.len(), 2); 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 let should_unsubscribe = state.remove_reference("nonexistent");
1475
1476 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 assert!(!state.remove_reference("never.added"));
1487 assert_eq!(state.get_reference_count("never.added"), 0);
1488
1489 state.add_reference("once.added");
1491 assert_eq!(state.get_reference_count("once.added"), 1);
1492
1493 assert!(state.remove_reference("once.added")); assert_eq!(state.get_reference_count("once.added"), 0);
1495
1496 assert!(!state.remove_reference("once.added")); assert!(!state.remove_reference("once.added")); assert_eq!(state.get_reference_count("once.added"), 0);
1499
1500 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 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 assert_eq!(state.get_reference_count(topic), 100);
1530
1531 for _ in 0..50 {
1533 state.remove_reference(topic);
1534 }
1535
1536 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 for i in 0..10 {
1559 let topic = format!("channel{i}.SYMBOL");
1560 state.add_reference(&topic);
1561 state.add_reference(&topic); state.mark_subscribe(&topic);
1563 state.confirm_subscribe(&topic);
1564 }
1565
1566 assert_eq!(state.len(), 10);
1567 assert!(!state.is_empty());
1568
1569 state.clear();
1571
1572 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 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 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 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 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 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 state_clone.remove_reference(&topic);
1633 });
1634 handles.push(handle);
1635 }
1636
1637 for handle in handles {
1638 handle.await.unwrap();
1639 }
1640
1641 for i in 0..20 {
1643 let topic = format!("tickers.SYMBOL{i}");
1644 assert_eq!(state.get_reference_count(&topic), 1);
1645 }
1646
1647 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 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 state_clone.add_reference(&topic1);
1668 state_clone.add_reference(&topic2);
1669
1670 state_clone.mark_subscribe(&topic1);
1672 state_clone.confirm_subscribe(&topic1);
1673 state_clone.mark_subscribe(&topic2);
1674
1675 if i % 3 == 0 {
1677 state_clone.mark_unsubscribe(&topic1);
1678 state_clone.confirm_unsubscribe(&topic1);
1679 }
1680
1681 state_clone.add_reference(&topic2);
1683 state_clone.remove_reference(&topic2);
1684
1685 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 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); state_clone.mark_subscribe(&topic);
1728 state_clone.confirm_subscribe(&topic);
1729
1730 state_clone.mark_unsubscribe(&topic);
1732 state_clone.mark_subscribe(&topic);
1734 state_clone.confirm_unsubscribe(&topic);
1736 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 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}"); state_clone.mark_subscribe(&topic);
1765 state_clone.confirm_subscribe(&topic);
1766
1767 for _ in 0..5 {
1769 state_clone.mark_failure(&topic);
1770 state_clone.confirm_subscribe(&topic); }
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 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_op(&state, op1, topic);
1803 apply_op(&state, op2, topic);
1804
1805 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 fn check_invariants(state: &SubscriptionState, label: &str) {
1834 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 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 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 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 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 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 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 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 fn topic_strategy() -> impl Strategy<Value = String> {
1970 prop_oneof![
1971 (any::<u8>(), any::<u8>())
1973 .prop_map(|(ch, sym)| { format!("channel{}.SYMBOL{}", ch % 5, sym % 10) }),
1974 any::<u8>().prop_map(|ch| format!("channel{}", ch % 5)),
1976 ]
1977 }
1978
1979 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 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 #[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 #[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 #[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 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 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 #[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 for op in &operations {
2191 apply_operation(&state, op);
2192 }
2193
2194 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 #[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}