pi_quic 0.10.1

//! QUIC协议库；指标增量仅在池启动、UDP冷绑定及首次客户端定时器前登记原Local。
//! 身份/去重/强保活归pi_async_rt，采集归宿主；不增加ID字段、注销或热路径操作。
//! NET-MET-BATCH-DESIGN-001/B-QUIC；纯登记专项见batch_runtime_metrics_tests。

use std::ptr;
use std::sync::Arc;
use std::task::Waker;
use std::time::Duration;
use std::net::SocketAddr;
use std::cell::UnsafeCell;
use std::result::Result as GenResult;
use std::io::{Error, Result, ErrorKind};
use std::sync::atomic::{AtomicU8, Ordering};

use futures::future::LocalBoxFuture;
use quinn_proto::{ConnectionHandle, ConnectionEvent, Dir, StreamId, Transmit};
use crossbeam_channel::Sender;
use crossbeam_utils::atomic::AtomicCell;
use bytes::BytesMut;
use pi_async_rt::{lock::spin_lock::SpinLock,
                  rt::{AsyncValue,
                       serial_local_thread::LocalTaskRuntime}};

use udp::SocketHandle as UdpSocketHandle;

pub mod acceptor;
pub mod connect;
pub mod connect_pool;
pub mod server;
pub mod client;
pub mod utils;

use crate::connect::QuicSocket;
use crate::utils::{QuicSocketStatus, QuicSocketReady, Hibernate, ContextHandle};

// 仅测试构建观测真实登记的原参数/Result，不新增生产状态或可见性。
#[cfg(all(test, feature = "runtime_metrics", not(target_arch = "wasm32")))]
std::thread_local! {
    static QUIC_METRICS_TEST_RESULT: std::cell::RefCell<Option<(usize, String,
        std::result::Result<pi_async_rt::runtime_metrics::RuntimeMetricId,
            pi_async_rt::runtime_metrics::RuntimeMetricRegistrationError>)>> = const { std::cell::RefCell::new(None) };
}

/// 供本crate冷装配复用的原生Local登记；不创建核心、不读队列、不派发任务。
/// runtime仅借用，返回()不表示成功；同名重复及真实核心身份由上游判断。
/// 名称O(L)时间/空间，另有上游短同步锁/分配/强保活及失败logger可能等待的成本；
/// 无await/异步锁，不持业务锁调用，不处理close/Drop/失败注销。功能关闭时无操作。
/// 专项：batch_runtime_metrics_tests::test_quic_registers_runtime_metrics，仅参数/Ok。
fn register_local_runtime_metrics(runtime: &LocalTaskRuntime<()>) {
    #[cfg(all(feature = "runtime_metrics", not(target_arch = "wasm32")))]
    {
        let metric_name = format!("pi_quic-local-{}", runtime.get_id());
        let registration = pi_async_rt::runtime_metrics::register_runtime_metrics(runtime, &metric_name);
        #[cfg(test)]
        QUIC_METRICS_TEST_RESULT.with(|value| {
            *value.borrow_mut() = Some((runtime.get_id(), metric_name.clone(), registration));
        });
        if let Err(error) = registration {
            log::warn!("QUIC 运行时指标注册失败, 名称: {}, 原因: {}", metric_name, error);
        }
    }
    #[cfg(not(all(feature = "runtime_metrics", not(target_arch = "wasm32"))))]
    let _ = runtime;
}

#[cfg(all(test, feature = "runtime_metrics", not(target_arch = "wasm32")))]
mod batch_runtime_metrics_tests {
    //! 只验证register_local_runtime_metrics向上游传入真实Local/稳定名称且原Result为Ok。
    //! 这是无网络局部单元，不执行UDP/QUIC循环，不声称完整网络冷装配已动态测试。
    //! 各实际调用点/原顺序由整批源码审查核对，注册算法/SDK/生命周期不重复测试。

    use super::*;

    /// 被测生产入口：register_local_runtime_metrics，与冷启动实际共用，不手动预登记或模拟上游。
    /// 本测试自有真实空Local，不派发采样/业务任务，不读本地队列；取真实返回ID后
    /// 在断言前显式注销清理，再随空runner释放，不为观测扩大生产可见性。
    /// O(L)命名核对，严格串行，无await/网络/无界等待。
    #[test]
    fn test_quic_registers_runtime_metrics() {
        QUIC_METRICS_TEST_RESULT.with(|value| { value.borrow_mut().take(); });
        let runner = pi_async_rt::rt::serial_local_thread::LocalTaskRunner::<()>::new();
        let runtime = runner.get_runtime();
        let expected_id = runtime.get_id();
        register_local_runtime_metrics(&runtime);
        let observed = QUIC_METRICS_TEST_RESULT.with(|value| value.borrow_mut().take());
        if let Some((_, _, Ok(id))) = &observed {
            let _ = pi_async_rt::runtime_metrics::unregister_runtime_metrics(*id);
        }
        let (actual_id, actual_name, registration) = observed.expect("实际冷登记必须发生");
        assert_eq!(actual_id, expected_id);
        assert_eq!(actual_name, format!("pi_quic-local-{}", expected_id));
        assert!(registration.is_ok(), "真实登记必须成功: {:?}", registration);
    }
}

/// QUIC连接异步服务；本次登记不改变其业务回调。
///
pub trait AsyncService: Send + Sync + 'static {
    /// 异步处理已连接
    fn handle_connected(&self,
                        handle: SocketHandle,
                        result: Result<()>) -> LocalBoxFuture<'static, ()>;

    /// 异步处理已打开扩展流
    fn handle_opened_expanding_stream(&self,
                                      handle: SocketHandle,
                                      stream_id: StreamId,
                                      stream_type: Dir,
                                      result: Result<()>) -> LocalBoxFuture<'static, ()>;

    /// 异步处理已读
    fn handle_readed(&self,
                     handle: SocketHandle,
                     stream_id: StreamId,
                     result: Result<usize>) -> LocalBoxFuture<'static, ()>;

    /// 异步处理已写
    fn handle_writed(&self,
                     handle: SocketHandle,
                     stream_id: StreamId,
                     result: Result<()>) -> LocalBoxFuture<'static, ()>;

    /// 异步处理已关闭
    fn handle_closed(&self,
                     handle: SocketHandle,
                     stream_id: Option<StreamId>,
                     code: u32,
                     result: Result<()>) -> LocalBoxFuture<'static, ()>;

    /// 异步处理已超时
    fn handle_timeouted(&self,
                        handle: SocketHandle,
                        result: Result<SocketEvent>) -> LocalBoxFuture<'static, ()>;
}

///
/// Quic连接句柄
///
pub struct SocketHandle(Arc<InnerSocketHandle>);

unsafe impl Send for SocketHandle {}
unsafe impl Sync for SocketHandle {}

impl Clone for SocketHandle {
    fn clone(&self) -> Self {
        SocketHandle(self.0.clone())
    }
}

impl SocketHandle {
    /// 构建一个Quic连接句柄
    pub fn new(uid: usize,
               status: Arc<AtomicU8>,
               socket: Arc<UnsafeCell<QuicSocket>>,
               local: SocketAddr,
               remote: SocketAddr,
               sender: Sender<QuicEvent>) -> Self {
        let inner = InnerSocketHandle {
            uid,
            status,
            inner: socket,
            local: AtomicCell::new(local),
            remote,
            sender,
        };

        SocketHandle(Arc::new(inner))
    }

    /// 线程安全的异步广播指定负载
    pub fn broadcast(handles: &[SocketHandle],
                     payload: Vec<u8>) -> Result<()> {
        if handles.len() == 0 {
            //连接为空，则忽略
            return Ok(());
        }

        for handle in handles {
            for stream_id in handle.get_stream_ids() {
                handle.write_ready(stream_id, payload.clone())?;
            }
        }

        Ok(())
    }

    /// 获取连接唯一id
    pub fn get_uid(&self) -> usize {
        self.0.uid
    }

    /// 获取连接状态
    pub fn get_status(&self) -> QuicSocketStatus {
        self
            .0
            .status
            .load(Ordering::Acquire)
            .into()
    }

    /// 获取本地Udp连接地址
    pub fn get_local_udp(&self) -> &SocketAddr {
        unsafe {
            (&*self.0.inner.get())
                .get_udp_handle()
                .get_local()
        }
    }

    /// 获取远端Udp连接地址
    pub fn get_remote_udp(&self) -> Option<&SocketAddr> {
        unsafe {
            (&*self.0.inner.get())
                .get_udp_handle()
                .get_remote()
        }
    }

    /// 获取本地连接地址
    pub fn get_local(&self) -> SocketAddr {
        self
            .0
            .local
            .load()
    }

    /// 设置本地连接地址
    pub fn set_local(&self, local: SocketAddr) {
        self
            .0
            .local
            .store(local);
    }

    /// 获取远端连接地址
    pub fn get_remote(&self) -> SocketAddr {
        self.0.remote
    }

    /// 获取Quic连接对应的Udp连接的唯一id
    pub fn get_udp_uid(&self) -> usize {
        unsafe {
            (&*self.0.inner.get()).get_udp_handle().get_uid()
        }
    }

    /// 判断当前连接是否是客户端连接
    pub fn is_client(&self) -> bool {
        unsafe {
            (&*self.0.inner.get()).is_client()
        }
    }

    /// 获取内部连接句柄
    pub fn get_connection_handle(&self) -> &ConnectionHandle {
        unsafe {
            (&*self.0.inner.get()).get_connection_handle()
        }
    }

    /// 获取Quic连接会话的句柄
    pub fn get_session<T>(&self) -> Option<ContextHandle<T>> {
        unsafe {
            (&*(&*self.0.inner.get()).get_context().get()).get::<T>()
        }
    }

    /// 设置Quic连接会话的句柄
    pub fn set_session<T>(&self, context: T) {
        unsafe {
            (&mut *(&*self.0.inner.get()).get_context().get()).set::<T>(context);
        }
    }

    /// 移除Quic连接会话的句柄
    pub fn remove_session<T>(&self) -> Result<Option<T>> {
        unsafe {
            match (&mut *(&*self.0.inner.get()).get_context().get()).remove::<T>() {
                Err(e) => Err(Error::new(ErrorKind::Other, e)),
                Ok(r) => Ok(r)
            }
        }
    }

    /// 判断当前连接是否通过0rtt建立连接
    pub fn is_0rtt(&self) -> bool {
        unsafe {
            (&*self.0.inner.get()).is_0rtt()
        }
    }

    /// 判断当前连接是否已关闭
    pub fn is_closed(&self) -> bool {
        unsafe {
            (&*self.0.inner.get()).is_closed()
        }
    }

    /// 获取当前连接的延迟估计
    pub fn get_latency(&self) -> Duration {
        unsafe {
            (&*self.0.inner.get()).get_latency()
        }
    }

    /// 获取连接的主流唯一id
    pub fn get_main_stream_id(&self) -> Option<&StreamId> {
        unsafe {
            (&*self.0.inner.get()).get_main_stream_id()
        }
    }

    /// 打开连接的主流，主流一定是双向流
    pub(crate) fn open_main_streams(&self) -> Result<()> {
        unsafe {
            (&mut *self.0.inner.get()).open_main_streams()
        }
    }

    /// 打开连接的扩展流，可以指定流的类型
    pub async fn open_expanding_stream(&self, stream_type: Dir) -> Result<StreamId> {
        unsafe {
            (&*self.0.inner.get()).open_expanding_stream(stream_type).await
        }
    }

    /// 设置当前连接感兴趣的事件
    pub fn set_ready(&self,
                     stream_id: StreamId,
                     ready: QuicSocketReady) {
        unsafe {
            (&*self.0.inner.get()).set_ready(stream_id, ready);
        }
    }

    /// 获取连接的所有流唯一id
    pub fn get_stream_ids(&self) -> Vec<StreamId> {
        unsafe {
            (&*self.0.inner.get()).get_stream_ids()
        }
    }

    /// 通知连接的指定流读就绪
    pub fn read_ready(&self,
                      stream_id: &StreamId,
                      adjust: usize) -> GenResult<AsyncValue<usize>, usize> {
        unsafe {
            (&mut *self.0.inner.get()).read_ready(stream_id, adjust)
        }
    }

    /// 获取连接的输入缓冲区的剩余未读字节数
    pub fn read_buffer_remaining(&self, stream_id: &StreamId) -> Option<usize> {
        unsafe {
            (&*self.0.inner.get()).read_buffer_remaining(stream_id)
        }
    }

    /// 获取连接的输入缓冲区的只读引用
    pub fn get_read_buffer(&self,
                           stream_id: &StreamId) -> Option<Arc<SpinLock<Option<BytesMut>>>> {
        unsafe {
            (&*self.0.inner.get()).get_read_buffer(stream_id)
        }
    }

    /// 通知连接写就绪，可以开始发送指定的数据
    pub fn write_ready<B>(&self,
                          stream_id: StreamId,
                          buf: B) -> Result<()>
        where B: AsRef<[u8]> + 'static {
        unsafe {
            (&*self.0.inner.get()).write_ready(stream_id, buf)
        }
    }

    /// 设置连接的超时定时器，同时只允许设置一个定时器，新的定时器会覆盖未超时的旧定时器
    pub fn set_timeout(&self, timeout: usize, event: SocketEvent) {
        unsafe {
            (&*self.0.inner.get()).set_timeout(timeout, event);
        }
    }

    /// 取消连接的未超时超时定时器
    pub fn unset_timeout(&self) {
        unsafe {
            (&*self.0.inner.get()).unset_timeout();
        }
    }

    /// 开始执行连接休眠时加入的任务，当前任务执行完成后自动执行下一个任务，直到任务队列为空
    pub fn run_hibernated_tasks(&self) {
        unsafe {
            (&*self.0.inner.get()).run_hibernated_tasks();
        }
    }

    /// 获取当前连接的休眠对象，返回空表示连接已关闭
    pub fn hibernate(&self,
                     stream_id: StreamId,
                     ready: QuicSocketReady) -> Option<Hibernate> {
        unsafe {
            (&*self.0.inner.get()).hibernate(self.clone(),
                                             stream_id,
                                             ready)
        }
    }

    /// 设置当前连接的休眠对象，设置成功返回真
    pub fn set_hibernate(&self, hibernate: Hibernate) -> bool {
        unsafe {
            (&*self.0.inner.get()).set_hibernate(hibernate)
        }
    }

    /// 设置当前连接在休眠时挂起的其它休眠对象的唤醒器
    pub fn set_hibernate_wakers(&self, waker: Waker) {
        unsafe {
            (&*self.0.inner.get()).set_hibernate_wakers(waker);
        }
    }

    /// 非阻塞的唤醒被休眠的当前连接，如果当前连接未被休眠，则忽略
    /// 还会唤醒当前连接正在休眠时，当前连接的所有其它休眠对象的唤醒器
    /// 唤醒过程可能会被阻塞，这不会导致线程阻塞而是返回假，调用者可以继续尝试唤醒，直到返回真
    pub fn wakeup(&self, result: Result<()>) -> bool {
        unsafe {
            (&mut *self.0.inner.get()).wakeup(result)
        }
    }

    /// 线程安全的获取连接所在运行时
    pub fn get_runtime(&self) -> Option<LocalTaskRuntime<()>> {
        unsafe {
            (&*self.0.inner.get()).get_runtime()
        }
    }

    /// 线程安全的派发一个异步任务到连接所在运行时
    pub fn spawn(&self, task: LocalBoxFuture<'static, ()>) {
        unsafe {
            (&*self.0.inner.get()).spawn(task);
        }
    }

    /// 关闭指定的流
    pub fn close_expanding_stream(&self, stream_id: StreamId) -> Result<()> {
        unsafe {
            (&*self.0.inner.get()).close_expanding_stream(stream_id)
        }
    }

    /// 线程安全的关闭Quic连接
    pub fn close(&self, code: u32, reason: Result<()>) -> Result<()> {
        unsafe {
            (&mut *self.0.inner.get()).close(code, reason)
        }
    }

    /// 强制关闭当前连接
    pub fn force_close(&mut self, code: u32, reason: String) {
        unsafe {
            (&mut *self.0.inner.get()).force_close(code, reason);
        }
    }
}

// 内部Quic连接句柄
struct InnerSocketHandle {
    uid:            usize,                                                      //Quic连接唯一id
    status:         Arc<AtomicU8>,                                              //Quic连接状态
    inner:          Arc<UnsafeCell<QuicSocket>>,                                //Quic连接指针
    local:          AtomicCell<SocketAddr>,                                     //Quic连接本地地址
    remote:         SocketAddr,                                                 //Quic连接远端地址
    sender:         Sender<QuicEvent>,                                          //Quic事件发送器
}

///
/// Quic连接的事件
///
#[derive(Debug)]
pub struct SocketEvent {
    inner: *mut (), //内部事件
}

unsafe impl Send for SocketEvent {}

impl SocketEvent {
    /// 创建空的事件
    pub fn empty() -> Self {
        SocketEvent {
            inner: ptr::null_mut(),
        }
    }

    /// 判断事件是否为空
    pub fn is_empty(&self) -> bool {
        self.inner.is_null()
    }

    /// 获取事件
    pub fn get<T: 'static>(&self) -> Option<T> {
        if self.is_empty() {
            return None;
        }

        Some(unsafe { *Box::from_raw(self.inner as *mut T) })
    }

    /// 设置事件，如果当前事件不为空，则设置失败
    pub fn set<T: 'static>(&mut self, event: T) -> bool {
        if !self.is_empty() {
            return false;
        }

        self.inner = Box::into_raw(Box::new(event)) as *mut ();
        true
    }

    /// 移除事件，如果当前有事件，则返回被移除的事件
    pub fn remove<T: 'static>(&mut self) -> Option<T> {
        if self.is_empty() {
            return None;
        }

        let result = self.get();
        self.inner = ptr::null_mut();
        result
    }
}

///
/// Quic事件
///
#[derive(Debug)]
pub enum QuicEvent {
    Accepted(QuicSocket),                                                       //Quic已接受连接
    ConnectionReceived(ConnectionHandle, ConnectionEvent),                      //Quic连接接收数据
    ConnectionSend(ConnectionHandle, Transmit),                                 //Quic连接发送数据
    StreamOpen(ConnectionHandle, Dir, AsyncValue<Result<StreamId>>), //Quic流打开
    StreamReady(ConnectionHandle, StreamId, QuicSocketReady),                   //Quic流就绪
    StreamWrite(ConnectionHandle, StreamId, Vec<u8>),                           //Quic流写数据
    RebindUdp(UdpSocketHandle),                                                 //Quic重绑定Udp连接句柄
    Timeout(ConnectionHandle, Option<(usize, SocketEvent)>),                    //Quic连接超时
    StreamClose(ConnectionHandle, Option<StreamId>, u32),                       //Quic流关闭
    ConnectionClose(ConnectionHandle),                                          //Quic连接关闭
}