为UdpSocket实现AsyncRead:如何在poll_read中调用async函数?
问题分析与解决方案
你的代码存在两个核心问题:
- 生命周期不匹配:
UdpSocket::recv的future持有对poll_read参数buf的引用,但buf的生命周期仅在当前poll_read调用期间,无法存入结构体长期保存。 - future复用逻辑错误:
unwrap_or不会将新创建的future存入this.fut,导致每次poll_read都会重新创建future,丢失pending状态。
另外,UDP是面向数据包的协议,而AsyncRead是字节流语义,直接映射两者会有语义冲突——UDP每次recv获取完整数据包,但AsyncRead允许分多次读取数据包内容。因此正确的实现需要在结构体内部维护缓冲区,先缓存UDP数据包,再流式输出。
正确实现代码
使用pin-project crate简化Pin的处理(避免手动unsafe操作),同时维护内部缓冲区:
首先在Cargo.toml添加依赖:
[dependencies] async-std = { version = "1.12", features = ["attributes"] } pin-project = "1.1" futures = "0.3"
然后实现代码:
use async_std::{ io::{AsyncRead, Poll, Result}, net::UdpSocket, }; use futures::Future; use pin_project::pin_project; use std::pin::Pin; use std::task::Context; use std::sync::Arc; #[pin_project] struct MyUdpWrapper { socket: Arc<UdpSocket>, // 缓存已接收但未读取的UDP数据包内容 buffer: Vec<u8>, // 缓冲区中已读取的位置 pos: usize, // 正在进行的recv操作的future #[pin] recv_fut: Option<Box<dyn Future<Output = Result<(usize, Vec<u8>)>> + '_>>, } impl MyUdpWrapper { fn new(socket: Arc<UdpSocket>) -> Self { Self { socket, buffer: Vec::new(), pos: 0, recv_fut: None, } } } impl AsyncRead for MyUdpWrapper { fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> { let this = self.as_mut().project(); // 先处理缓冲区中剩余的数据 if *this.pos < this.buffer.len() { let copy_len = buf.len().min(this.buffer.len() - *this.pos); buf[..copy_len].copy_from_slice(&this.buffer[*this.pos..*this.pos + copy_len]); *this.pos += copy_len; return Poll::Ready(Ok(copy_len)); } // 缓冲区已空,重置位置 *this.pos = 0; this.buffer.clear(); // 检查是否有正在进行的recv操作 if let Some(recv_fut) = this.recv_fut.as_mut().as_pin_mut() { match recv_fut.poll(cx) { Poll::Ready(Ok((n, data))) => { // 接收完成,将数据存入缓冲区 *this.buffer = data; *this.recv_fut = None; // 递归调用,处理新存入的数据 self.poll_read(cx, buf) } Poll::Ready(Err(e)) => { *this.recv_fut = None; Poll::Ready(Err(e)) } Poll::Pending => Poll::Pending, } } else { // 发起新的recv操作,使用足够大的缓冲区接收UDP数据包 let socket_clone = this.socket.clone(); let recv_fut = Box::new(async move { let mut data = vec![0; 65535]; // UDP最大数据包大小 let n = socket_clone.recv(&mut data).await?; data.truncate(n); Ok((n, data)) }); *this.recv_fut = Some(recv_fut); // 递归调用,poll新创建的future self.poll_read(cx, buf) } } }
更优方案:避免手动实现AsyncRead
如果不需要自定义逻辑,可直接用futures库的适配器将UDP数据包流转换为AsyncRead,代码更简洁:
use async_std::net::UdpSocket; use futures::{stream, Stream}; use futures::io::{AsyncRead, StreamReader}; use std::sync::Arc; fn udp_to_read(socket: Arc<UdpSocket>) -> impl AsyncRead { let byte_stream: impl Stream<Item = Vec<u8>> = stream::unfold(socket, |socket| async { let mut buf = vec![0; 65535]; match socket.recv(&mut buf).await { Ok(n) => { buf.truncate(n); Some((buf, socket)) } Err(_) => None, } }).flatten(); StreamReader::new(byte_stream) }
内容的提问来源于stack exchange,提问作者Ervadac
相关产品推荐
相关产品推荐

