smux_rust 0.2.1

use crate::alloc::Allocator;
use crate::config::Config;
use crate::error::SmuxError;
use crate::frame::{Cmd, Frame, RawHeader, UpdHeader, HEADER_SIZE, UPD_SIZE};
use crate::shaper::{ClassId, ShaperQueue, WriteRequest};
use crate::stream::{Stream, StreamInner};
use bytes::Bytes;
use std::collections::HashMap;
use std::io;
use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::{mpsc, Mutex, Notify};
use tokio::time::{sleep, Duration as TokioDuration};

/// 默认接受队列大小
const DEFAULT_ACCEPT_BACKLOG: usize = 1024;
const MIN_SHAPER_NOTIFY_SIZE: usize = 16;
const OPEN_CLOSE_TIMEOUT: Duration = Duration::from_secs(30);

/// Session 管理多个流的多路复用连接
pub struct Session {
    /// 底层连接的读端
    conn_read: Arc<Mutex<Box<dyn AsyncRead + Send + Unpin>>>,
    /// 底层连接的写端
    conn_write: Arc<Mutex<Box<dyn AsyncWrite + Send + Unpin>>>,
    /// 配置
    pub config: Config,
    /// 下一个流ID
    next_stream_id: Mutex<u32>,
    /// 是否为客户端
    is_client: bool,
    /// Token bucket（接收缓冲区）
    bucket: AtomicI32,
    /// Bucket 通知
    bucket_notify: Arc<Notify>,
    /// 所有流
    streams: Arc<Mutex<HashMap<u32, Arc<StreamInner>>>>,
    /// 会话关闭标志
    die: Arc<AtomicBool>,
    /// Socket 读取错误
    socket_read_error: Arc<Mutex<Option<io::Error>>>,
    /// Socket 写入错误
    socket_write_error: Arc<Mutex<Option<io::Error>>>,
    /// Socket 读取错误通知
    ch_socket_read_error: Arc<Notify>,
    /// Socket 写入错误通知
    ch_socket_write_error: Arc<Notify>,
    /// 协议错误
    proto_error: Arc<Mutex<Option<SmuxError>>>,
    /// 协议错误通知
    ch_proto_error: Arc<Notify>,
    /// 接受队列
    ch_accepts: Arc<Mutex<mpsc::Receiver<Arc<StreamInner>>>>,
    ch_accepts_tx: mpsc::Sender<Arc<StreamInner>>,
    /// 数据就绪标志
    data_ready: AtomicBool,
    /// GoAway 标志
    go_away: AtomicBool,
    /// 截止时间
    deadline: Arc<Mutex<Option<Instant>>>,
    /// 请求ID（单调递增）
    request_id: AtomicU32,
    /// 整形队列
    shaper: Arc<Mutex<ShaperQueue>>,
    /// 整形队列待处理通知
    ch_shaper_pending: Arc<Notify>,
    /// 内存分配器
    allocator: Arc<Allocator>,
    /// 本地地址（可选）
    local_addr: Option<String>,
    /// 远程地址（可选）
    remote_addr: Option<String>,
}

/// 异步读写 trait
pub trait AsyncReadWrite: AsyncRead + AsyncWrite + Send + Unpin {}
impl<T: AsyncRead + AsyncWrite + Send + Unpin> AsyncReadWrite for T {}

impl Session {
    /// 创建新的会话
    pub(crate) async fn new(
        config: Config,
        conn: Box<dyn AsyncReadWrite + Send + Unpin>,
        is_client: bool,
    ) -> Result<Arc<Self>, SmuxError> {
        config.verify().map_err(|_e| SmuxError::Io(io::ErrorKind::InvalidInput))?;

        let (accepts_tx, accepts_rx) = mpsc::channel(DEFAULT_ACCEPT_BACKLOG);

        // 使用 tokio::io::split 分离读写端
        let (read_half, write_half) = tokio::io::split(conn);

        let session = Arc::new(Self {
            conn_read: Arc::new(Mutex::new(Box::new(read_half))),
            conn_write: Arc::new(Mutex::new(Box::new(write_half))),
            config: config.clone(),
            next_stream_id: Mutex::new(if is_client { 1 } else { 0 }),
            is_client,
            bucket: AtomicI32::new(config.max_receive_buffer as i32),
            bucket_notify: Arc::new(Notify::new()),
            streams: Arc::new(Mutex::new(HashMap::new())),
            die: Arc::new(AtomicBool::new(false)),
            socket_read_error: Arc::new(Mutex::new(None)),
            socket_write_error: Arc::new(Mutex::new(None)),
            ch_socket_read_error: Arc::new(Notify::new()),
            ch_socket_write_error: Arc::new(Notify::new()),
            proto_error: Arc::new(Mutex::new(None)),
            ch_proto_error: Arc::new(Notify::new()),
            ch_accepts: Arc::new(Mutex::new(accepts_rx)),
            ch_accepts_tx: accepts_tx,
            data_ready: AtomicBool::new(false),
            go_away: AtomicBool::new(false),
            deadline: Arc::new(Mutex::new(None)),
            request_id: AtomicU32::new(0),
            shaper: Arc::new(Mutex::new(ShaperQueue::new())),
            ch_shaper_pending: Arc::new(Notify::new()),
            allocator: Arc::new(Allocator::new()),
            local_addr: None,
            remote_addr: None,
        });

        // 启动后台任务
        let session_clone = Arc::clone(&session);
        tokio::spawn(async move {
            session_clone.recv_loop().await;
        });

        let session_clone = Arc::clone(&session);
        tokio::spawn(async move {
            session_clone.send_loop().await;
        });

        if !config.keep_alive_disabled {
            let session_clone = Arc::clone(&session);
            tokio::spawn(async move {
                session_clone.keepalive().await;
            });
        }

        // 触发一次初始通知，确保 send_loop 立即开始处理队列
        session.ch_shaper_pending.notify_one();

        Ok(session)
    }

    /// 打开新流
    pub async fn open_stream(self: &Arc<Self>) -> Result<Stream, SmuxError> {
        if self.is_closed() {
            return Err(SmuxError::Io(io::ErrorKind::BrokenPipe));
        }

        // 生成流ID
        let sid = {
            let mut next_id = self.next_stream_id.lock().await;
            if self.go_away.load(Ordering::Relaxed) {
                return Err(SmuxError::GoAway);
            }

            *next_id += 2;
            let sid = *next_id;
            if sid % 2 == 0 && !self.is_client {
                // 流ID溢出
                self.go_away.store(true, Ordering::Relaxed);
                return Err(SmuxError::GoAway);
            }
            sid
        };

        let stream = Arc::new(StreamInner::new(sid, self.config.max_frame_size, Arc::clone(self)));

        // 发送 SYN 帧
        let frame = Frame::new(self.config.version, Cmd::Syn, sid);
        self.write_control_frame(frame).await?;

        // 添加到流映射
        {
            let mut streams = self.streams.lock().await;
            streams.insert(sid, stream.clone());
        }

        Ok(Stream::new(stream))
    }

    /// 接受流
    pub async fn accept_stream(self: &Arc<Self>) -> Result<Stream, SmuxError> {
        let deadline_opt = *self.deadline.lock().await;

        loop {
            // 检查错误
            if let Some(err) = self.socket_read_error.lock().await.as_ref() {
                return Err(SmuxError::Io(err.kind()));
            }

            if let Some(err) = self.proto_error.lock().await.as_ref() {
                return Err(err.clone());
            }

            if self.is_closed() {
                return Err(SmuxError::Io(io::ErrorKind::BrokenPipe));
            }

            // 等待接受流或超时
            if let Some(deadline) = deadline_opt {
                let now = Instant::now();
                if now >= deadline {
                    return Err(SmuxError::Timeout);
                }
                let timeout_duration = deadline.duration_since(now);
                let mut rx = self.ch_accepts.lock().await;
                tokio::select! {
                    stream = rx.recv() => {
                        if let Some(stream) = stream {
                            return Ok(crate::stream::Stream::new(stream));
                        }
                        continue;
                    }
                    _ = sleep(TokioDuration::from_secs(timeout_duration.as_secs())) => {
                        return Err(SmuxError::Timeout);
                    }
                    _ = self.ch_socket_read_error.notified() => {
                        continue;
                    }
                    _ = self.ch_proto_error.notified() => {
                        continue;
                    }
                }
            } else {
                let mut rx = self.ch_accepts.lock().await;
                tokio::select! {
                    stream = rx.recv() => {
                        if let Some(stream) = stream {
                            return Ok(crate::stream::Stream::new(stream));
                        }
                        continue;
                    }
                    _ = self.ch_socket_read_error.notified() => {
                        continue;
                    }
                    _ = self.ch_proto_error.notified() => {
                        continue;
                    }
                }
            };
        }
    }

    /// 关闭会话
    pub async fn close(&self) -> io::Result<()> {
        if self.die.swap(true, Ordering::Relaxed) {
            return Err(io::Error::from(io::ErrorKind::BrokenPipe));
        }

        // 关闭所有流
        let streams = self.streams.lock().await;
        for stream in streams.values() {
            let _ = stream.close().await;
        }

        Ok(())
    }

    /// 检查是否已关闭
    pub fn is_closed(&self) -> bool {
        self.die.load(Ordering::Relaxed)
    }

    /// 返回当前活跃流的数量
    pub async fn num_streams(&self) -> usize {
        let streams = self.streams.lock().await;
        streams.len()
    }

    /// 设置会话的截止时间
    pub async fn set_deadline(&self, deadline: Option<Instant>) {
        *self.deadline.lock().await = deadline;
    }

    /// 获取本地地址
    pub fn local_addr(&self) -> Option<&str> {
        self.local_addr.as_deref()
    }

    /// 获取远程地址
    pub fn remote_addr(&self) -> Option<&str> {
        self.remote_addr.as_deref()
    }

    /// 设置本地地址（用于外部设置地址信息）
    pub fn set_local_addr(&mut self, addr: String) {
        self.local_addr = Some(addr);
    }

    /// 设置远程地址（用于外部设置地址信息）
    pub fn set_remote_addr(&mut self, addr: String) {
        self.remote_addr = Some(addr);
    }

    /// 等待多个流中任意一个变为可读
    /// 返回第一个有数据可读的流的索引
    /// 
    /// 这是一个高级功能，用于实现高效的多路复用 I/O
    /// 
    /// # 示例
    /// ```no_run
    /// # use smux_rust::*;
    /// # async fn example(session: std::sync::Arc<Session>, s1: &Stream, s2: &Stream, s3: &Stream) -> Result<(), SmuxError> {
    /// let streams = vec![s1, s2, s3];
    /// let idx = session.poll_wait(&streams).await?;
    /// println!("流 {} 有数据可读", idx);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn poll_wait(&self, streams: &[&crate::stream::Stream]) -> Result<usize, SmuxError> {
        use tokio::time::{sleep, Duration};
        
        if streams.is_empty() {
            return Err(SmuxError::Io(io::ErrorKind::InvalidInput));
        }

        loop {
            // 检查会话是否关闭
            if self.is_closed() {
                return Err(SmuxError::Io(io::ErrorKind::BrokenPipe));
            }

            // 首先快速检查是否有流立即可读
            for (idx, stream) in streams.iter().enumerate() {
                let inner = &stream.inner;
                
                // 检查缓冲区是否有数据
                if inner.has_buffered_data().await {
                    return Ok(idx);
                }
                
                // 检查流是否已关闭（EOF）
                if inner.is_fin().await {
                    return Ok(idx);
                }
            }

            // 没有立即可读的流，等待任意一个流的通知
            // 使用 futures::future::select_all 会更高效，但为了不引入新依赖
            // 我们创建一个 channel 来收集通知
            let (tx, mut rx) = tokio::sync::mpsc::channel(streams.len());
            
            // 为每个流启动一个监听任务
            let mut handles = Vec::new();
            for (idx, stream) in streams.iter().enumerate() {
                let notifier = stream.inner.get_read_event_notifier();
                let tx_clone = tx.clone();
                
                let handle = tokio::spawn(async move {
                    notifier.notified().await;
                    let _ = tx_clone.send(idx).await;
                });
                
                handles.push(handle);
            }
            
            // 释放原始的发送端
            drop(tx);
            
            // 等待任意一个流就绪，或超时
            let result = tokio::select! {
                Some(idx) = rx.recv() => {
                    // 有流收到通知
                    Ok(idx)
                }
                _ = sleep(Duration::from_millis(50)) => {
                    // 超时，重新检查（这避免了可能的竞争条件）
                    continue;
                }
                _ = self.ch_socket_read_error.notified() => {
                    // Socket 错误
                    Err(SmuxError::Io(io::ErrorKind::BrokenPipe))
                }
            };
            
            // 取消所有监听任务
            for handle in handles {
                handle.abort();
            }
            
            if let Ok(idx) = result {
                return Ok(idx);
            } else if let Err(e) = result {
                return Err(e);
            }
        }
    }

    /// 返回令牌到 bucket
    pub(crate) fn return_tokens(&self, n: usize) {
        let new_bucket = self.bucket.fetch_add(n as i32, Ordering::Relaxed) + n as i32;
        if new_bucket > 0 {
            self.bucket_notify.notify_one();
        }
    }

    /// 通知 bucket 可用
    fn notify_bucket(&self) {
        self.bucket_notify.notify_one();
    }

    /// 流关闭回调
    pub(crate) async fn stream_closed(&self, sid: u32) {
        let mut streams = self.streams.lock().await;
        if let Some(stream) = streams.remove(&sid) {
            let n = stream.recycle_tokens().await;
            if n > 0 {
                self.return_tokens(n);
            }
        }
    }

    /// 写入控制帧
    async fn write_control_frame(&self, frame: Frame) -> Result<usize, SmuxError> {
        self.write_frame_internal(frame, Some(OPEN_CLOSE_TIMEOUT), ClassId::Ctrl).await
    }

    /// 内部写入帧
    pub(crate) async fn write_frame(&self, frame: Frame, is_control: bool) -> Result<usize, SmuxError> {
        let class = if is_control {
            ClassId::Ctrl
        } else {
            ClassId::Data
        };
        self.write_frame_internal(frame, None, class).await
    }

    /// 内部写入帧（支持截止时间）
    async fn write_frame_internal(
        &self,
        frame: Frame,
        _timeout: Option<Duration>,
        class: ClassId,
    ) -> Result<usize, SmuxError> {
        let seq = self.request_id.fetch_add(1, Ordering::Relaxed);
        let data_len = frame.data.len();
        let req = WriteRequest {
            class,
            frame,
            seq,
        };

        // 添加到整形队列
        {
            let mut shaper = self.shaper.lock().await;
            shaper.push(req).await;
            let len = shaper.len().await;
            
            
            // 对于控制帧（尤其是 SYN/FIN），立即通知以避免死锁
            // 对于数据帧，等待队列积累到一定数量再通知以提高效率
            if class == ClassId::Ctrl || len >= MIN_SHAPER_NOTIFY_SIZE {
                self.ch_shaper_pending.notify_one();
            }
        }

        Ok(data_len)
    }

    /// 接收循环
    async fn recv_loop(self: Arc<Self>) {
        let mut header_buf = vec![0u8; HEADER_SIZE];

        loop {
            // 等待 bucket 可用
            while self.bucket.load(Ordering::Relaxed) <= 0 && !self.is_closed() {
                self.bucket_notify.notified().await;
            }

            if self.is_closed() {
                return;
            }

            // 读取帧头
            let read_result = {
                let mut conn_read = self.conn_read.lock().await;
                conn_read.read_exact(&mut header_buf).await
            };
            
            match read_result {
                Ok(_) => {}
                Err(e) => {
                    *self.socket_read_error.lock().await = Some(e);
                    self.ch_socket_read_error.notify_waiters();
                    return;
                }
            }

            self.data_ready.store(true, Ordering::Relaxed);

            let header = match RawHeader::from_bytes(&header_buf) {
                Some(h) => h,
                None => {
                    *self.proto_error.lock().await = Some(SmuxError::InvalidProtocol);
                    self.ch_proto_error.notify_waiters();
                    return;
                }
            };

            if header.ver != self.config.version {
                *self.proto_error.lock().await = Some(SmuxError::InvalidProtocol);
                self.ch_proto_error.notify_waiters();
                return;
            }

            let sid = header.sid;

            match header.cmd {
                Cmd::Nop => {
                    // 无操作
                }
                Cmd::Syn => {
                    // 流打开
                    let mut streams = self.streams.lock().await;
                    if !streams.contains_key(&sid) {
                        let stream = Arc::new(StreamInner::new(
                            sid,
                            self.config.max_frame_size,
                            Arc::clone(&self),
                        ));
                        streams.insert(sid, stream.clone());

                        // 发送到接受队列（忽略错误，如果队列满了）
                        let _ = self.ch_accepts_tx.send(stream).await;
                    }
                }
                Cmd::Fin => {
                    // 流关闭
                    let streams = self.streams.lock().await;
                    if let Some(stream) = streams.get(&sid) {
                        stream.fin();
                    }
                }
                Cmd::Psh => {
                    // 数据帧
                    if header.length == 0 {
                        continue;
                    }

                    let data_len = header.length as usize;
                    let data = match self.allocator.get(data_len).await {
                        Some(mut buf) => {
                            buf.resize(data_len, 0);
                            let read_result = {
                                let mut conn_read = self.conn_read.lock().await;
                                conn_read.read_exact(&mut buf).await
                            };
                            
                            match read_result {
                                Ok(_) => Bytes::from(buf),
                                Err(e) => {
                                    *self.socket_read_error.lock().await = Some(e);
                                    self.ch_socket_read_error.notify_waiters();
                                    return;
                                }
                            }
                        }
                        None => continue,
                    };

                    let streams = self.streams.lock().await;
                    if let Some(stream) = streams.get(&sid) {
                        stream.push_bytes(data.clone()).await;
                        let written = data.len();
                        self.bucket.fetch_sub(written as i32, Ordering::Relaxed);
                    }
                }
                Cmd::Upd => {
                    // 窗口更新
                    let mut upd_buf = vec![0u8; UPD_SIZE];
                    let read_result = {
                        let mut conn_read = self.conn_read.lock().await;
                        conn_read.read_exact(&mut upd_buf).await
                    };
                    
                    match read_result {
                        Ok(_) => {}
                        Err(e) => {
                            *self.socket_read_error.lock().await = Some(e);
                            self.ch_socket_read_error.notify_waiters();
                            return;
                        }
                    }

                    let upd = match UpdHeader::from_bytes(&upd_buf) {
                        Some(u) => u,
                        None => continue,
                    };

                    let streams = self.streams.lock().await;
                    if let Some(stream) = streams.get(&sid) {
                        stream.update(upd.consumed, upd.window).await;
                    }
                }
            }
        }
    }

    /// 发送循环
    async fn send_loop(self: Arc<Self>) {
        loop {
            if self.is_closed() {
                return;
            }

            // 等待待处理的请求，但设置超时避免死锁
            tokio::select! {
                _ = self.ch_shaper_pending.notified() => {
                }
                _ = sleep(TokioDuration::from_millis(10)) => {
                    // 超时，检查队列（避免错过通知）
                }
            }

            // 从整形队列中弹出请求并发送
            loop {
                let req = {
                    let mut shaper = self.shaper.lock().await;
                    shaper.pop().await
                };

                if let Some(req) = req {
                    let encoded = req.frame.encode();
                    
                    let write_result = {
                        let mut conn_write = self.conn_write.lock().await;
                        conn_write.write_all(&encoded).await
                    };
                    
                    match write_result {
                        Ok(_) => {
                        }
                        Err(e) => {
                            *self.socket_write_error.lock().await = Some(e);
                            self.ch_socket_write_error.notify_waiters();
                            return;
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    /// 保活循环
    async fn keepalive(self: Arc<Self>) {
        let ping_interval = self.config.keep_alive_interval;
        let timeout_interval = self.config.keep_alive_timeout;

        loop {
            sleep(TokioDuration::from_secs(ping_interval.as_secs())).await;

            if self.is_closed() {
                return;
            }

            // 发送 NOP 帧
            let frame = Frame::new(self.config.version, Cmd::Nop, 0);
            let _ = self.write_frame_internal(frame, Some(ping_interval), ClassId::Ctrl).await;
            self.notify_bucket();

            // 检查超时
            sleep(TokioDuration::from_secs(
                (timeout_interval - ping_interval).as_secs(),
            ))
            .await;
            if !self.data_ready.swap(false, Ordering::Relaxed) {
                if self.bucket.load(Ordering::Relaxed) > 0 {
                    let _ = self.close().await;
                    return;
                }
            }
        }
    }
}