add: Add `UdpSocket::{try_send,try_send_to}` methods (#1979)

2 files changed, 85 insertions, 9 deletions

diff --git a/tokio/src/net/udp/socket.rs b/tokio/src/net/udp/socket.rs
index 97090a20..16e53773 100644
--- a/tokio/src/net/udp/socket.rs
+++ b/tokio/src/net/udp/socket.rs
@@ -116,6 +116,20 @@ impl UdpSocket {
         poll_fn(|cx| self.poll_send(cx, buf)).await
     }
 
+    /// Try to send data on the socket to the remote address to which it is
+    /// connected.
+    ///
+    /// # Returns
+    ///
+    /// If successfull, the number of bytes sent is returned. Users
+    /// should ensure that when the remote cannot receive, the
+    /// [`ErrorKind::WouldBlock`] is properly handled.
+    ///
+    /// [`ErrorKind::WouldBlock`]: std::io::error::ErrorKind::WouldBlock
+    pub fn try_send(&self, buf: &[u8]) -> io::Result<usize> {
+        self.io.get_ref().send(buf)
+    }
+
     // Poll IO functions that takes `&self` are provided for the split API.
     //
     // They are not public because (taken from the doc of `PollEvented`):
@@ -185,6 +199,21 @@ impl UdpSocket {
         }
     }
 
+    /// Try to send data on the socket to the given address.
+    ///
+    /// # Returns
+    ///
+    /// If successfull, the future resolves to the number of bytes sent.
+    ///
+    /// Users should ensure that when the remote cannot receive, the
+    /// [`ErrorKind::WouldBlock`] is properly handled. An error can also occur
+    /// if the IP version of the socket does not match that of `target`.
+    ///
+    /// [`ErrorKind::WouldBlock`]: std::io::error::ErrorKind::WouldBlock
+    pub fn try_send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
+        self.io.get_ref().send_to(buf, &target)
+    }
+
     // TODO: Public or not?
     #[doc(hidden)]
     pub fn poll_send_to(
diff --git a/tokio/tests/udp.rs b/tokio/tests/udp.rs
index 71c282a5..62a2234f 100644
--- a/tokio/tests/udp.rs
+++ b/tokio/tests/udp.rs
@@ -3,6 +3,9 @@
 
 use tokio::net::UdpSocket;
 
+const MSG: &[u8] = b"hello";
+const MSG_LEN: usize = MSG.len();
+
 #[tokio::test]
 async fn send_recv() -> std::io::Result<()> {
     let mut sender = UdpSocket::bind("127.0.0.1:0").await?;
@@ -11,13 +14,12 @@ async fn send_recv() -> std::io::Result<()> {
     sender.connect(receiver.local_addr()?).await?;
     receiver.connect(sender.local_addr()?).await?;
 
-    let message = b"hello!";
-    sender.send(message).await?;
+    sender.send(MSG).await?;
 
     let mut recv_buf = [0u8; 32];
     let len = receiver.recv(&mut recv_buf[..]).await?;
 
-    assert_eq!(&recv_buf[..len], message);
+    assert_eq!(&recv_buf[..len], MSG);
     Ok(())
 }
 
@@ -26,14 +28,13 @@ async fn send_to_recv_from() -> std::io::Result<()> {
     let mut sender = UdpSocket::bind("127.0.0.1:0").await?;
     let mut receiver = UdpSocket::bind("127.0.0.1:0").await?;
 
-    let message = b"hello!";
     let receiver_addr = receiver.local_addr()?;
-    sender.send_to(message, &receiver_addr).await?;
+    sender.send_to(MSG, &receiver_addr).await?;
 
     let mut recv_buf = [0u8; 32];
     let (len, addr) = receiver.recv_from(&mut recv_buf[..]).await?;
 
-    assert_eq!(&recv_buf[..len], message);
+    assert_eq!(&recv_buf[..len], MSG);
     assert_eq!(addr, sender.local_addr()?);
     Ok(())
 }
@@ -43,14 +44,13 @@ async fn split() -> std::io::Result<()> {
     let socket = UdpSocket::bind("127.0.0.1:0").await?;
     let (mut r, mut s) = socket.split();
 
-    let msg = b"hello";
     let addr = s.as_ref().local_addr()?;
     tokio::spawn(async move {
-        s.send_to(msg, &addr).await.unwrap();
+        s.send_to(MSG, &addr).await.unwrap();
     });
     let mut recv_buf = [0u8; 32];
     let (len, _) = r.recv_from(&mut recv_buf[..]).await?;
-    assert_eq!(&recv_buf[..len], msg);
+    assert_eq!(&recv_buf[..len], MSG);
     Ok(())
 }
 
@@ -71,3 +71,50 @@ async fn reunite_error() -> std::io::Result<()> {
     assert!(s.reunite(r1).is_err());
     Ok(())
 }
+
+// # Note
+//
+// This test is purposely written such that each time `sender` sends data on
+// the socket, `receiver` awaits the data. On Unix, it would be okay waiting
+// until the end of the test to receive all the data. On Windows, this would
+// **not** be okay because it's resources are completion based (via IOCP).
+// If data is sent and not yet received, attempting to send more data will
+// result in `ErrorKind::WouldBlock` until the first operation completes.
+#[tokio::test]
+async fn try_send_spawn() {
+    const MSG2: &[u8] = b"world!";
+    const MSG2_LEN: usize = MSG2.len();
+
+    let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+    let mut receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+
+    receiver
+        .connect(sender.local_addr().unwrap())
+        .await
+        .unwrap();
+
+    let sent = &sender
+        .try_send_to(MSG, receiver.local_addr().unwrap())
+        .unwrap();
+    assert_eq!(sent, &MSG_LEN);
+    let mut buf = [0u8; 32];
+    let mut received = receiver.recv(&mut buf[..]).await.unwrap();
+
+    sender
+        .connect(receiver.local_addr().unwrap())
+        .await
+        .unwrap();
+    let sent = &sender.try_send(MSG2).unwrap();
+    assert_eq!(sent, &MSG2_LEN);
+    received += receiver.recv(&mut buf[..]).await.unwrap();
+
+    std::thread::spawn(move || {
+        let sent = &sender.try_send(MSG).unwrap();
+        assert_eq!(sent, &MSG_LEN);
+    })
+    .join()
+    .unwrap();
+    received += receiver.recv(&mut buf[..]).await.unwrap();
+
+    assert_eq!(received, MSG_LEN * 2 + MSG2_LEN);
+}