1mod connection_lifecycle;
29mod disconnected_cleanup;
30mod push_regs;
31mod qot_commit;
32mod session_detail;
33mod unsubscribe_all_commit;
34mod views;
35
36use std::collections::{HashMap, HashSet};
37use std::sync::Arc;
38use std::sync::atomic::{AtomicU64, Ordering};
39use std::time::{Duration, Instant};
40
41use dashmap::DashMap;
42use futu_core::qot_stock_key::QotSecurityKey;
43pub use futu_domain_qot_subscription::QOT_MIN_UNSUB_ELAPSED_SECS;
44use futu_domain_qot_subscription::{
45 CryptoSubscriptionProbe, UnsubscribeAllGlobalEmptyProbe,
46 is_crypto_stock_broker_globally_unsubscribed, is_crypto_stock_globally_unsubscribed,
47 plan_unsubscribe_all_global_empty_keys, qot_min_unsub_freshness_from_elapsed_secs,
48};
49use parking_lot::RwLock;
50
51use crate::conn::ClientCloseControl;
52
53pub type ConnectionDisconnectObserver = Arc<dyn Fn(u64) + Send + Sync>;
54pub type ConnectionOpenObserver = Arc<dyn Fn(u64, u64) + Send + Sync>;
55
56pub struct SubscriptionManager {
58 connection_open_observers: RwLock<Vec<ConnectionOpenObserver>>,
59 disconnect_observers: RwLock<Vec<ConnectionDisconnectObserver>>,
63
64 client_close_controls: DashMap<u64, ClientCloseControl>,
70 connection_generations: DashMap<u64, u64>,
72
73 notify_subs: RwLock<HashSet<u64>>,
75
76 trd_acc_subs: RwLock<HashMap<u64, HashSet<u64>>>,
78
79 api_page_req_keys: RwLock<HashMap<u64, HashSet<[u8; 16]>>>,
82
83 qot_subs: RwLock<HashMap<(QotSecurityKey, i32), HashSet<u64>>>,
86
87 qot_push_regs: RwLock<QotPushRegistrations>,
91
92 qot_sub_sessions: RwLock<QotSessionState>,
97
98 qot_orderbook_detail: RwLock<HashMap<QotSecurityKey, HashMap<u64, bool>>>,
101
102 qot_broker_detail: RwLock<HashMap<QotSecurityKey, HashMap<u64, bool>>>,
105
106 total_quota: RwLock<u32>,
110
111 qot_sub_times: RwLock<HashMap<(QotSecurityKey, i32), Instant>>,
115
116 qot_disconnected_conns: RwLock<HashSet<u64>>,
122
123 qot_disconnect_sync_generation: AtomicU64,
129 qot_owner_token_high_water: AtomicU64,
130 qot_owner_tokens: RwLock<HashMap<(QotSecurityKey, i32, u64), u64>>,
131}
132
133#[derive(Default)]
134struct QotPushRegistrations {
135 by_tuple: HashMap<(QotSecurityKey, i32, i32), HashSet<u64>>,
136 qot_push_regs_by_cache_key: HashMap<String, HashMap<(i32, i32), HashSet<u64>>>,
137}
138
139#[derive(Default)]
140struct QotSessionState {
141 by_key: HashMap<(QotSecurityKey, i32), HashMap<u64, i32>>,
142 by_cache_key: HashMap<(String, i32), HashMap<u64, i32>>,
145}
146
147pub const TOTAL_QUOTA: u32 = 4000;
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
161pub enum SubResult {
162 NewGlobal,
164 AlreadyGlobal,
166 NoChange,
168}
169
170#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172pub enum UnsubResult {
173 LastSubscriber,
176 StillSubscribed,
178 NotSubscribed,
181}
182
183impl SubscriptionManager {
184 pub fn new() -> Self {
185 Self {
186 connection_open_observers: RwLock::new(Vec::new()),
187 disconnect_observers: RwLock::new(Vec::new()),
188 client_close_controls: DashMap::new(),
189 connection_generations: DashMap::new(),
190 notify_subs: RwLock::new(HashSet::new()),
191 trd_acc_subs: RwLock::new(HashMap::new()),
192 api_page_req_keys: RwLock::new(HashMap::new()),
193 qot_subs: RwLock::new(HashMap::new()),
194 qot_push_regs: RwLock::new(QotPushRegistrations::default()),
195 qot_sub_sessions: RwLock::new(QotSessionState::default()),
196 qot_orderbook_detail: RwLock::new(HashMap::new()),
197 qot_broker_detail: RwLock::new(HashMap::new()),
198 total_quota: RwLock::new(TOTAL_QUOTA),
199 qot_sub_times: RwLock::new(HashMap::new()),
200 qot_disconnected_conns: RwLock::new(HashSet::new()),
201 qot_disconnect_sync_generation: AtomicU64::new(0),
202 qot_owner_token_high_water: AtomicU64::new(0),
203 qot_owner_tokens: RwLock::new(HashMap::new()),
204 }
205 }
206
207 pub fn subscribe_notify(&self, conn_id: u64) {
210 self.notify_subs.write().insert(conn_id);
211 }
212
213 pub fn unsubscribe_notify(&self, conn_id: u64) {
214 self.notify_subs.write().remove(&conn_id);
215 }
216
217 pub fn is_subscribed_notify(&self, conn_id: u64) -> bool {
218 self.notify_subs.read().contains(&conn_id)
219 }
220
221 pub fn register_api_page_req_key(&self, conn_id: u64, key: &[u8]) -> bool {
222 let Ok(key) = <[u8; 16]>::try_from(key) else {
223 return false;
224 };
225 self.api_page_req_keys
226 .write()
227 .entry(conn_id)
228 .or_default()
229 .insert(key);
230 true
231 }
232
233 #[must_use]
234 pub fn is_api_page_req_key_registered(&self, conn_id: u64, key: &[u8]) -> bool {
235 let Ok(key) = <[u8; 16]>::try_from(key) else {
236 return false;
237 };
238 self.api_page_req_keys
239 .read()
240 .get(&conn_id)
241 .is_some_and(|keys| keys.contains(&key))
242 }
243
244 pub fn subscribe_trd_acc(&self, conn_id: u64, acc_id: u64) {
247 self.trd_acc_subs
248 .write()
249 .entry(acc_id)
250 .or_default()
251 .insert(conn_id);
252 }
253
254 pub fn unsubscribe_trd_acc(&self, conn_id: u64, acc_id: u64) {
255 if let Some(subs) = self.trd_acc_subs.write().get_mut(&acc_id) {
256 subs.remove(&conn_id);
257 }
258 }
259
260 pub fn get_acc_subscribers(&self, acc_id: u64) -> Vec<u64> {
261 match self.trd_acc_subs.read().get(&acc_id) {
262 Some(subscribers) => subscribers.iter().copied().collect(),
263 None => Vec::new(),
264 }
265 }
266
267 pub fn make_qot_key(market: i32, code: &str, sub_type: i32) -> String {
271 format!("{market}_{code}:{sub_type}")
272 }
273
274 #[inline]
275 fn broker_key(sec_key: &QotSecurityKey) -> QotSecurityKey {
276 sec_key.clone()
277 }
278
279 pub fn subscribe_qot_broker(
285 &self,
286 conn_id: u64,
287 sec_key: &QotSecurityKey,
288 sub_type: i32,
289 ) -> SubResult {
290 qot_commit::subscribe_broker(self, conn_id, Self::broker_key(sec_key), sub_type)
291 }
292
293 pub fn unsubscribe_qot_broker(
296 &self,
297 conn_id: u64,
298 sec_key: &QotSecurityKey,
299 sub_type: i32,
300 ) -> UnsubResult {
301 qot_commit::unsubscribe_broker(self, conn_id, Self::broker_key(sec_key), sub_type)
302 }
303
304 pub fn is_qot_subscribed_broker(
306 &self,
307 conn_id: u64,
308 sec_key: &QotSecurityKey,
309 sub_type: i32,
310 ) -> bool {
311 self.qot_subs
312 .read()
313 .get(&(Self::broker_key(sec_key), sub_type))
314 .is_some_and(|subs| subs.contains(&conn_id))
315 }
316
317 pub fn is_globally_subscribed_broker(&self, sec_key: &QotSecurityKey, sub_type: i32) -> bool {
320 self.qot_subs
321 .read()
322 .get(&(Self::broker_key(sec_key), sub_type))
323 .is_some_and(|subs| !subs.is_empty())
324 }
325
326 pub fn qot_min_unsub_elapsed_broker(&self, sec_key: &QotSecurityKey, sub_type: i32) -> bool {
328 self.qot_sub_times
329 .read()
330 .get(&(Self::broker_key(sec_key), sub_type))
331 .map(|instant| {
332 qot_min_unsub_freshness_from_elapsed_secs(instant.elapsed().as_secs()).min_elapsed
333 })
334 .unwrap_or(true)
335 }
336
337 pub fn qot_min_unsub_remaining_secs_broker(
339 &self,
340 sec_key: &QotSecurityKey,
341 sub_type: i32,
342 ) -> u64 {
343 self.qot_sub_times
344 .read()
345 .get(&(Self::broker_key(sec_key), sub_type))
346 .map(|instant| {
347 qot_min_unsub_freshness_from_elapsed_secs(instant.elapsed().as_secs())
348 .remaining_secs
349 })
350 .unwrap_or(0)
351 }
352
353 pub fn qot_disconnect_sync_generation(&self) -> u64 {
355 self.qot_disconnect_sync_generation.load(Ordering::SeqCst)
356 }
357
358 #[doc(hidden)]
359 pub fn backdate_qot_sub_time_broker_for_test(
360 &self,
361 sec_key: &QotSecurityKey,
362 sub_type: i32,
363 elapsed: Duration,
364 ) {
365 let map_key = (Self::broker_key(sec_key), sub_type);
366 let instant = Instant::now()
367 .checked_sub(elapsed)
368 .unwrap_or_else(Instant::now);
369 self.qot_sub_times.write().insert(map_key, instant);
370 }
371
372 pub fn unsubscribe_all_qot_collect_global_empty(&self, conn_id: u64) -> Vec<(String, i32)> {
377 unsubscribe_all_commit::collect_global_empty(self, conn_id)
378 }
379
380 pub fn cleanup_due_disconnected_qot(&self) -> Vec<(String, i32)> {
388 disconnected_cleanup::cleanup_due(self)
389 }
390
391 pub fn unsubscribe_all_qot_dry_run(&self, conn_id: u64) -> Vec<(String, i32)> {
403 let qot = self.qot_subs.read();
404 let probes = qot
405 .iter()
406 .map(|((key, sub_type), set)| UnsubscribeAllGlobalEmptyProbe {
407 key: key.cache_key(),
408 sub_type: *sub_type,
409 conn_is_subscribed: set.contains(&conn_id),
410 subscriber_count: set.len(),
411 });
412 plan_unsubscribe_all_global_empty_keys(probes)
413 }
414
415 pub fn unsubscribe_all_qot_commit(&self, conn_id: u64) -> Vec<(String, i32)> {
419 self.unsubscribe_all_qot_collect_global_empty(conn_id)
420 }
421
422 pub fn get_qot_subscribers_broker(&self, sec_key: &QotSecurityKey, sub_type: i32) -> Vec<u64> {
426 match self
427 .qot_subs
428 .read()
429 .get(&(Self::broker_key(sec_key), sub_type))
430 {
431 Some(subscribers) => subscribers.iter().copied().collect(),
432 None => Vec::new(),
433 }
434 }
435
436 pub fn qot_owner_lease_broker(
438 &self,
439 sec_key: &QotSecurityKey,
440 sub_type: i32,
441 ) -> Vec<(u64, u64, i32, u64)> {
442 let disconnected = self.qot_disconnected_conns.read();
443 let mut owners = self
444 .get_qot_subscribers_broker(sec_key, sub_type)
445 .into_iter()
446 .filter(|conn_id| !disconnected.contains(conn_id))
447 .filter_map(|conn_id| {
448 let connection_generation = self
449 .connection_generations
450 .get(&conn_id)
451 .map(|generation| *generation)
452 .unwrap_or(0);
453 let owner_token = self
454 .qot_owner_tokens
455 .read()
456 .get(&(Self::broker_key(sec_key), sub_type, conn_id))
457 .copied()?;
458 Some((
459 conn_id,
460 connection_generation,
461 self.get_conn_session_broker(conn_id, sec_key, sub_type),
462 owner_token,
463 ))
464 })
465 .collect::<Vec<_>>();
466 owners.sort_unstable();
467 owners
468 }
469
470 pub fn qot_owner_lease_is_current(
471 &self,
472 sec_key: &QotSecurityKey,
473 sub_type: i32,
474 captured: &[(u64, u64, i32, u64)],
475 request_section: Option<i32>,
476 ) -> bool {
477 self.with_current_qot_owner_lease(sec_key, sub_type, captured, request_section, || ())
478 .is_some()
479 }
480
481 pub fn with_current_qot_owner_lease<R>(
482 &self,
483 sec_key: &QotSecurityKey,
484 sub_type: i32,
485 captured: &[(u64, u64, i32, u64)],
486 request_section: Option<i32>,
487 publish: impl FnOnce() -> R,
488 ) -> Option<R> {
489 let key = Self::broker_key(sec_key);
490 let qot = self.qot_subs.read();
491 let subscribers = qot.get(&(key.clone(), sub_type))?;
492 let disconnected = self.qot_disconnected_conns.read();
493 let tokens = self.qot_owner_tokens.read();
494 let sessions = self.qot_sub_sessions.read();
495 let session_map = sessions.by_key.get(&(key.clone(), sub_type));
496 let current = captured.iter().any(
497 |(conn_id, connection_generation, captured_session, owner_token)| {
498 if !subscribers.contains(conn_id) || disconnected.contains(conn_id) {
499 return false;
500 }
501 let current_connection_generation = self
502 .connection_generations
503 .get(conn_id)
504 .map(|generation| *generation)
505 .unwrap_or(0);
506 let same_connection = current_connection_generation == *connection_generation;
507 let same_owner = tokens
508 .get(&(key.clone(), sub_type, *conn_id))
509 .is_some_and(|token| *token == *owner_token);
510 let session = session_map
511 .and_then(|map| map.get(conn_id))
512 .copied()
513 .unwrap_or(1);
514 same_connection
515 && same_owner
516 && session == *captured_session
517 && request_section.is_none_or(|section| match section {
518 2 | 3 => matches!(session, 2 | 3),
519 5 => session == 3,
520 _ => matches!(session, 0..=3),
521 })
522 },
523 );
524 current.then(publish)
525 }
526
527 fn next_qot_owner_token(&self) -> u64 {
528 let mut current = self.qot_owner_token_high_water.load(Ordering::SeqCst);
529 loop {
530 let next = if current == u64::MAX { 1 } else { current + 1 };
531 match self.qot_owner_token_high_water.compare_exchange(
532 current,
533 next,
534 Ordering::SeqCst,
535 Ordering::SeqCst,
536 ) {
537 Ok(_) => return next,
538 Err(observed) => current = observed,
539 }
540 }
541 }
542
543 pub(super) fn assign_qot_owner_token(&self, key: &QotSecurityKey, sub_type: i32, conn_id: u64) {
544 let token = self.next_qot_owner_token();
545 self.qot_owner_tokens
546 .write()
547 .insert((Self::broker_key(key), sub_type, conn_id), token);
548 }
549
550 pub(super) fn remove_qot_owner_token(&self, key: &QotSecurityKey, sub_type: i32, conn_id: u64) {
551 self.qot_owner_tokens
552 .write()
553 .remove(&(Self::broker_key(key), sub_type, conn_id));
554 }
555
556 pub(super) fn remove_all_qot_owner_tokens(&self, conn_id: u64) {
557 self.qot_owner_tokens
558 .write()
559 .retain(|(_, _, owner), _| *owner != conn_id);
560 }
561
562 pub fn crypto_stock_globally_unsubscribed(&self, stock_id: u64) -> bool {
572 let qot = self.qot_subs.read();
573 let probes = qot.iter().map(|((key, _sub_type), subs)| {
574 CryptoSubscriptionProbe::from_runtime_facts(
575 key.stock_key.stock_id,
576 key.stock_key.broker_id,
577 subs.len(),
578 )
579 });
580 is_crypto_stock_globally_unsubscribed(stock_id, probes)
581 }
582
583 pub fn crypto_stock_broker_globally_unsubscribed(&self, stock_id: u64, broker_id: u32) -> bool {
592 let target_broker = std::num::NonZeroU32::new(broker_id);
593 let qot = self.qot_subs.read();
594 let probes = qot.iter().map(|((key, _sub_type), subs)| {
595 CryptoSubscriptionProbe::from_runtime_facts(
596 key.stock_key.stock_id,
597 key.stock_key.broker_id,
598 subs.len(),
599 )
600 });
601 is_crypto_stock_broker_globally_unsubscribed(stock_id, target_broker, probes)
602 }
603
604 pub fn set_conn_session_broker(
607 &self,
608 conn_id: u64,
609 sec_key: &QotSecurityKey,
610 sub_type: i32,
611 session: i32,
612 ) {
613 let previous = self.get_conn_session_broker(conn_id, sec_key, sub_type);
614 session_detail::set_conn_session(
615 self,
616 conn_id,
617 Self::broker_key(sec_key),
618 sub_type,
619 session,
620 );
621 if previous != session && self.is_qot_subscribed_broker(conn_id, sec_key, sub_type) {
622 self.assign_qot_owner_token(sec_key, sub_type, conn_id);
623 }
624 }
625
626 pub fn get_global_session_broker(&self, sec_key: &QotSecurityKey, sub_type: i32) -> i32 {
627 session_detail::global_session(self, Self::broker_key(sec_key), sub_type)
628 }
629
630 pub fn get_conn_session_broker(
632 &self,
633 conn_id: u64,
634 sec_key: &QotSecurityKey,
635 sub_type: i32,
636 ) -> i32 {
637 session_detail::conn_session(self, conn_id, Self::broker_key(sec_key), sub_type)
638 }
639
640 pub fn get_conn_session_by_cache_key(
643 &self,
644 conn_id: u64,
645 cache_key: &str,
646 sub_type: i32,
647 ) -> i32 {
648 self.qot_sub_sessions
649 .read()
650 .by_cache_key
651 .get(&(cache_key.to_owned(), sub_type))
652 .and_then(|sessions| sessions.get(&conn_id))
653 .copied()
654 .unwrap_or(1)
655 }
656
657 pub fn set_conn_orderbook_detail_broker(
658 &self,
659 conn_id: u64,
660 sec_key: &QotSecurityKey,
661 detail: bool,
662 ) {
663 session_detail::set_conn_orderbook_detail(self, conn_id, Self::broker_key(sec_key), detail);
664 }
665
666 pub fn is_global_orderbook_detail_broker(&self, sec_key: &QotSecurityKey) -> bool {
667 session_detail::global_orderbook_detail(self, Self::broker_key(sec_key))
668 }
669
670 pub fn set_conn_broker_detail_broker(
671 &self,
672 conn_id: u64,
673 sec_key: &QotSecurityKey,
674 detail: bool,
675 ) {
676 session_detail::set_conn_broker_detail(self, conn_id, Self::broker_key(sec_key), detail);
677 }
678
679 pub fn is_global_broker_detail_broker(&self, sec_key: &QotSecurityKey) -> bool {
680 session_detail::global_broker_detail(self, Self::broker_key(sec_key))
681 }
682
683 pub fn register_connection_open_observer(&self, observer: ConnectionOpenObserver) {
686 self.connection_open_observers.write().push(observer);
687 }
688
689 pub(crate) fn on_connect(&self, conn_id: u64, session_generation: u64) {
690 self.connection_generations
691 .insert(conn_id, session_generation);
692 let observers = self.connection_open_observers.read().clone();
693 for observer in observers {
694 observer(conn_id, session_generation);
695 }
696 }
697
698 pub fn register_disconnect_observer(&self, observer: ConnectionDisconnectObserver) {
699 self.disconnect_observers.write().push(observer);
700 }
701
702 pub(crate) fn register_client_close_control(
703 &self,
704 conn_id: u64,
705 close_control: ClientCloseControl,
706 ) {
707 self.client_close_controls.insert(conn_id, close_control);
708 }
709
710 pub(crate) fn request_client_close(&self, conn_id: u64) -> Option<bool> {
714 self.client_close_controls
715 .get(&conn_id)
716 .map(|control| control.request_close())
717 }
718
719 pub(crate) fn remove_client_close_control(&self, conn_id: u64) {
720 self.client_close_controls.remove(&conn_id);
721 }
722
723 pub fn on_disconnect(&self, conn_id: u64) -> Vec<(String, i32)> {
724 connection_lifecycle::on_disconnect(self, conn_id)
725 }
726}
727
728impl Default for SubscriptionManager {
729 fn default() -> Self {
730 Self::new()
731 }
732}
733
734#[inline]
735fn sub_type_orderbook() -> i32 {
736 2
737}
738
739#[inline]
740fn sub_type_broker() -> i32 {
741 14
742}
743
744#[cfg(test)]
745mod tests;