net: add UdpSocket::try_peek_from(), fix docs on poll_peek_from

tokio/src/net/udp.rs

@@ -1364,7 +1364,7 @@ impl UdpSocket {
     }
 
     /// Receives data from the socket, without removing it from the input queue.
-    /// On success, returns the number of bytes read.
+    /// On success, returns the sending address of the datagram.
     ///
     /// # Notes
     ///
@@ -1418,6 +1418,33 @@ impl UdpSocket {
         Poll::Ready(Ok(addr))
     }
 
+    /// Tries to receive data on the socket without removing it from the input queue.
+    /// On success, returns the number of bytes read and the sending address of the
+    /// datagram.
+    ///
+    /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is
+    /// returned. This function is usually paired with `readable()`.
+    ///
+    /// # Notes
+    ///
+    /// On Windows, if the data is larger than the buffer specified, the buffer
+    /// is filled with the first part of the data, and peek returns the error
+    /// WSAEMSGSIZE(10040). The excess data is lost.
+    /// Make sure to always use a sufficiently large buffer to hold the
+    /// maximum UDP packet size, which can be up to 65536 bytes in size.
+    ///
+    /// MacOS will return an error if you pass a zero-sized buffer.
+    ///
+    /// If you're merely interested in learning the sender of the data at the head of the queue,
+    /// try [`try_peek_sender`].
+    ///
+    /// [`try_peek_sender`]: method@Self::try_peek_sender
+    pub fn try_peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        self.io
+            .registration()
+            .try_io(Interest::READABLE, || self.io.peek_from(buf))
+    }
+
     /// Retrieve the sender of the data at the head of the input queue, waiting if empty.
     ///
     /// This is equivalent to calling [`peek_from`] with a zero-sized buffer,

tokio/tests/udp.rs

@@ -106,6 +106,45 @@ async fn send_to_peek_from() -> std::io::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn send_to_try_peek_from() -> std::io::Result<()> {
+    let sender = UdpSocket::bind("127.0.0.1:0").await?;
+    let receiver = UdpSocket::bind("127.0.0.1:0").await?;
+
+    let receiver_addr = receiver.local_addr()?;
+    poll_fn(|cx| sender.poll_send_to(cx, MSG, receiver_addr)).await?;
+
+    // peek
+    let mut recv_buf = [0u8; 32];
+
+    loop {
+        match receiver.try_peek_from(&mut recv_buf) {
+            Ok((n, addr)) => {
+                assert_eq!(&recv_buf[..n], MSG);
+                assert_eq!(addr, sender.local_addr()?);
+                break;
+            }
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                receiver.readable().await?;
+            }
+            Err(e) => return Err(e),
+        }
+    }
+
+    // peek
+    let mut recv_buf = [0u8; 32];
+    let (n, addr) = receiver.peek_from(&mut recv_buf).await?;
+    assert_eq!(&recv_buf[..n], MSG);
+    assert_eq!(addr, sender.local_addr()?);
+
+    let mut recv_buf = [0u8; 32];
+    let (n, addr) = receiver.recv_from(&mut recv_buf).await?;
+    assert_eq!(&recv_buf[..n], MSG);
+    assert_eq!(addr, sender.local_addr()?);
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn send_to_peek_from_poll() -> std::io::Result<()> {
     let sender = UdpSocket::bind("127.0.0.1:0").await?;