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;
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.state_lock.read();
138 snapshot(&self.confirmed)
139 }
140
141 #[must_use]
143 pub fn pending_subscribe(&self) -> SubscriptionSnapshot {
144 let _guard = self.state_lock.read();
145 snapshot(&self.pending_subscribe)
146 }
147
148 #[must_use]
150 pub fn pending_unsubscribe(&self) -> SubscriptionSnapshot {
151 let _guard = self.state_lock.read();
152 snapshot(&self.pending_unsubscribe)
153 }
154
155 #[must_use]
159 pub fn len(&self) -> usize {
160 let _guard = self.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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.state_lock.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
474#[must_use]
476pub fn split_topic(topic: &str, delimiter: char) -> (&str, Option<&str>) {
477 topic
478 .split_once(delimiter)
479 .map_or((topic, None), |(channel, symbol)| (channel, Some(symbol)))
480}
481
482fn snapshot(map: &DashMap<Ustr, AHashSet<Ustr>>) -> SubscriptionSnapshot {
483 SubscriptionSnapshot(
484 map.iter()
485 .map(|entry| (*entry.key(), entry.value().clone()))
486 .collect(),
487 )
488}
489
490fn track_topic(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) -> bool {
495 map.entry(Ustr::from(channel))
496 .or_default()
497 .insert(topic_symbol(symbol))
498}
499
500fn untrack_topic(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) {
504 let symbol = topic_symbol(symbol);
505
506 if let dashmap::mapref::entry::Entry::Occupied(mut entry) = map.entry(Ustr::from(channel)) {
509 entry.get_mut().remove(&symbol);
510 if entry.get().is_empty() {
511 entry.remove();
512 }
513 }
514}
515
516fn is_tracked(map: &DashMap<Ustr, AHashSet<Ustr>>, channel: &str, symbol: Option<&str>) -> bool {
518 let symbol = topic_symbol(symbol);
519 map.get(&Ustr::from(channel))
520 .is_some_and(|entry| entry.contains(&symbol))
521}
522
523fn topic_symbol(symbol: Option<&str>) -> Ustr {
524 symbol.map_or(*CHANNEL_LEVEL_MARKER, Ustr::from)
525}
526
527#[cfg(test)]
528mod tests {
529 use std::sync::{
530 Barrier,
531 atomic::{AtomicUsize, Ordering},
532 };
533
534 use rstest::rstest;
535
536 use super::*;
537
538 #[rstest]
539 fn test_split_topic_with_symbol() {
540 let (channel, symbol) = split_topic("tickers.BTCUSDT", '.');
541 assert_eq!(channel, "tickers");
542 assert_eq!(symbol, Some("BTCUSDT"));
543
544 let (channel, symbol) = split_topic("orderBookL2:XBTUSD", ':');
545 assert_eq!(channel, "orderBookL2");
546 assert_eq!(symbol, Some("XBTUSD"));
547 }
548
549 #[rstest]
550 fn test_split_topic_without_symbol() {
551 let (channel, symbol) = split_topic("orderbook", '.');
552 assert_eq!(channel, "orderbook");
553 assert_eq!(symbol, None);
554 }
555
556 #[rstest]
557 fn test_topic_without_symbol() {
558 let state = SubscriptionState::new('.');
559 state.mark_subscribe("orderbook");
560 state.confirm_subscribe("orderbook");
561
562 assert_eq!(state.len(), 1);
563 assert_eq!(state.all_topics(), vec!["orderbook"]);
564 }
565
566 #[rstest]
567 fn test_different_delimiters() {
568 let state_dot = SubscriptionState::new('.');
569 state_dot.mark_subscribe("tickers.BTCUSDT");
570 assert_eq!(
571 state_dot.pending_subscribe_topics(),
572 vec!["tickers.BTCUSDT"]
573 );
574
575 let state_colon = SubscriptionState::new(':');
576 state_colon.mark_subscribe("orderBookL2:XBTUSD");
577 assert_eq!(
578 state_colon.pending_subscribe_topics(),
579 vec!["orderBookL2:XBTUSD"]
580 );
581 }
582
583 #[rstest]
584 fn test_different_delimiter_does_not_affect_storage() {
585 let state_dot = SubscriptionState::new('.');
587 let state_colon = SubscriptionState::new(':');
588
589 state_dot.mark_subscribe("channel.SYMBOL");
591 state_colon.mark_subscribe("channel:SYMBOL");
592
593 assert_eq!(state_dot.pending_subscribe_topics(), vec!["channel.SYMBOL"]);
595 assert_eq!(
596 state_colon.pending_subscribe_topics(),
597 vec!["channel:SYMBOL"]
598 );
599 }
600
601 #[rstest]
602 fn test_multiple_symbols_same_channel() {
603 let state = SubscriptionState::new('.');
604 state.mark_subscribe("tickers.BTCUSDT");
605 state.mark_subscribe("tickers.ETHUSDT");
606 state.confirm_subscribe("tickers.BTCUSDT");
607 state.confirm_subscribe("tickers.ETHUSDT");
608
609 assert_eq!(state.len(), 2);
610 assert_eq!(
611 state.all_topics(),
612 vec!["tickers.BTCUSDT", "tickers.ETHUSDT"]
613 );
614 }
615
616 #[rstest]
617 fn test_mixed_channel_and_symbol_subscriptions() {
618 let state = SubscriptionState::new('.');
619
620 state.mark_subscribe("tickers");
622 state.confirm_subscribe("tickers");
623 assert_eq!(state.len(), 1);
624 assert_eq!(state.all_topics(), vec!["tickers"]);
625
626 state.mark_subscribe("tickers.BTCUSDT");
628 state.confirm_subscribe("tickers.BTCUSDT");
629 assert_eq!(state.len(), 2);
630
631 assert_eq!(state.all_topics(), vec!["tickers", "tickers.BTCUSDT"]);
633
634 state.mark_subscribe("tickers.ETHUSDT");
636 state.confirm_subscribe("tickers.ETHUSDT");
637 assert_eq!(state.len(), 3);
638
639 assert_eq!(
640 state.all_topics(),
641 vec!["tickers", "tickers.BTCUSDT", "tickers.ETHUSDT"]
642 );
643
644 state.mark_unsubscribe("tickers");
646 state.confirm_unsubscribe("tickers");
647 assert_eq!(state.len(), 2);
648
649 assert_eq!(
650 state.all_topics(),
651 vec!["tickers.BTCUSDT", "tickers.ETHUSDT"]
652 );
653 }
654
655 #[rstest]
656 fn test_symbol_subscription_before_channel() {
657 let state = SubscriptionState::new('.');
658
659 state.mark_subscribe("tickers.BTCUSDT");
661 state.confirm_subscribe("tickers.BTCUSDT");
662 assert_eq!(state.len(), 1);
663
664 state.mark_subscribe("tickers");
666 state.confirm_subscribe("tickers");
667 assert_eq!(state.len(), 2);
668
669 assert_eq!(state.all_topics(), vec!["tickers", "tickers.BTCUSDT"]);
671 }
672
673 #[rstest]
674 fn test_edge_case_empty_channel_name() {
675 let state = SubscriptionState::new('.');
676
677 state.mark_subscribe("");
679 state.confirm_subscribe("");
680
681 assert_eq!(state.len(), 1);
682 assert_eq!(state.all_topics(), vec![""]);
683 }
684
685 #[rstest]
686 fn test_special_characters_in_topics() {
687 let state = SubscriptionState::new('.');
688
689 let special_topics = vec![
691 "channel.symbol-with-dash",
692 "channel.SYMBOL_WITH_UNDERSCORE",
693 "channel.symbol123",
694 "channel.symbol@special",
695 ];
696
697 for topic in &special_topics {
698 state.mark_subscribe(topic);
699 state.confirm_subscribe(topic);
700 }
701
702 assert_eq!(state.len(), special_topics.len());
703
704 let all_topics = state.all_topics();
705
706 for topic in &special_topics {
707 assert!(
708 all_topics.contains(&(*topic).to_string()),
709 "Missing topic: {topic}"
710 );
711 }
712 }
713
714 #[rstest]
715 fn test_edge_case_malformed_topics() {
716 let state = SubscriptionState::new('.');
717
718 state.mark_subscribe("channel.symbol.extra");
720 state.confirm_subscribe("channel.symbol.extra");
721 let topics = state.all_topics();
722 assert!(topics.contains(&"channel.symbol.extra".to_string()));
723
724 state.mark_subscribe(".channel");
726 state.confirm_subscribe(".channel");
727 assert_eq!(state.len(), 2);
728
729 state.mark_subscribe("channel.");
732 state.confirm_subscribe("channel.");
733 assert_eq!(state.len(), 3);
734
735 state.mark_subscribe("tickers");
737 state.confirm_subscribe("tickers");
738 assert_eq!(state.len(), 4);
739
740 let all = state.all_topics();
742 assert_eq!(all.len(), 4);
743 assert!(all.contains(&"channel.symbol.extra".to_string()));
744 assert!(all.contains(&".channel".to_string()));
745 assert!(all.contains(&"channel".to_string())); assert!(all.contains(&"tickers".to_string()));
747 }
748
749 #[rstest]
750 fn test_new_state_is_empty() {
751 let state = SubscriptionState::new('.');
752 assert!(state.is_empty());
753 assert_eq!(state.len(), 0);
754 }
755
756 #[rstest]
757 fn test_subscription_map_accessors_return_isolated_snapshots() {
758 let state = SubscriptionState::new('.');
759 state.mark_subscribe("tickers.BTCUSDT");
760 state.confirm_subscribe("tickers.BTCUSDT");
761
762 let confirmed = state.confirmed();
763 let pending_subscribe = state.pending_subscribe();
764 let pending_unsubscribe = state.pending_unsubscribe();
765
766 state.mark_unsubscribe("tickers.BTCUSDT");
767
768 assert_eq!(
769 confirmed.get(&Ustr::from("tickers")),
770 Some(&AHashSet::from_iter([Ustr::from("BTCUSDT")]))
771 );
772 assert!(pending_subscribe.is_empty());
773 assert!(pending_unsubscribe.is_empty());
774 assert!(state.confirmed().is_empty());
775 assert_eq!(
776 state.pending_unsubscribe().get(&Ustr::from("tickers")),
777 Some(&AHashSet::from_iter([Ustr::from("BTCUSDT")]))
778 );
779 }
780
781 #[rstest]
782 fn test_is_subscribed_empty_state() {
783 let state = SubscriptionState::new('.');
784 let channel = Ustr::from("tickers");
785 let symbol = Ustr::from("BTCUSDT");
786
787 assert!(!state.is_subscribed(&channel, &symbol));
788 }
789
790 #[rstest]
791 fn test_is_subscribed_pending() {
792 let state = SubscriptionState::new('.');
793 let channel = Ustr::from("tickers");
794 let symbol = Ustr::from("BTCUSDT");
795
796 state.mark_subscribe("tickers.BTCUSDT");
797
798 assert!(state.is_subscribed(&channel, &symbol));
799 }
800
801 #[rstest]
802 fn test_is_subscribed_confirmed() {
803 let state = SubscriptionState::new('.');
804 let channel = Ustr::from("tickers");
805 let symbol = Ustr::from("BTCUSDT");
806
807 state.mark_subscribe("tickers.BTCUSDT");
808 state.confirm_subscribe("tickers.BTCUSDT");
809
810 assert!(state.is_subscribed(&channel, &symbol));
811 }
812
813 #[rstest]
814 fn test_is_subscribed_after_unsubscribe() {
815 let state = SubscriptionState::new('.');
816 let channel = Ustr::from("tickers");
817 let symbol = Ustr::from("BTCUSDT");
818
819 state.mark_subscribe("tickers.BTCUSDT");
820 state.confirm_subscribe("tickers.BTCUSDT");
821 state.mark_unsubscribe("tickers.BTCUSDT");
822
823 assert!(!state.is_subscribed(&channel, &symbol));
825 }
826
827 #[rstest]
828 fn test_is_subscribed_after_confirm_unsubscribe() {
829 let state = SubscriptionState::new('.');
830 let channel = Ustr::from("tickers");
831 let symbol = Ustr::from("BTCUSDT");
832
833 state.mark_subscribe("tickers.BTCUSDT");
834 state.confirm_subscribe("tickers.BTCUSDT");
835 state.mark_unsubscribe("tickers.BTCUSDT");
836 state.confirm_unsubscribe("tickers.BTCUSDT");
837
838 assert!(!state.is_subscribed(&channel, &symbol));
839 }
840
841 #[rstest]
842 fn test_all_topics_includes_confirmed_and_pending_subscribe() {
843 let state = SubscriptionState::new('.');
844 state.mark_subscribe("tickers.BTCUSDT");
845 state.confirm_subscribe("tickers.BTCUSDT");
846 state.mark_subscribe("tickers.ETHUSDT");
847
848 assert_eq!(
849 state.all_topics(),
850 vec!["tickers.BTCUSDT", "tickers.ETHUSDT"]
851 );
852 }
853
854 #[rstest]
855 fn test_all_topics_excludes_pending_unsubscribe() {
856 let state = SubscriptionState::new('.');
857 state.mark_subscribe("tickers.BTCUSDT");
858 state.confirm_subscribe("tickers.BTCUSDT");
859 state.mark_unsubscribe("tickers.BTCUSDT");
860
861 let topics = state.all_topics();
862 assert!(topics.is_empty());
863 }
864
865 #[rstest]
866 fn test_all_topics_is_sorted_within_each_group() {
867 let state = SubscriptionState::new('.');
868
869 for topic in [
871 "trades.SOLUSDT",
872 "tickers.ETHUSDT",
873 "trades.BTCUSDT",
874 "tickers.BTCUSDT",
875 ] {
876 state.mark_subscribe(topic);
877 state.confirm_subscribe(topic);
878 }
879
880 state.mark_subscribe("orders.XRPUSDT");
881 state.mark_subscribe("orders.ADAUSDT");
882
883 assert_eq!(
885 state.all_topics(),
886 vec![
887 "tickers.BTCUSDT",
888 "tickers.ETHUSDT",
889 "trades.BTCUSDT",
890 "trades.SOLUSDT",
891 "orders.ADAUSDT",
892 "orders.XRPUSDT",
893 ]
894 );
895 }
896
897 #[rstest]
898 fn test_pending_subscribe_excludes_pending_unsubscribe() {
899 let state = SubscriptionState::new('.');
900
901 state.mark_subscribe("tickers.BTCUSDT");
903 state.confirm_subscribe("tickers.BTCUSDT");
904
905 state.mark_unsubscribe("tickers.BTCUSDT");
907
908 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
910 assert!(state.all_topics().is_empty());
911 assert_eq!(state.len(), 0);
912 }
913
914 #[rstest]
915 fn test_mark_subscribe() {
916 let state = SubscriptionState::new('.');
917 state.mark_subscribe("tickers.BTCUSDT");
918
919 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
920 assert_eq!(state.len(), 0); }
922
923 #[rstest]
924 fn test_try_mark_subscribe_returns_true_once_across_concurrent_calls() {
925 const CALLERS: usize = 32;
926
927 let state = Arc::new(SubscriptionState::new('.'));
928 let start = Arc::new(Barrier::new(CALLERS));
929 let send_count = Arc::new(AtomicUsize::new(0));
930
931 std::thread::scope(|scope| {
932 for _ in 0..CALLERS {
933 let state = Arc::clone(&state);
934 let start = Arc::clone(&start);
935 let send_count = Arc::clone(&send_count);
936
937 scope.spawn(move || {
938 start.wait();
939
940 if state.try_mark_subscribe("tickers.BTCUSDT") {
941 send_count.fetch_add(1, Ordering::SeqCst);
942 }
943 });
944 }
945 });
946
947 assert_eq!(send_count.load(Ordering::SeqCst), 1);
948 assert_eq!(state.pending_subscribe_topics(), ["tickers.BTCUSDT"]);
949 assert!(state.pending_unsubscribe_topics().is_empty());
950 }
951
952 #[rstest]
953 fn test_try_mark_subscribe_respects_lifecycle_state() {
954 let state = SubscriptionState::new('.');
955 let topic = "tickers.BTCUSDT";
956
957 assert!(state.try_mark_subscribe(topic));
958 assert!(!state.try_mark_subscribe(topic));
959
960 state.confirm_subscribe(topic);
961 assert!(!state.try_mark_subscribe(topic));
962
963 state.mark_unsubscribe(topic);
964 assert!(state.try_mark_subscribe(topic));
965 assert_eq!(state.pending_subscribe_topics(), [topic]);
966 assert!(state.pending_unsubscribe_topics().is_empty());
967 }
968
969 #[rstest]
970 fn test_confirm_subscribe() {
971 let state = SubscriptionState::new('.');
972 state.mark_subscribe("tickers.BTCUSDT");
973 state.confirm_subscribe("tickers.BTCUSDT");
974
975 assert!(state.pending_subscribe_topics().is_empty());
976 assert_eq!(state.len(), 1);
977 }
978
979 #[rstest]
980 fn test_mark_unsubscribe() {
981 let state = SubscriptionState::new('.');
982 state.mark_subscribe("tickers.BTCUSDT");
983 state.confirm_subscribe("tickers.BTCUSDT");
984 state.mark_unsubscribe("tickers.BTCUSDT");
985
986 assert_eq!(state.len(), 0); assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
988 }
989
990 #[rstest]
991 fn test_confirm_unsubscribe() {
992 let state = SubscriptionState::new('.');
993 state.mark_subscribe("tickers.BTCUSDT");
994 state.confirm_subscribe("tickers.BTCUSDT");
995 state.mark_unsubscribe("tickers.BTCUSDT");
996 state.confirm_unsubscribe("tickers.BTCUSDT");
997
998 assert!(state.is_empty());
999 }
1000
1001 #[rstest]
1002 fn test_unsubscribe_before_subscribe_confirmed() {
1003 let state = SubscriptionState::new('.');
1004
1005 state.mark_subscribe("tickers.BTCUSDT");
1007 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1008
1009 state.mark_unsubscribe("tickers.BTCUSDT");
1011
1012 assert!(state.pending_subscribe_topics().is_empty());
1014 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1015
1016 state.confirm_unsubscribe("tickers.BTCUSDT");
1018
1019 assert!(state.is_empty());
1021 assert!(state.all_topics().is_empty());
1022 assert_eq!(state.len(), 0);
1023 }
1024
1025 #[rstest]
1026 fn test_late_subscribe_confirmation_after_unsubscribe() {
1027 let state = SubscriptionState::new('.');
1028
1029 state.mark_subscribe("tickers.BTCUSDT");
1031
1032 state.mark_unsubscribe("tickers.BTCUSDT");
1034
1035 state.confirm_subscribe("tickers.BTCUSDT");
1037
1038 assert_eq!(state.len(), 0);
1040 assert!(state.pending_subscribe_topics().is_empty());
1041
1042 state.confirm_unsubscribe("tickers.BTCUSDT");
1044
1045 assert!(state.is_empty());
1047 assert!(state.all_topics().is_empty());
1048 }
1049
1050 #[rstest]
1051 fn test_late_subscribe_ack_after_unsubscribe_ack_does_not_restore_topic() {
1052 let state = SubscriptionState::new('.');
1053 state.mark_subscribe("tickers.BTCUSDT");
1054 state.mark_unsubscribe("tickers.BTCUSDT");
1055
1056 state.confirm_unsubscribe("tickers.BTCUSDT");
1057 state.confirm_subscribe("tickers.BTCUSDT");
1058
1059 assert!(state.is_empty());
1060 assert!(state.all_topics().is_empty());
1061 assert!(state.pending_subscribe_topics().is_empty());
1062 assert!(state.pending_unsubscribe_topics().is_empty());
1063 }
1064
1065 #[rstest]
1066 fn test_resubscribe_before_unsubscribe_ack() {
1067 let state = SubscriptionState::new('.');
1071
1072 state.mark_subscribe("tickers.BTCUSDT");
1073 state.confirm_subscribe("tickers.BTCUSDT");
1074 assert_eq!(state.len(), 1);
1075
1076 state.mark_unsubscribe("tickers.BTCUSDT");
1077 assert_eq!(state.len(), 0);
1078 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1079
1080 state.mark_subscribe("tickers.BTCUSDT");
1082 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1083
1084 state.confirm_unsubscribe("tickers.BTCUSDT");
1086 assert!(state.pending_unsubscribe_topics().is_empty());
1087 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]); state.confirm_subscribe("tickers.BTCUSDT");
1091 assert_eq!(state.len(), 1);
1092 assert!(state.pending_subscribe_topics().is_empty());
1093
1094 let all = state.all_topics();
1096 assert_eq!(all.len(), 1);
1097 assert!(all.contains(&"tickers.BTCUSDT".to_string()));
1098 }
1099
1100 #[rstest]
1101 fn test_stale_unsubscribe_ack_after_resubscribe_confirmed() {
1102 let state = SubscriptionState::new('.');
1107
1108 state.mark_subscribe("tickers.BTCUSDT");
1110 state.confirm_subscribe("tickers.BTCUSDT");
1111 assert_eq!(state.len(), 1);
1112
1113 state.mark_unsubscribe("tickers.BTCUSDT");
1115 assert_eq!(state.len(), 0);
1116 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1117
1118 state.mark_subscribe("tickers.BTCUSDT");
1120 assert!(state.pending_unsubscribe_topics().is_empty()); assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1122
1123 state.confirm_subscribe("tickers.BTCUSDT");
1125 assert_eq!(state.len(), 1); assert!(state.pending_subscribe_topics().is_empty());
1127
1128 state.confirm_unsubscribe("tickers.BTCUSDT");
1131
1132 assert_eq!(state.len(), 1); assert!(state.pending_unsubscribe_topics().is_empty());
1135 assert!(state.pending_subscribe_topics().is_empty());
1136
1137 let all = state.all_topics();
1139 assert_eq!(all.len(), 1);
1140 assert!(all.contains(&"tickers.BTCUSDT".to_string()));
1141 }
1142
1143 #[rstest]
1144 fn test_unsubscribe_clears_all_states() {
1145 let state = SubscriptionState::new('.');
1146
1147 state.mark_subscribe("tickers.BTCUSDT");
1149 state.confirm_subscribe("tickers.BTCUSDT");
1150 assert_eq!(state.len(), 1);
1151
1152 state.mark_unsubscribe("tickers.BTCUSDT");
1154
1155 assert_eq!(state.len(), 0);
1157 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1158
1159 state.confirm_subscribe("tickers.BTCUSDT");
1161
1162 state.confirm_unsubscribe("tickers.BTCUSDT");
1164
1165 assert!(state.is_empty());
1167 assert_eq!(state.len(), 0);
1168 assert!(state.pending_subscribe_topics().is_empty());
1169 assert!(state.pending_unsubscribe_topics().is_empty());
1170 assert!(state.all_topics().is_empty());
1171 }
1172
1173 #[rstest]
1174 fn test_concurrent_subscribe_confirmation_and_unsubscribe_preserve_intent() {
1175 const ITERATIONS: usize = 10_000;
1176
1177 let state = Arc::new(SubscriptionState::new('.'));
1178 let start = Arc::new(Barrier::new(3));
1179 let finish = Arc::new(Barrier::new(3));
1180 let failed_iteration = Arc::new(AtomicUsize::new(usize::MAX));
1181
1182 std::thread::scope(|scope| {
1183 let confirming_state = Arc::clone(&state);
1184 let confirming_start = Arc::clone(&start);
1185 let confirming_finish = Arc::clone(&finish);
1186
1187 scope.spawn(move || {
1188 for _ in 0..ITERATIONS {
1189 confirming_start.wait();
1190 confirming_state.confirm_subscribe("tickers.BTCUSDT");
1191 confirming_finish.wait();
1192 }
1193 });
1194
1195 let unsubscribing_state = Arc::clone(&state);
1196 let unsubscribing_start = Arc::clone(&start);
1197 let unsubscribing_finish = Arc::clone(&finish);
1198
1199 scope.spawn(move || {
1200 for _ in 0..ITERATIONS {
1201 unsubscribing_start.wait();
1202 unsubscribing_state.mark_unsubscribe("tickers.BTCUSDT");
1203 unsubscribing_finish.wait();
1204 }
1205 });
1206
1207 for iteration in 0..ITERATIONS {
1208 state.clear();
1209 state.mark_subscribe("tickers.BTCUSDT");
1210 start.wait();
1211 finish.wait();
1212
1213 if !state.all_topics().is_empty()
1214 || state.pending_unsubscribe_topics() != ["tickers.BTCUSDT"]
1215 {
1216 _ = failed_iteration.compare_exchange(
1217 usize::MAX,
1218 iteration,
1219 Ordering::SeqCst,
1220 Ordering::SeqCst,
1221 );
1222 }
1223 }
1224 });
1225
1226 assert_eq!(
1227 failed_iteration.load(Ordering::SeqCst),
1228 usize::MAX,
1229 "late subscribe confirmation restored the unsubscribe intent"
1230 );
1231 }
1232
1233 #[rstest]
1234 fn test_state_machine_invalid_transitions() {
1235 let state = SubscriptionState::new('.');
1236
1237 state.confirm_subscribe("tickers.BTCUSDT");
1239 assert_eq!(state.len(), 0);
1240
1241 state.confirm_unsubscribe("tickers.ETHUSDT");
1243 assert_eq!(state.len(), 0); state.mark_subscribe("orderbook");
1247 state.confirm_subscribe("orderbook");
1248 state.confirm_subscribe("orderbook"); assert_eq!(state.len(), 1);
1250
1251 state.mark_unsubscribe("nonexistent");
1253 state.confirm_unsubscribe("nonexistent");
1254 assert_eq!(state.len(), 1); }
1256
1257 #[rstest]
1258 fn test_mark_failure() {
1259 let state = SubscriptionState::new('.');
1260 state.mark_subscribe("tickers.BTCUSDT");
1261 state.confirm_subscribe("tickers.BTCUSDT");
1262 state.mark_failure("tickers.BTCUSDT");
1263
1264 assert_eq!(state.len(), 0);
1265 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1266 }
1267
1268 #[rstest]
1269 fn test_mark_failure_moves_to_pending() {
1270 let state = SubscriptionState::new('.');
1271
1272 state.mark_subscribe("tickers.BTCUSDT");
1274 state.confirm_subscribe("tickers.BTCUSDT");
1275 assert_eq!(state.len(), 1);
1276 assert!(state.pending_subscribe_topics().is_empty());
1277
1278 state.mark_failure("tickers.BTCUSDT");
1280
1281 assert_eq!(state.len(), 0);
1283 assert_eq!(state.pending_subscribe_topics(), vec!["tickers.BTCUSDT"]);
1284
1285 assert_eq!(state.all_topics(), vec!["tickers.BTCUSDT"]);
1287 }
1288
1289 #[rstest]
1290 fn test_mark_failure_respects_pending_unsubscribe() {
1291 let state = SubscriptionState::new('.');
1292
1293 state.mark_subscribe("tickers.BTCUSDT");
1295 state.confirm_subscribe("tickers.BTCUSDT");
1296 assert_eq!(state.len(), 1);
1297
1298 state.mark_unsubscribe("tickers.BTCUSDT");
1300 assert_eq!(state.len(), 0);
1301 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1302
1303 state.mark_failure("tickers.BTCUSDT");
1305
1306 assert!(state.pending_subscribe_topics().is_empty());
1308 assert_eq!(state.pending_unsubscribe_topics(), vec!["tickers.BTCUSDT"]);
1309
1310 assert!(state.all_topics().is_empty());
1312
1313 state.confirm_unsubscribe("tickers.BTCUSDT");
1315 assert!(state.is_empty());
1316 }
1317
1318 #[rstest]
1319 fn test_reconnection_scenario() {
1320 let state = SubscriptionState::new('.');
1321
1322 state.add_reference("tickers.BTCUSDT");
1324 state.mark_subscribe("tickers.BTCUSDT");
1325 state.confirm_subscribe("tickers.BTCUSDT");
1326
1327 state.add_reference("tickers.ETHUSDT");
1328 state.mark_subscribe("tickers.ETHUSDT");
1329 state.confirm_subscribe("tickers.ETHUSDT");
1330
1331 state.add_reference("orderbook");
1332 state.mark_subscribe("orderbook");
1333 state.confirm_subscribe("orderbook");
1334
1335 assert_eq!(state.len(), 3);
1336
1337 let topics_to_resubscribe = state.all_topics();
1339 assert_eq!(topics_to_resubscribe.len(), 3);
1340 assert!(topics_to_resubscribe.contains(&"tickers.BTCUSDT".to_string()));
1341 assert!(topics_to_resubscribe.contains(&"tickers.ETHUSDT".to_string()));
1342 assert!(topics_to_resubscribe.contains(&"orderbook".to_string()));
1343
1344 for topic in &topics_to_resubscribe {
1346 state.mark_subscribe(topic);
1347 }
1348
1349 for topic in &topics_to_resubscribe {
1351 state.confirm_subscribe(topic);
1352 }
1353
1354 assert_eq!(state.len(), 3);
1356 assert_eq!(state.all_topics().len(), 3);
1357 }
1358
1359 #[rstest]
1360 fn test_reconnection_with_partial_state() {
1361 let state = SubscriptionState::new('.');
1362
1363 state.add_reference("confirmed.BTCUSDT");
1366 state.mark_subscribe("confirmed.BTCUSDT");
1367 state.confirm_subscribe("confirmed.BTCUSDT");
1368
1369 state.add_reference("pending.ETHUSDT");
1371 state.mark_subscribe("pending.ETHUSDT");
1372
1373 state.mark_subscribe("cancelled.XRPUSDT");
1375 state.confirm_subscribe("cancelled.XRPUSDT");
1376 state.mark_unsubscribe("cancelled.XRPUSDT");
1377
1378 assert_eq!(state.len(), 1); let all = state.all_topics();
1381 assert_eq!(all.len(), 2); assert!(all.contains(&"confirmed.BTCUSDT".to_string()));
1383 assert!(all.contains(&"pending.ETHUSDT".to_string()));
1384 assert!(!all.contains(&"cancelled.XRPUSDT".to_string())); let topics_to_resubscribe = state.reset_after_reconnect();
1388 assert_eq!(
1389 topics_to_resubscribe,
1390 [
1391 "confirmed.BTCUSDT".to_string(),
1392 "pending.ETHUSDT".to_string()
1393 ]
1394 );
1395 assert_eq!(state.reset_after_reconnect(), topics_to_resubscribe);
1396 assert_eq!(state.len(), 0);
1397 assert_eq!(
1398 state.pending_subscribe_topics(),
1399 [
1400 "confirmed.BTCUSDT".to_string(),
1401 "pending.ETHUSDT".to_string()
1402 ]
1403 );
1404 assert!(state.pending_unsubscribe_topics().is_empty());
1405 assert_eq!(state.get_reference_count("confirmed.BTCUSDT"), 1);
1406 assert_eq!(state.get_reference_count("pending.ETHUSDT"), 1);
1407
1408 for topic in &topics_to_resubscribe {
1410 state.confirm_subscribe(topic);
1411 }
1412
1413 assert_eq!(state.len(), 2); let final_topics = state.all_topics();
1416 assert_eq!(final_topics.len(), 2);
1417 assert!(final_topics.contains(&"confirmed.BTCUSDT".to_string()));
1418 assert!(final_topics.contains(&"pending.ETHUSDT".to_string()));
1419 assert!(!final_topics.contains(&"cancelled.XRPUSDT".to_string()));
1420 }
1421
1422 #[rstest]
1423 fn test_reference_counting_single_topic() {
1424 let state = SubscriptionState::new('.');
1425
1426 assert!(state.add_reference("tickers.BTCUSDT"));
1427 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 1);
1428
1429 assert!(!state.add_reference("tickers.BTCUSDT"));
1430 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 2);
1431
1432 assert!(!state.remove_reference("tickers.BTCUSDT"));
1433 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 1);
1434
1435 assert!(state.remove_reference("tickers.BTCUSDT"));
1436 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 0);
1437 }
1438
1439 #[rstest]
1440 fn test_reference_counting_multiple_topics() {
1441 let state = SubscriptionState::new('.');
1442
1443 assert!(state.add_reference("tickers.BTCUSDT"));
1444 assert!(state.add_reference("tickers.ETHUSDT"));
1445
1446 assert!(!state.add_reference("tickers.BTCUSDT"));
1447 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 2);
1448 assert_eq!(state.get_reference_count("tickers.ETHUSDT"), 1);
1449
1450 assert!(!state.remove_reference("tickers.BTCUSDT"));
1451 assert!(state.remove_reference("tickers.ETHUSDT"));
1452 }
1453
1454 #[rstest]
1455 fn test_remove_reference_nonexistent_topic() {
1456 let state = SubscriptionState::new('.');
1457
1458 let should_unsubscribe = state.remove_reference("nonexistent");
1460
1461 assert!(!should_unsubscribe);
1463 assert_eq!(state.get_reference_count("nonexistent"), 0);
1464 }
1465
1466 #[rstest]
1467 fn test_reference_count_underflow_safety() {
1468 let state = SubscriptionState::new('.');
1469
1470 assert!(!state.remove_reference("never.added"));
1472 assert_eq!(state.get_reference_count("never.added"), 0);
1473
1474 state.add_reference("once.added");
1476 assert_eq!(state.get_reference_count("once.added"), 1);
1477
1478 assert!(state.remove_reference("once.added")); assert_eq!(state.get_reference_count("once.added"), 0);
1480
1481 assert!(!state.remove_reference("once.added")); assert!(!state.remove_reference("once.added")); assert_eq!(state.get_reference_count("once.added"), 0);
1484
1485 assert!(state.add_reference("once.added"));
1487 assert_eq!(state.get_reference_count("once.added"), 1);
1488 }
1489
1490 #[rstest]
1491 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1492 async fn test_concurrent_reference_counting_same_topic() {
1493 let state = Arc::new(SubscriptionState::new('.'));
1494 let topic = "tickers.BTCUSDT";
1495 let mut handles = vec![];
1496
1497 for _ in 0..10 {
1499 let state_clone = Arc::clone(&state);
1500
1501 let handle = tokio::spawn(async move {
1502 for _ in 0..10 {
1503 state_clone.add_reference(topic);
1504 }
1505 });
1506 handles.push(handle);
1507 }
1508
1509 for handle in handles {
1510 handle.await.unwrap();
1511 }
1512
1513 assert_eq!(state.get_reference_count(topic), 100);
1515
1516 for _ in 0..50 {
1518 state.remove_reference(topic);
1519 }
1520
1521 assert_eq!(state.get_reference_count(topic), 50);
1523 }
1524
1525 #[rstest]
1526 fn test_clear() {
1527 let state = SubscriptionState::new('.');
1528 state.mark_subscribe("tickers.BTCUSDT");
1529 state.confirm_subscribe("tickers.BTCUSDT");
1530 state.add_reference("tickers.BTCUSDT");
1531
1532 state.clear();
1533
1534 assert!(state.is_empty());
1535 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 0);
1536 }
1537
1538 #[rstest]
1539 fn test_clear_resets_all_state() {
1540 let state = SubscriptionState::new('.');
1541
1542 for i in 0..10 {
1544 let topic = format!("channel{i}.SYMBOL");
1545 state.add_reference(&topic);
1546 state.add_reference(&topic); state.mark_subscribe(&topic);
1548 state.confirm_subscribe(&topic);
1549 }
1550
1551 assert_eq!(state.len(), 10);
1552 assert!(!state.is_empty());
1553
1554 state.clear();
1556
1557 assert_eq!(state.len(), 0);
1559 assert!(state.is_empty());
1560 assert!(state.all_topics().is_empty());
1561 assert!(state.pending_subscribe_topics().is_empty());
1562 assert!(state.pending_unsubscribe_topics().is_empty());
1563
1564 for i in 0..10 {
1566 let topic = format!("channel{i}.SYMBOL");
1567 assert_eq!(state.get_reference_count(&topic), 0);
1568 }
1569 }
1570
1571 #[rstest]
1572 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1573 async fn test_concurrent_subscribe_same_topic() {
1574 let state = Arc::new(SubscriptionState::new('.'));
1575 let mut handles = vec![];
1576
1577 for _ in 0..10 {
1579 let state_clone = Arc::clone(&state);
1580 let handle = tokio::spawn(async move {
1581 state_clone.add_reference("tickers.BTCUSDT");
1582 state_clone.mark_subscribe("tickers.BTCUSDT");
1583 state_clone.confirm_subscribe("tickers.BTCUSDT");
1584 });
1585 handles.push(handle);
1586 }
1587
1588 for handle in handles {
1589 handle.await.unwrap();
1590 }
1591
1592 assert_eq!(state.get_reference_count("tickers.BTCUSDT"), 10);
1594 assert_eq!(state.len(), 1);
1595 }
1596
1597 #[rstest]
1598 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1599 async fn test_concurrent_subscribe_unsubscribe() {
1600 let state = Arc::new(SubscriptionState::new('.'));
1601 let mut handles = vec![];
1602
1603 for i in 0..20 {
1606 let state_clone = Arc::clone(&state);
1607
1608 let handle = tokio::spawn(async move {
1609 let topic = format!("tickers.SYMBOL{i}");
1610 state_clone.add_reference(&topic);
1612 state_clone.add_reference(&topic);
1613 state_clone.mark_subscribe(&topic);
1614 state_clone.confirm_subscribe(&topic);
1615
1616 state_clone.remove_reference(&topic);
1618 });
1619 handles.push(handle);
1620 }
1621
1622 for handle in handles {
1623 handle.await.unwrap();
1624 }
1625
1626 for i in 0..20 {
1628 let topic = format!("tickers.SYMBOL{i}");
1629 assert_eq!(state.get_reference_count(&topic), 1);
1630 }
1631
1632 assert_eq!(state.len(), 20);
1634 assert!(!state.is_empty());
1635 }
1636
1637 #[rstest]
1638 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1639 async fn test_concurrent_stress_mixed_operations() {
1640 let state = Arc::new(SubscriptionState::new('.'));
1641 let mut handles = vec![];
1642
1643 for i in 0..50 {
1645 let state_clone = Arc::clone(&state);
1646
1647 let handle = tokio::spawn(async move {
1648 let topic1 = format!("channel.SYMBOL{i}");
1649 let topic2 = format!("channel.SYMBOL{}", i + 100);
1650
1651 state_clone.add_reference(&topic1);
1653 state_clone.add_reference(&topic2);
1654
1655 state_clone.mark_subscribe(&topic1);
1657 state_clone.confirm_subscribe(&topic1);
1658 state_clone.mark_subscribe(&topic2);
1659
1660 if i % 3 == 0 {
1662 state_clone.mark_unsubscribe(&topic1);
1663 state_clone.confirm_unsubscribe(&topic1);
1664 }
1665
1666 state_clone.add_reference(&topic2);
1668 state_clone.remove_reference(&topic2);
1669
1670 state_clone.confirm_subscribe(&topic2);
1672 });
1673 handles.push(handle);
1674 }
1675
1676 for handle in handles {
1677 handle.await.unwrap();
1678 }
1679
1680 let actual = state.all_topics().into_iter().collect::<AHashSet<_>>();
1681 let expected = (0..50)
1682 .flat_map(|i| {
1683 let topic2 = format!("channel.SYMBOL{}", i + 100);
1684 (i % 3 != 0)
1685 .then(|| format!("channel.SYMBOL{i}"))
1686 .into_iter()
1687 .chain(std::iter::once(topic2))
1688 })
1689 .collect::<AHashSet<_>>();
1690
1691 assert_eq!(actual, expected);
1692 assert_eq!(state.len(), 83);
1693 assert!(state.pending_subscribe_topics().is_empty());
1694 assert!(state.pending_unsubscribe_topics().is_empty());
1695 assert_eq!(state.reference_counts.len(), 100);
1696 }
1697
1698 #[rstest]
1699 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1700 async fn test_stress_rapid_resubscribe_pattern() {
1701 let state = Arc::new(SubscriptionState::new('.'));
1703 let mut handles = vec![];
1704
1705 for i in 0..100 {
1706 let state_clone = Arc::clone(&state);
1707
1708 let handle = tokio::spawn(async move {
1709 let topic = format!("rapid.SYMBOL{}", i % 10); state_clone.mark_subscribe(&topic);
1713 state_clone.confirm_subscribe(&topic);
1714
1715 state_clone.mark_unsubscribe(&topic);
1717 state_clone.mark_subscribe(&topic);
1719 state_clone.confirm_unsubscribe(&topic);
1721 state_clone.confirm_subscribe(&topic);
1723 });
1724 handles.push(handle);
1725 }
1726
1727 for handle in handles {
1728 handle.await.unwrap();
1729 }
1730
1731 check_invariants(&state, "After rapid resubscribe stress test");
1732 }
1733
1734 #[rstest]
1735 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
1736 async fn test_stress_failure_recovery_loop() {
1737 let state = Arc::new(SubscriptionState::new('.'));
1740 let mut handles = vec![];
1741
1742 for i in 0..30 {
1743 let state_clone = Arc::clone(&state);
1744
1745 let handle = tokio::spawn(async move {
1746 let topic = format!("failure.SYMBOL{i}"); state_clone.mark_subscribe(&topic);
1750 state_clone.confirm_subscribe(&topic);
1751
1752 for _ in 0..5 {
1754 state_clone.mark_failure(&topic);
1755 state_clone.confirm_subscribe(&topic); }
1757 });
1758 handles.push(handle);
1759 }
1760
1761 for handle in handles {
1762 handle.await.unwrap();
1763 }
1764
1765 check_invariants(&state, "After failure recovery loops");
1766
1767 assert_eq!(state.len(), 30);
1769 }
1770
1771 #[rstest]
1772 fn test_exhaustive_two_step_transitions() {
1773 let operations = [
1774 "mark_subscribe",
1775 "confirm_subscribe",
1776 "mark_unsubscribe",
1777 "confirm_unsubscribe",
1778 "mark_failure",
1779 ];
1780
1781 for &op1 in &operations {
1782 for &op2 in &operations {
1783 let state = SubscriptionState::new('.');
1784 let topic = "test.TOPIC";
1785
1786 apply_op(&state, op1, topic);
1788 apply_op(&state, op2, topic);
1789
1790 check_invariants(&state, &format!("{op1} -> {op2}"));
1792 check_topic_exclusivity(&state, topic, &format!("{op1} -> {op2}"));
1793 }
1794 }
1795 }
1796
1797 fn apply_op(state: &SubscriptionState, op: &str, topic: &str) {
1798 match op {
1799 "mark_subscribe" => state.mark_subscribe(topic),
1800 "confirm_subscribe" => state.confirm_subscribe(topic),
1801 "mark_unsubscribe" => state.mark_unsubscribe(topic),
1802 "confirm_unsubscribe" => state.confirm_unsubscribe(topic),
1803 "mark_failure" => state.mark_failure(topic),
1804 _ => panic!("Unknown operation: {op}"),
1805 }
1806 }
1807
1808 fn check_invariants(state: &SubscriptionState, label: &str) {
1819 let confirmed_topics: AHashSet<String> = state
1821 .topics_from_map(&state.confirmed)
1822 .into_iter()
1823 .collect();
1824 let pending_sub_topics: AHashSet<String> =
1825 state.pending_subscribe_topics().into_iter().collect();
1826 let pending_unsub_topics: AHashSet<String> =
1827 state.pending_unsubscribe_topics().into_iter().collect();
1828
1829 let confirmed_and_pending_sub: Vec<_> =
1831 confirmed_topics.intersection(&pending_sub_topics).collect();
1832 assert!(
1833 confirmed_and_pending_sub.is_empty(),
1834 "{label}: Topic in both confirmed and pending_subscribe: {confirmed_and_pending_sub:?}"
1835 );
1836
1837 let confirmed_and_pending_unsub: Vec<_> = confirmed_topics
1838 .intersection(&pending_unsub_topics)
1839 .collect();
1840 assert!(
1841 confirmed_and_pending_unsub.is_empty(),
1842 "{label}: Topic in both confirmed and pending_unsubscribe: {confirmed_and_pending_unsub:?}"
1843 );
1844
1845 let pending_sub_and_unsub: Vec<_> = pending_sub_topics
1846 .intersection(&pending_unsub_topics)
1847 .collect();
1848 assert!(
1849 pending_sub_and_unsub.is_empty(),
1850 "{label}: Topic in both pending_subscribe and pending_unsubscribe: {pending_sub_and_unsub:?}"
1851 );
1852
1853 let all_topics: AHashSet<String> = state.all_topics().into_iter().collect();
1855 let expected_all: AHashSet<String> = confirmed_topics
1856 .union(&pending_sub_topics)
1857 .cloned()
1858 .collect();
1859 assert_eq!(
1860 all_topics, expected_all,
1861 "{label}: all_topics() doesn't match confirmed ∪ pending_subscribe"
1862 );
1863
1864 for topic in &pending_unsub_topics {
1866 assert!(
1867 !all_topics.contains(topic),
1868 "{label}: pending_unsubscribe topic {topic} incorrectly in all_topics()"
1869 );
1870 }
1871
1872 let expected_len: usize = state
1874 .confirmed
1875 .iter()
1876 .map(|entry| entry.value().len())
1877 .sum();
1878 assert_eq!(
1879 state.len(),
1880 expected_len,
1881 "{label}: len() mismatch. Expected {expected_len}, was {}",
1882 state.len()
1883 );
1884
1885 let should_be_empty = state.confirmed.is_empty()
1887 && pending_sub_topics.is_empty()
1888 && pending_unsub_topics.is_empty();
1889 assert_eq!(
1890 state.is_empty(),
1891 should_be_empty,
1892 "{label}: is_empty() inconsistent. Maps empty: {should_be_empty}, is_empty(): {}",
1893 state.is_empty()
1894 );
1895
1896 for entry in state.reference_counts.iter() {
1898 let count = entry.value().get();
1899 assert!(
1900 count > 0,
1901 "{label}: Reference count should be NonZeroUsize (> 0), was {count} for {:?}",
1902 entry.key()
1903 );
1904 }
1905 }
1906
1907 fn check_topic_exclusivity(state: &SubscriptionState, topic: &str, label: &str) {
1909 let (channel, symbol) = split_topic(topic, state.delimiter);
1910
1911 let in_confirmed = is_tracked(&state.confirmed, channel, symbol);
1912 let in_pending_sub = is_tracked(&state.pending_subscribe, channel, symbol);
1913 let in_pending_unsub = is_tracked(&state.pending_unsubscribe, channel, symbol);
1914
1915 let count = [in_confirmed, in_pending_sub, in_pending_unsub]
1916 .iter()
1917 .filter(|&&x| x)
1918 .count();
1919
1920 assert!(
1921 count <= 1,
1922 "{label}: Topic {topic} in {count} states (should be 0 or 1). \
1923 confirmed: {in_confirmed}, pending_sub: {in_pending_sub}, pending_unsub: {in_pending_unsub}"
1924 );
1925 }
1926
1927 #[cfg(test)]
1928 mod property_tests {
1929 use ahash::AHashMap;
1930 use proptest::prelude::*;
1931
1932 use super::*;
1933
1934 #[derive(Debug, Clone)]
1935 enum Operation {
1936 MarkSubscribe(String),
1937 ConfirmSubscribe(String),
1938 MarkUnsubscribe(String),
1939 ConfirmUnsubscribe(String),
1940 MarkFailure(String),
1941 AddReference(String),
1942 RemoveReference(String),
1943 Clear,
1944 }
1945
1946 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
1947 enum ModelState {
1948 Confirmed,
1949 PendingSubscribe,
1950 PendingUnsubscribe,
1951 }
1952
1953 fn topic_strategy() -> impl Strategy<Value = String> {
1955 prop_oneof![
1956 (any::<u8>(), any::<u8>())
1958 .prop_map(|(ch, sym)| { format!("channel{}.SYMBOL{}", ch % 5, sym % 10) }),
1959 any::<u8>().prop_map(|ch| format!("channel{}", ch % 5)),
1961 ]
1962 }
1963
1964 fn operation_strategy() -> impl Strategy<Value = Operation> {
1966 topic_strategy().prop_flat_map(|topic| {
1967 prop_oneof![
1968 Just(Operation::MarkSubscribe(topic.clone())),
1969 Just(Operation::ConfirmSubscribe(topic.clone())),
1970 Just(Operation::MarkUnsubscribe(topic.clone())),
1971 Just(Operation::ConfirmUnsubscribe(topic.clone())),
1972 Just(Operation::MarkFailure(topic.clone())),
1973 Just(Operation::AddReference(topic.clone())),
1974 Just(Operation::RemoveReference(topic)),
1975 Just(Operation::Clear),
1976 ]
1977 })
1978 }
1979
1980 fn apply_operation(state: &SubscriptionState, op: &Operation) {
1982 match op {
1983 Operation::MarkSubscribe(topic) => state.mark_subscribe(topic),
1984 Operation::ConfirmSubscribe(topic) => state.confirm_subscribe(topic),
1985 Operation::MarkUnsubscribe(topic) => state.mark_unsubscribe(topic),
1986 Operation::ConfirmUnsubscribe(topic) => state.confirm_unsubscribe(topic),
1987 Operation::MarkFailure(topic) => state.mark_failure(topic),
1988 Operation::AddReference(topic) => {
1989 state.add_reference(topic);
1990 }
1991 Operation::RemoveReference(topic) => {
1992 state.remove_reference(topic);
1993 }
1994 Operation::Clear => state.clear(),
1995 }
1996 }
1997
1998 fn apply_model_operation(model: &mut AHashMap<String, ModelState>, op: &Operation) {
1999 match op {
2000 Operation::MarkSubscribe(topic) => {
2001 if model.get(topic) != Some(&ModelState::Confirmed) {
2002 model.insert(topic.clone(), ModelState::PendingSubscribe);
2003 }
2004 }
2005 Operation::ConfirmSubscribe(topic) => {
2006 if matches!(
2007 model.get(topic),
2008 Some(ModelState::PendingSubscribe | ModelState::Confirmed)
2009 ) {
2010 model.insert(topic.clone(), ModelState::Confirmed);
2011 }
2012 }
2013 Operation::MarkUnsubscribe(topic) => {
2014 model.insert(topic.clone(), ModelState::PendingUnsubscribe);
2015 }
2016 Operation::ConfirmUnsubscribe(topic) => {
2017 if model.get(topic) == Some(&ModelState::PendingUnsubscribe) {
2018 model.remove(topic);
2019 }
2020 }
2021 Operation::MarkFailure(topic) => {
2022 if matches!(
2023 model.get(topic),
2024 Some(ModelState::PendingSubscribe | ModelState::Confirmed)
2025 ) {
2026 model.insert(topic.clone(), ModelState::PendingSubscribe);
2027 }
2028 }
2029 Operation::AddReference(_) | Operation::RemoveReference(_) => {}
2030 Operation::Clear => model.clear(),
2031 }
2032 }
2033
2034 fn assert_state_matches_model(
2035 state: &SubscriptionState,
2036 model: &AHashMap<String, ModelState>,
2037 ) {
2038 let topics_for = |expected_state| {
2039 model
2040 .iter()
2041 .filter(|&(_topic, state)| *state == expected_state)
2042 .map(|(topic, _state)| topic.clone())
2043 .collect::<AHashSet<_>>()
2044 };
2045 let confirmed = state
2046 .topics_from_map(&state.confirmed)
2047 .into_iter()
2048 .collect::<AHashSet<_>>();
2049 let pending_subscribe = state
2050 .pending_subscribe_topics()
2051 .into_iter()
2052 .collect::<AHashSet<_>>();
2053 let pending_unsubscribe = state
2054 .pending_unsubscribe_topics()
2055 .into_iter()
2056 .collect::<AHashSet<_>>();
2057 let expected_confirmed = topics_for(ModelState::Confirmed);
2058 let expected_pending_subscribe = topics_for(ModelState::PendingSubscribe);
2059 let expected_pending_unsubscribe = topics_for(ModelState::PendingUnsubscribe);
2060 let expected_all = expected_confirmed
2061 .union(&expected_pending_subscribe)
2062 .cloned()
2063 .collect::<AHashSet<_>>();
2064 let all = state.all_topics().into_iter().collect::<AHashSet<_>>();
2065
2066 assert_eq!(confirmed, expected_confirmed);
2067 assert_eq!(pending_subscribe, expected_pending_subscribe);
2068 assert_eq!(pending_unsubscribe, expected_pending_unsubscribe);
2069 assert_eq!(all, expected_all);
2070 assert_eq!(state.len(), confirmed.len());
2071 assert_eq!(state.is_empty(), model.is_empty());
2072 }
2073
2074 proptest! {
2075 #![proptest_config(ProptestConfig::with_cases(500))]
2076
2077 #[rstest]
2079 fn prop_invariants_hold_after_operations(
2080 operations in prop::collection::vec(operation_strategy(), 1..50)
2081 ) {
2082 let state = SubscriptionState::new('.');
2083 let mut model = AHashMap::new();
2084
2085 for (i, op) in operations.iter().enumerate() {
2086 apply_operation(&state, op);
2087 apply_model_operation(&mut model, op);
2088
2089 check_invariants(&state, &format!("After op {i}: {op:?}"));
2090 assert_state_matches_model(&state, &model);
2091 }
2092
2093 check_invariants(&state, "Final state");
2094 assert_state_matches_model(&state, &model);
2095 }
2096
2097 #[rstest]
2099 fn prop_reference_counting_matches_reference(
2100 ops in prop::collection::vec(
2101 topic_strategy().prop_flat_map(|t| {
2102 prop_oneof![
2103 Just(Operation::AddReference(t.clone())),
2104 Just(Operation::RemoveReference(t)),
2105 ]
2106 }),
2107 1..100
2108 )
2109 ) {
2110 let state = SubscriptionState::new('.');
2111 let mut expected = AHashMap::new();
2112
2113 for op in &ops {
2114 match op {
2115 Operation::AddReference(topic) => {
2116 let count = expected.entry(topic.clone()).or_insert(0usize);
2117 let should_subscribe = *count == 0;
2118 *count += 1;
2119 prop_assert_eq!(state.add_reference(topic), should_subscribe);
2120 }
2121 Operation::RemoveReference(topic) => {
2122 let count = expected.get(topic).copied().unwrap_or(0);
2123 let should_unsubscribe = count == 1;
2124 if should_unsubscribe {
2125 expected.remove(topic);
2126 } else if count > 1 {
2127 *expected.get_mut(topic).unwrap() -= 1;
2128 }
2129 prop_assert_eq!(state.remove_reference(topic), should_unsubscribe);
2130 }
2131 _ => unreachable!("reference-count strategy only generates reference operations"),
2132 }
2133
2134 prop_assert_eq!(state.reference_counts.len(), expected.len());
2135 for (topic, count) in &expected {
2136 prop_assert_eq!(state.get_reference_count(topic), *count);
2137 }
2138 }
2139 }
2140
2141 #[rstest]
2143 fn prop_all_topics_is_union(
2144 operations in prop::collection::vec(operation_strategy(), 1..50)
2145 ) {
2146 let state = SubscriptionState::new('.');
2147
2148 for op in &operations {
2149 apply_operation(&state, op);
2150
2151 let all_topics: AHashSet<String> = state.all_topics().into_iter().collect();
2153 let confirmed: AHashSet<String> = state.topics_from_map(&state.confirmed).into_iter().collect();
2154 let pending_sub: AHashSet<String> = state.pending_subscribe_topics().into_iter().collect();
2155 let expected: AHashSet<String> = confirmed.union(&pending_sub).cloned().collect();
2156
2157 assert_eq!(all_topics, expected);
2158
2159 let pending_unsub: AHashSet<String> = state.pending_unsubscribe_topics().into_iter().collect();
2161 for topic in pending_unsub {
2162 assert!(!all_topics.contains(&topic));
2163 }
2164 }
2165 }
2166
2167 #[rstest]
2169 fn prop_clear_resets_completely(
2170 operations in prop::collection::vec(operation_strategy(), 1..30)
2171 ) {
2172 let state = SubscriptionState::new('.');
2173
2174 for op in &operations {
2176 apply_operation(&state, op);
2177 }
2178
2179 state.clear();
2181
2182 assert!(state.is_empty());
2183 assert_eq!(state.len(), 0);
2184 assert!(state.all_topics().is_empty());
2185 assert!(state.pending_subscribe_topics().is_empty());
2186 assert!(state.pending_unsubscribe_topics().is_empty());
2187 assert!(state.confirmed.is_empty());
2188 assert!(state.pending_subscribe.is_empty());
2189 assert!(state.pending_unsubscribe.is_empty());
2190 assert!(state.reference_counts.is_empty());
2191 }
2192
2193 #[rstest]
2195 fn prop_topic_mutual_exclusivity(
2196 operations in prop::collection::vec(operation_strategy(), 1..50),
2197 topic in topic_strategy()
2198 ) {
2199 let state = SubscriptionState::new('.');
2200
2201 for (i, op) in operations.iter().enumerate() {
2202 apply_operation(&state, op);
2203 check_topic_exclusivity(&state, &topic, &format!("After op {i}: {op:?}"));
2204 }
2205 }
2206 }
2207 }
2208}