如何在Rust中将AsyncReader拆分为两个独立读取器？

要实现你需要的双路异步读取功能，核心是构建一个基于共享缓存的分叉读取器，让两个读取实例共享底层流的数据，同时各自维护读取进度，且缓存能动态扩容。以下是完整的实现方案：

核心设计 #

• 用线程安全的共享状态存储底层流、缓存队列和两个读取器的读取偏移量
• 读取器优先从缓存取数据，缓存不足时从底层流读取并填充缓存
• 采用VecDeque作为缓存容器，自动适配动态扩容需求

完整代码实现 #

1. 定义共享状态与读取器类型 #

use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

// 共享状态：包含底层流、缓存、两个读取器的偏移
struct SharedState<R>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    inner: R,
    buffer: VecDeque<u8>,
    read_offsets: [usize; 2], // 两个读取器各自已读取的字节数
}

impl<R> SharedState<R>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    fn new(inner: R) -> Self {
        Self {
            inner,
            buffer: VecDeque::new(),
            read_offsets: [0, 0],
        }
    }
}

// 分叉后的读取器，持有共享状态的引用和自身ID
#[derive(Clone)]
pub struct ForkedReader<R>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    state: Arc<Mutex<SharedState<R>>>,
    reader_id: usize,
}

2. 实现AsyncRead trait #

impl<R> AsyncRead for ForkedReader<R>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        let mut state = match self.state.lock() {
            Ok(s) => s,
            Err(_) => return Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "Poisoned mutex",
            ))),
        };

        let current_offset = state.read_offsets[self.reader_id];
        let available = state.buffer.len() - current_offset;

        // 优先从缓存读取数据
        if available > 0 {
            let take = std::cmp::min(available, buf.remaining());
            let data: Vec<u8> = state.buffer.range(current_offset..current_offset + take).cloned().collect();
            buf.put_slice(&data);
            state.read_offsets[self.reader_id] += take;
            return Poll::Ready(Ok(()));
        }

        // 缓存为空，从底层流读取数据填充缓存
        let mut temp_buf = vec![0; 4096]; // 初始缓冲区大小可按需调整
        match std::pin::Pin::new(&mut state.inner).poll_read(cx, &mut temp_buf.as_mut_slice()) {
            Poll::Ready(Ok(n)) => {
                if n == 0 {
                    return Poll::Ready(Ok(())); // 到达流末尾
                }
                state.buffer.extend(&temp_buf[0..n]);
                // 读取刚填充的缓存数据
                let take = std::cmp::min(n, buf.remaining());
                let data: Vec<u8> = state.buffer.range(current_offset..current_offset + take).cloned().collect();
                buf.put_slice(&data);
                state.read_offsets[self.reader_id] += take;
                Poll::Ready(Ok(()))
            }
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    }
}

3. 实现fork函数 #

/// 将一个AsyncReader拆分为两个独立的读取器
pub fn fork<R>(reader: R) -> (ForkedReader<R>, ForkedReader<R>)
where
    R: AsyncRead + Unpin + Send + 'static,
{
    let state = Arc::new(Mutex::new(SharedState::new(reader)));
    (
        ForkedReader {
            state: state.clone(),
            reader_id: 0,
        },
        ForkedReader {
            state,
            reader_id: 1,
        },
    )
}

4. 业务逻辑适配 #

use tokio_rustls::{server::TlsStream, LazyConfigAcceptor};

async fn detect_server_name<R>(mut reader: R) -> Option<String>
where
    R: AsyncRead + Unpin,
{
    let acceptor = LazyConfigAcceptor::new(TlsStream::new(reader));
    match acceptor.await {
        Ok(acceptor) => acceptor.client_hello().server_name().map(|s| s.to_string()),
        Err(_) => None,
    }
}

// 代理连接处理示例
async fn handle_connection(stream: tokio::net::TcpStream, proxy_map: &ProxyMap) -> std::io::Result<()> {
    let (reader, reader_copy) = fork(stream);
    
    // 检测ClientHello中的服务器名称
    let server_name = detect_server_name(reader_copy).await.unwrap_or_default();
    
    // 根据服务器名称选择代理并转发流量
    if let Some(proxy) = proxy_map.get(&server_name) {
        proxy.forward(reader).await?;
    }

    Ok(())
}

关键说明 #

• 动态扩容：VecDeque会自动按需扩容，无需手动管理缓存大小，适配不同读取量需求
• 线程安全：用Arc<Mutex>保证共享状态在多线程异步环境下的安全访问
• 无重复读取：底层流的数据仅被读取一次，存入缓存后供两个读取器按需获取，避免资源浪费
• 异步友好：完全适配tokio的异步IO模型，不会出现不必要的阻塞

内容的提问来源于stack exchange，提问作者Starua