futu_server/

subscription.rs

// 订阅管理：行情订阅 + 交易账户推送订阅 + 通知订阅

use std::collections::{HashMap, HashSet};

use parking_lot::RwLock;

/// 订阅管理器
pub struct SubscriptionManager {
    /// 通知订阅：哪些连接订阅了系统通知
    notify_subs: RwLock<HashSet<u64>>,

    /// 交易账户推送订阅：acc_id → Set<conn_id>
    trd_acc_subs: RwLock<HashMap<u64, HashSet<u64>>>,

    /// 行情订阅："market_code:sub_type" → Set<conn_id>
    qot_subs: RwLock<HashMap<String, HashSet<u64>>>,

    /// 每连接的行情订阅额度使用
    conn_quota: RwLock<HashMap<u64, u32>>,
}

/// 总订阅额度上限
pub const TOTAL_QUOTA: u32 = 4000;

impl SubscriptionManager {
    pub fn new() -> Self {
        Self {
            notify_subs: RwLock::new(HashSet::new()),
            trd_acc_subs: RwLock::new(HashMap::new()),
            qot_subs: RwLock::new(HashMap::new()),
            conn_quota: RwLock::new(HashMap::new()),
        }
    }

    // ===== 通知订阅 =====

    pub fn subscribe_notify(&self, conn_id: u64) {
        self.notify_subs.write().insert(conn_id);
    }

    pub fn unsubscribe_notify(&self, conn_id: u64) {
        self.notify_subs.write().remove(&conn_id);
    }

    pub fn is_subscribed_notify(&self, conn_id: u64) -> bool {
        self.notify_subs.read().contains(&conn_id)
    }

    // ===== 交易账户推送 =====

    pub fn subscribe_trd_acc(&self, conn_id: u64, acc_id: u64) {
        self.trd_acc_subs
            .write()
            .entry(acc_id)
            .or_default()
            .insert(conn_id);
    }

    pub fn unsubscribe_trd_acc(&self, conn_id: u64, acc_id: u64) {
        if let Some(subs) = self.trd_acc_subs.write().get_mut(&acc_id) {
            subs.remove(&conn_id);
        }
    }

    pub fn get_acc_subscribers(&self, acc_id: u64) -> Vec<u64> {
        self.trd_acc_subs
            .read()
            .get(&acc_id)
            .map(|s| s.iter().copied().collect())
            .unwrap_or_default()
    }

    // ===== 行情订阅 =====

    /// 生成行情订阅 key
    pub fn make_qot_key(market: i32, code: &str, sub_type: i32) -> String {
        format!("{market}_{code}:{sub_type}")
    }

    /// 订阅行情
    pub fn subscribe_qot(&self, conn_id: u64, security_key: &str, sub_type: i32) {
        let key = format!("{security_key}:{sub_type}");
        self.qot_subs
            .write()
            .entry(key)
            .or_default()
            .insert(conn_id);

        // 更新额度
        *self.conn_quota.write().entry(conn_id).or_insert(0) += 1;
    }

    /// 退订行情
    pub fn unsubscribe_qot(&self, conn_id: u64, security_key: &str, sub_type: i32) {
        let key = format!("{security_key}:{sub_type}");
        if let Some(subs) = self.qot_subs.write().get_mut(&key) {
            if subs.remove(&conn_id) {
                let mut quota = self.conn_quota.write();
                if let Some(q) = quota.get_mut(&conn_id) {
                    *q = q.saturating_sub(1);
                }
            }
        }
    }

    /// 获取订阅了指定行情的连接列表
    pub fn get_qot_subscribers(&self, security_key: &str, sub_type: i32) -> Vec<u64> {
        let key = format!("{security_key}:{sub_type}");
        self.qot_subs
            .read()
            .get(&key)
            .map(|s| s.iter().copied().collect())
            .unwrap_or_default()
    }

    /// 获取连接的已用订阅额度
    pub fn get_conn_used_quota(&self, conn_id: u64) -> u32 {
        self.conn_quota.read().get(&conn_id).copied().unwrap_or(0)
    }

    /// 获取总已用额度
    pub fn get_total_used_quota(&self) -> u32 {
        self.conn_quota.read().values().sum()
    }

    /// 获取指定连接订阅的行情列表: sub_type → Vec<security_key>
    pub fn get_conn_qot_subs(&self, conn_id: u64) -> HashMap<i32, Vec<String>> {
        let qot = self.qot_subs.read();
        let mut result: HashMap<i32, Vec<String>> = HashMap::new();
        for (key, conn_ids) in qot.iter() {
            if conn_ids.contains(&conn_id) {
                // key 格式: "market_code:sub_type"
                if let Some(colon) = key.rfind(':') {
                    if let Ok(sub_type) = key[colon + 1..].parse::<i32>() {
                        let sec_key = &key[..colon];
                        result
                            .entry(sub_type)
                            .or_default()
                            .push(sec_key.to_string());
                    }
                }
            }
        }
        result
    }

    /// 获取所有连接 ID（有行情订阅的）
    pub fn get_all_qot_conn_ids(&self) -> HashSet<u64> {
        let qot = self.qot_subs.read();
        let mut ids = HashSet::new();
        for conn_ids in qot.values() {
            ids.extend(conn_ids);
        }
        ids
    }

    // ===== 连接断开清理 =====

    pub fn on_disconnect(&self, conn_id: u64) {
        self.notify_subs.write().remove(&conn_id);

        // 清理交易账户订阅
        let mut trd = self.trd_acc_subs.write();
        for subs in trd.values_mut() {
            subs.remove(&conn_id);
        }

        // 清理行情订阅
        let mut qot = self.qot_subs.write();
        for subs in qot.values_mut() {
            subs.remove(&conn_id);
        }

        self.conn_quota.write().remove(&conn_id);
    }
}

impl Default for SubscriptionManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_notify_subscription() {
        let mgr = SubscriptionManager::new();
        assert!(!mgr.is_subscribed_notify(1));
        mgr.subscribe_notify(1);
        assert!(mgr.is_subscribed_notify(1));
        mgr.unsubscribe_notify(1);
        assert!(!mgr.is_subscribed_notify(1));
    }

    #[test]
    fn test_trd_acc_subscription() {
        let mgr = SubscriptionManager::new();
        mgr.subscribe_trd_acc(1, 100);
        mgr.subscribe_trd_acc(2, 100);
        mgr.subscribe_trd_acc(1, 200);

        let subs = mgr.get_acc_subscribers(100);
        assert_eq!(subs.len(), 2);
        assert!(subs.contains(&1));
        assert!(subs.contains(&2));

        mgr.unsubscribe_trd_acc(1, 100);
        assert_eq!(mgr.get_acc_subscribers(100).len(), 1);
    }

    #[test]
    fn test_qot_subscription() {
        let mgr = SubscriptionManager::new();
        mgr.subscribe_qot(1, "1_00700", 1); // Basic
        mgr.subscribe_qot(2, "1_00700", 1);
        mgr.subscribe_qot(1, "1_00700", 2); // OrderBook

        assert_eq!(mgr.get_qot_subscribers("1_00700", 1).len(), 2);
        assert_eq!(mgr.get_qot_subscribers("1_00700", 2).len(), 1);
        assert_eq!(mgr.get_conn_used_quota(1), 2);
    }

    #[test]
    fn test_disconnect_cleanup() {
        let mgr = SubscriptionManager::new();
        mgr.subscribe_notify(1);
        mgr.subscribe_trd_acc(1, 100);
        mgr.subscribe_qot(1, "1_00700", 1);

        mgr.on_disconnect(1);

        assert!(!mgr.is_subscribed_notify(1));
        assert!(mgr.get_acc_subscribers(100).is_empty());
        assert!(mgr.get_qot_subscribers("1_00700", 1).is_empty());
        assert_eq!(mgr.get_conn_used_quota(1), 0);
    }
}