Rust中如何安全关闭自定义读写VirtualStream并优化实现？

我实现了一个可读可写的VirtualStream，在多线程环境中多次用try_clone克隆它。现在需要解决两个问题：

1. 如何安全关闭该VirtualStream，以正常响应HTTP请求？
2. 有没有更优的VirtualStream实现方案？

以下是我的实现代码及测试情况：

use std::io;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::sync::Arc;
use std::thread;

use std::borrow::BorrowMut;
use std::sync::mpsc::{channel, Receiver, RecvError, Sender};

pub struct VirtualStream {
    tx1: Sender<Vec<u8>>,
    rp1: Arc<Receiver<Vec<u8>>>,
    tx2: Sender<Vec<u8>>,
    rp2: Arc<Receiver<Vec<u8>>>,
}

unsafe impl Sync for VirtualStream {}

unsafe impl Send for VirtualStream {}

impl VirtualStream {
    pub fn new() -> Self {
        let (tx1, rp1): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
        let (tx2, rp2): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();

        VirtualStream {
            rp1: Arc::new(rp1),
            tx1,
            rp2: Arc::new(rp2),
            tx2,
        }
    }

    pub fn produce(&self, buf: &[u8]) {
        // provide data
        let _ = self.tx1.send(buf.to_vec());
    }

    pub fn accept(&self) -> Result<Vec<u8>, RecvError> {
        // consume  data
        self.rp2.recv().map_err(|e| e.into())
    }
    pub fn shutdown(&mut self) -> std::io::Result<()> {
        // close stream
        println!("shuwdown....");
        Ok(())
    }
    pub fn try_clone(&self) -> Option<Self> {
        let tx1 = self.tx1.clone();
        let tx2 = self.tx2.clone();
        let rp1 = self.rp1.clone();
        let rp2 = self.rp2.clone();
        let cloned = VirtualStream { rp1, tx1, rp2, tx2 };
        Some(cloned)
    }
}

impl Read for VirtualStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.rp1.recv() {
            Ok(data) => {
                let len = data.len().min(buf.len());
                buf[0..len].copy_from_slice(&data[0..len]);
                println!("read ok");
                Ok(len)
            }
            Err(_) => {
                println!("read error");
                Err(io::ErrorKind::WouldBlock.into())
            }
        }
    }
}

impl Write for VirtualStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let writer = self.tx2.borrow_mut();
        writer.send(buf.to_vec()).unwrap();
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        // no cache
        Ok(())
    }
}

fn main() {
    let listener = TcpListener::bind("127.0.0.1:4000").unwrap();
    println!("Forward 127.0.0.1:4000 to 127.0.0.1:8000");

    for stream in listener.incoming() {
        if let Ok(mut client) = stream {
            thread::spawn(move || {
                let mut client2 = client.try_clone().unwrap();
                let mut client3 = client.try_clone().unwrap();
                let mut v_stream0 = VirtualStream::new();
                // let virtual_stream1 = v_stream0.try_clone().unwrap();
                let mut virtual_stream2 = v_stream0.try_clone().unwrap();
                let mut virtual_stream3 = v_stream0.try_clone().unwrap();

                thread::spawn(move || loop {
                    println!("listen v_stream0 data");
                    match v_stream0.accept() {
                        Ok(data) => {
                            let string_result = std::str::from_utf8(&data).unwrap();
                            println!("http request packet:{:?}", string_result);
                            v_stream0.produce("HTTP/1.1 200 OK\n\nhello world123".as_bytes());
                            v_stream0.shutdown().expect("Failed to shutdown");
                            println!("after shutdown called");
                        }
                        Err(err) => {
                            // closed stream, then shutdown client
                            print!("close detect...{:?}", err);
                            client
                                .shutdown(std::net::Shutdown::Write)
                                .expect("falied close stream");
                        }
                    }
                });
                thread::spawn(move || {
                    let _ = io::copy(&mut virtual_stream2, &mut client2);
                    println!("copy virtual stream to client complete or error...1");
                });
                let _ = io::copy(&mut client3, &mut virtual_stream3);
                println!("copy client to virtual stream complete or error...2");
            });
        }
    }
}

测试情况 #

Telnet测试 #

☁  ~  telnet 127.0.0.1 4000
Trying 127.0.0.1...
Connected to localhost.
Escape character is '^]'.
GET / HTTP/1.1
HTTP/1.1 200 OK

hello world123

Curl测试 #

☁  examples [master] ⚡  curl -vs http://127.0.0.1:4000/
*   Trying 127.0.0.1...
* TCP_NODELAY set
* Connected to 127.0.0.1 (127.0.0.1) port 4000 (#0)
> GET / HTTP/1.1
> Host: 127.0.0.1:4000
> User-Agent: curl/7.64.1
> Accept: */*
>
< HTTP/1.1 200 OK
* no chunk, no close, no size. Assume close to signal end
<

一、安全关闭VirtualStream的实现 #

当前shutdown方法为空实现，无法通知所有克隆实例终止读写操作，导致上层io::copy无法正常结束，curl等工具会挂起。需做以下修改：

核心修改点 #

1. 添加共享关闭状态：用Arc<AtomicBool>记录流的关闭状态，所有克隆实例共享该状态，确保状态一致。
2. 关闭时终止通道通信：mpsc通道的Sender全部被drop后，Receiver会返回断开错误，因此shutdown时需drop当前持有的Sender，触发所有Receiver的终止信号。
3. 修正IO方法的错误返回：检测到关闭状态或通道断开时，返回明确的连接错误（如ConnectionReset/BrokenPipe），让io::copy能正确终止。

修改后的核心代码 #

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

pub struct VirtualStream {
    tx1: Option<Sender<Vec<u8>>>,
    rp1: Arc<Receiver<Vec<u8>>>,
    tx2: Option<Sender<Vec<u8>>>,
    rp2: Arc<Receiver<Vec<u8>>>,
    closed: Arc<AtomicBool>,
}

impl VirtualStream {
    pub fn new() -> Self {
        let (tx1, rp1) = channel();
        let (tx2, rp2) = channel();

        VirtualStream {
            rp1: Arc::new(rp1),
            tx1: Some(tx1),
            rp2: Arc::new(rp2),
            tx2: Some(tx2),
            closed: Arc::new(AtomicBool::new(false)),
        }
    }

    pub fn shutdown(&mut self) -> std::io::Result<()> {
        if self.closed.swap(true, Ordering::SeqCst) {
            return Ok(()); // 已关闭过，直接返回
        }
        // Drop Sender，触发通道断开
        self.tx1.take();
        self.tx2.take();
        println!("shutdown completed");
        Ok(())
    }

    pub fn try_clone(&self) -> Option<Self> {
        if self.closed.load(Ordering::SeqCst) {
            return None; // 已关闭的流禁止克隆
        }
        Some(VirtualStream {
            rp1: self.rp1.clone(),
            tx1: self.tx1.clone(),
            rp2: self.rp2.clone(),
            tx2: self.tx2.clone(),
            closed: self.closed.clone(),
        })
    }
}

impl Read for VirtualStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(io::Error::new(io::ErrorKind::ConnectionReset, "stream closed"));
        }

        match self.rp1.recv_timeout(std::time::Duration::from_millis(100)) {
            Ok(data) => {
                let len = data.len().min(buf.len());
                buf[0..len].copy_from_slice(&data[0..len]);
                Ok(len)
            }
            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
                Err(io::ErrorKind::WouldBlock.into())
            }
            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
                self.closed.store(true, Ordering::SeqCst);
                Err(io::Error::new(io::ErrorKind::ConnectionReset, "stream closed"))
            }
        }
    }
}

impl Write for VirtualStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "stream closed"));
        }

        match self.tx2.as_ref().unwrap().send(buf.to_vec()) {
            Ok(_) => Ok(buf.len()),
            Err(_) => {
                self.closed.store(true, Ordering::SeqCst);
                Err(io::Error::new(io::ErrorKind::BrokenPipe, "stream closed"))
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

关键说明 #

• 用AtomicBool保证多线程下关闭状态的原子性，避免竞态条件。
• recv_timeout替代阻塞式recv，符合IO操作的超时规范，防止线程永久阻塞。
• 关闭时返回明确的错误类型，让io::copy正确终止，进而关闭客户端连接，解决curl挂起问题。

二、更优的VirtualStream实现方案 #

当前基于标准库mpsc的实现存在两个明显缺陷：

1. 标准库mpsc的Receiver不是线程安全的，多线程同时调用recv会直接panic。
2. 每次发送数据都要复制Vec<u8>，性能开销较大。

推荐两种更优方案：

方案1：用crossbeam-channel替代标准库mpsc #

crossbeam-channel支持多生产者多消费者，线程安全，性能优于标准库mpsc，API更简洁：

use crossbeam_channel::{unbounded, Receiver, Sender};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

pub struct VirtualStream {
    tx_read: Option<Sender<Vec<u8>>>,
    rx_read: Arc<Receiver<Vec<u8>>>,
    tx_write: Option<Sender<Vec<u8>>>,
    rx_write: Arc<Receiver<Vec<u8>>>,
    closed: Arc<AtomicBool>,
}

impl VirtualStream {
    pub fn new() -> Self {
        let (tx_read, rx_read) = unbounded();
        let (tx_write, rx_write) = unbounded();
        VirtualStream {
            tx_read: Some(tx_read),
            rx_read: Arc::new(rx_read),
            tx_write: Some(tx_write),
            rx_write: Arc::new(rx_write),
            closed: Arc::new(AtomicBool::new(false)),
        }
    }

    // 其余方法（produce/accept/shutdown/try_clone）及Read/Write实现逻辑与之前一致，替换为crossbeam的API即可
}

方案2：环形缓冲区+零拷贝（高并发场景） #

若追求极致性能，可结合bytes库的Bytes类型和环形缓冲区实现零拷贝数据传递，避免内存复制：

use bytes::Bytes;
use ringbuf::RingBuffer;
use std::sync::{Arc, Mutex};
use std::io::{Read, Write, Error, ErrorKind};
use std::sync::atomic::{AtomicBool, Ordering};

pub struct VirtualStream {
    read_buf: Arc<Mutex<RingBuffer<Bytes>>>,
    write_buf: Arc<Mutex<RingBuffer<Bytes>>>,
    closed: Arc<AtomicBool>,
}

impl VirtualStream {
    pub fn new() -> Self {
        VirtualStream {
            read_buf: Arc::new(Mutex::new(RingBuffer::new(1024 * 16))),
            write_buf: Arc::new(Mutex::new(RingBuffer::new(1024 * 16))),
            closed: Arc::new(AtomicBool::new(false)),
        }
    }

    pub fn produce(&self, buf: &[u8]) -> io::Result<()> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(ErrorKind::BrokenPipe.into());
        }
        let mut buf_guard = self.read_buf.lock().unwrap();
        buf_guard.push(Bytes::copy_from_slice(buf)).map_err(|_| ErrorKind::WouldBlock.into())
    }

    // 其余方法及Read/Write实现需基于环形缓冲区的读写逻辑编写
}

该方案适合高并发场景，大幅减少内存拷贝开销。

内容的提问来源于stack exchange，提问作者Hai.Xu