DevFlow/crates/df-relay/src/conn.rs

//! df-relay 连接注册表
//!
//! 设计依据:设计文档「Layer2」—— 按 device_id 配对路由(非全局广播)。
//!
//! 注册表结构:`HashMap<conn_id, ConnHandle>` + `device_id → conn_id` 索引。
//! 配对语义:小程序 device_id ↔ 桌面端绑定该 device_id 的 token(同 device_id
//! 的小程序与桌面端互为对端)。relay 按 device_id 找到对端连接转发消息。
//!
//! 并发:tokio::sync::Mutex 单锁(注册表操作高频但临界区小,先简单后优化)。

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use tokio::sync::{mpsc, Mutex};

use crate::broadcast::{BroadcastMessage, ClientKind};

/// 连接唯一标识(u64 自增,进程内唯一)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct ConnId(pub u64);

impl ConnId {
    /// 伪值(用于无来源场景,如服务端构造的控制消息)
    pub const NIL: ConnId = ConnId(0);
}

static NEXT_CONN_ID: AtomicU64 = AtomicU64::new(1);

/// 分配下一个 conn_id(进程内递增,从 1 起,0 留作 NIL)
pub fn next_conn_id() -> ConnId {
    ConnId(NEXT_CONN_ID.fetch_add(1, Ordering::Relaxed))
}

/// 连接句柄(注册表条目)
///
/// 每个连接持一个 mpsc::UnboundedSender,广播时向其投递 BroadcastMessage,
/// 由连接自己的读取任务转发给对端 socket。
#[derive(Debug)]
pub struct ConnHandle {
    /// 连接 ID
    pub id: ConnId,
    /// 客户端类型
    pub kind: ClientKind,
    /// 配对绑定的设备 ID(桌面端首次配置生成,小程序携带)
    pub device_id: String,
    /// 广播投递通道(连接读取任务消费 → 写入 socket)
    pub sender: mpsc::UnboundedSender<BroadcastMessage>,
}

impl ConnHandle {
    pub fn new(
        id: ConnId,
        kind: ClientKind,
        device_id: String,
        sender: mpsc::UnboundedSender<BroadcastMessage>,
    ) -> Self {
        Self {
            id,
            kind,
            device_id,
            sender,
        }
    }
}

/// 连接注册表(RelayState 内部状态)
#[derive(Default)]
pub struct ConnRegistry {
    /// conn_id → 连接句柄
    by_id: HashMap<ConnId, ConnHandle>,
    /// device_id → 绑定的连接集合(同 device_id 的小程序 + 桌面端)
    /// 设计为 Vec:兼容同 device_id 多连接(多端登录),MVP 一般 1+1。
    by_device: HashMap<String, Vec<ConnId>>,
}

impl ConnRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    /// 注册连接
    pub fn add(&mut self, handle: ConnHandle) {
        let id = handle.id;
        let device_id = handle.device_id.clone();
        self.by_device.entry(device_id).or_default().push(id);
        self.by_id.insert(id, handle);
    }

    /// 注销连接(返回其 device_id 用于清理索引)
    pub fn remove(&mut self, id: ConnId) -> Option<String> {
        let handle = self.by_id.remove(&id)?;
        let device_id = handle.device_id.clone();
        if let Some(ids) = self.by_device.get_mut(&device_id) {
            ids.retain(|c| *c != id);
            if ids.is_empty() {
                self.by_device.remove(&device_id);
            }
        }
        Some(device_id)
    }

    /// 查询 device_id 是否有在线连接(离线降级判断用)
    pub fn is_device_online(&self, device_id: &str) -> bool {
        self.by_device
            .get(device_id)
            .map(|ids| !ids.is_empty())
            .unwrap_or(false)
    }

    /// 投递消息给指定 device_id 配对的对端连接(排除来源 source)
    ///
    /// 路由规则:
    /// - Event(Device 发):转发给同 device_id 的所有 Miniapp 连接
    /// - Command(Miniapp 发):转发给同 device_id 的所有 Device 连接
    /// - Control:转发给同 device_id 的所有连接(除来源)
    ///
    /// 返回成功投递的连接数(用于日志/降级判断)。
    pub fn route(&self, msg: &BroadcastMessage) -> usize {
        let Some(ids) = self.by_device.get(&msg.device_id) else {
            return 0;
        };
        let mut delivered = 0usize;
        for &id in ids {
            if id == msg.source {
                // 不回环给来源
                continue;
            }
            let Some(handle) = self.by_id.get(&id) else {
                continue;
            };
            match (msg.kind, msg.from, handle.kind) {
                // Event:仅投递给 Miniapp
                (crate::broadcast::MessageKind::Event, _, ClientKind::Miniapp) => {}
                // Command:仅投递给 Device
                (crate::broadcast::MessageKind::Command, _, ClientKind::Device) => {}
                // Control:全投递(除来源)
                (crate::broadcast::MessageKind::Control, _, _) => {}
                // 其余方向不匹配(Event 不投递给 Device,Command 不投递给 Miniapp)
                _ => continue,
            }
            if handle.sender.send(msg.clone()).is_ok() {
                delivered += 1;
            }
        }
        delivered
    }

    /// 当前总连接数(诊断用)
    pub fn len(&self) -> usize {
        self.by_id.len()
    }

    /// 是否为空(诊断用)
    pub fn is_empty(&self) -> bool {
        self.by_id.is_empty()
    }
}

/// 共享中继状态(axum State 传递)
///
/// 注册表用 tokio Mutex 包裹:广播路径在锁内调用 route,锁粒度小。
#[derive(Clone, Default)]
pub struct RelayState {
    inner: Arc<Mutex<ConnRegistry>>,
}

impl RelayState {
    pub fn new() -> Self {
        Self::default()
    }

    /// 从 Arc 直接构造(用于 axum state)
    pub fn from_arc(inner: Arc<Mutex<ConnRegistry>>) -> Self {
        Self { inner }
    }

    /// 暴露内部 Arc(测试/外部 start 使用)
    pub fn inner(&self) -> Arc<Mutex<ConnRegistry>> {
        self.inner.clone()
    }

    /// 注册连接
    pub async fn add_conn(&self, handle: ConnHandle) {
        self.inner.lock().await.add(handle);
    }

    /// 注销连接
    pub async fn remove_conn(&self, id: ConnId) -> Option<String> {
        self.inner.lock().await.remove(id)
    }

    /// 查询 device 在线状态
    pub async fn is_device_online(&self, device_id: &str) -> bool {
        self.inner.lock().await.is_device_online(device_id)
    }

    /// 投递消息给对端连接(返回成功投递数)
    pub async fn route(&self, msg: &BroadcastMessage) -> usize {
        self.inner.lock().await.route(msg)
    }
}