You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

为UdpSocket实现AsyncRead:如何在poll_read中调用async函数?

问题分析与解决方案

你的代码存在两个核心问题:

  1. 生命周期不匹配:UdpSocket::recv的future持有对poll_read参数buf的引用,但buf的生命周期仅在当前poll_read调用期间,无法存入结构体长期保存。
  2. 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

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.27 01:42:57