tokio::timeout on Stream always occurs - rust

Question

I'm trying to accept a UDP message but only if it happens within 5 seconds, I have a Stream abstraction built with both manually implementing Stream and by using the combinators in the futures library. Each way, after the recv_from future resolves, the duration will expire and the stream will return an Err(Elapsed(())). This is not the expected behaviour, if a value is returned, no error is supposed to be returned.
The expected behaviour is that the stream will resolve either the timeout or the Vec, but not one, then the other 5 seconds later.
use futures::{pin_mut, ready, stream::unfold, FutureExt};
use tokio::{
net::{udp, UdpSocket},
stream::{Stream, StreamExt},
time::{self, Duration},
};
use std::{
io,
net::SocketAddr,
pin::Pin,
task::{Context, Poll},
};
#[derive(Debug)]
pub(crate) struct UdpStream {
stream: udp::RecvHalf,
}
impl UdpStream {
fn new(stream: udp::RecvHalf) -> Self {
Self { stream }
}
fn stream(self) -> impl Stream<Item = io::Result<(Vec<u8>, SocketAddr)>> {
unfold(self.stream, |mut stream| async move {
let mut buf = [0; 4096];
match time::timeout(Duration::from_secs(5), stream.recv_from(&mut buf)).await {
Ok(Ok((len, src))) => {
Some((Ok((buf.iter().take(len).cloned().collect(), src)), stream))
}
e => {
println!("{:?}", e);
None
}
}
})
}
}
impl Stream for UdpStream {
type Item = io::Result<(Vec<u8>, SocketAddr)>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let socket = &mut self.stream;
pin_mut!(socket);
let mut buf = [0u8; 4096];
let (len, src) = ready!(Box::pin(socket.recv_from(&mut buf)).poll_unpin(cx))?;
Poll::Ready(Some(Ok((buf.iter().take(len).cloned().collect(), src))))
}
}
async fn listen_udp(addr: SocketAddr) -> io::Result<()> {
let udp = UdpSocket::bind(addr).await?;
let (mut udp_recv, mut udp_send) = udp.split();
let mut msg_stream = Box::pin(UdpStream::new(udp_recv).stream());
// use the manually implemented stream with this:
// let mut msg_stream = UdpStream::new(udp_recv).timeout(Duration::from_secs(5));
while let Some(msg) = msg_stream.next().await {
match msg {
Ok((buf, src)) => {
udp_send.send_to(&buf, &src).await?;
println!("Message recv: {:?}", buf);
}
Err(e) => {
eprintln!("timed out: {:?}", e);
}
}
}
Ok(())
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
listen_udp("127.0.0.1:9953".parse()?).await?;
Ok(())
}
You can try by running this code and making udp requests with echo "foo" | nc 127.0.0.1 9953 -u or with dig
cargo.toml
[package]
name = "udp_test"
version = "0.1.0"
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
tokio = { version = "0.2", features = ["full"] }
futures = "0.3"

Answer 1

After your stream returns the first (and the only) element, it goes back to wait for the next; it never ends until the timeout.
Basically a stream abstraction is not necessary here. A future wrapped in a timeout will do:
use std::{io, net::SocketAddr};
use tokio::{
net::UdpSocket,
time::{self, Duration},
};
async fn listen_udp(addr: SocketAddr) -> io::Result<()> {
let mut udp = UdpSocket::bind(addr).await?;
let mut buf = [0; 4096];
match time::timeout(Duration::from_secs(5), udp.recv_from(&mut buf)).await? {
Ok((count, src)) => {
udp.send_to(&buf[..count], &src).await?;
println!("Message recv: {:?}", &buf[..count]);
}
Err(e) => {
eprintln!("timed out: {:?}", e);
}
}
Ok(())
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
listen_udp("127.0.0.1:9953".parse()?).await?;
Ok(())
}