net: add UdpSocket readiness and non-blocking ops (#3138)

Adds `ready()`, `readable()`, and `writable()` async methods for waiting
for socket readiness. Adds `try_send`, `try_send_to`, `try_recv`, and
`try_recv_from` for performing non-blocking operations on the socket.

This is the UDP equivalent of #3130.

1 files changed, 90 insertions, 0 deletions

diff --git a/tokio/tests/udp.rs b/tokio/tests/udp.rs
index 8b79cb85..291267e0 100644
--- a/tokio/tests/udp.rs
+++ b/tokio/tests/udp.rs
@@ -2,6 +2,7 @@
 #![cfg(feature = "full")]
 
 use futures::future::poll_fn;
+use std::io;
 use std::sync::Arc;
 use tokio::{io::ReadBuf, net::UdpSocket};
 
@@ -238,6 +239,8 @@ async fn try_send_spawn() {
         .await
         .unwrap();
 
+    sender.writable().await.unwrap();
+
     let sent = &sender
         .try_send_to(MSG, receiver.local_addr().unwrap())
         .unwrap();
@@ -263,3 +266,90 @@ async fn try_send_spawn() {
 
     assert_eq!(received, MSG_LEN * 2 + MSG2_LEN);
 }
+
+#[tokio::test]
+async fn try_send_recv() {
+    // Create listener
+    let server = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+
+    // Create socket pair
+    let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+
+    // Connect the two
+    client.connect(server.local_addr().unwrap()).await.unwrap();
+    server.connect(client.local_addr().unwrap()).await.unwrap();
+
+    for _ in 0..5 {
+        loop {
+            client.writable().await.unwrap();
+
+            match client.try_send(b"hello world") {
+                Ok(n) => {
+                    assert_eq!(n, 11);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+
+        loop {
+            server.readable().await.unwrap();
+
+            let mut buf = [0; 512];
+
+            match server.try_recv(&mut buf) {
+                Ok(n) => {
+                    assert_eq!(n, 11);
+                    assert_eq!(&buf[0..11], &b"hello world"[..]);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+    }
+}
+
+#[tokio::test]
+async fn try_send_to_recv_from() {
+    // Create listener
+    let server = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+    let saddr = server.local_addr().unwrap();
+
+    // Create socket pair
+    let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+    let caddr = client.local_addr().unwrap();
+
+    for _ in 0..5 {
+        loop {
+            client.writable().await.unwrap();
+
+            match client.try_send_to(b"hello world", saddr) {
+                Ok(n) => {
+                    assert_eq!(n, 11);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+
+        loop {
+            server.readable().await.unwrap();
+
+            let mut buf = [0; 512];
+
+            match server.try_recv_from(&mut buf) {
+                Ok((n, addr)) => {
+                    assert_eq!(n, 11);
+                    assert_eq!(addr, caddr);
+                    assert_eq!(&buf[0..11], &b"hello world"[..]);
+                    break;
+                }
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => panic!("{:?}", e),
+            }
+        }
+    }
+}