1 files changed, 53 insertions, 0 deletions

diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index ac650364e..76602edc2 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -103,6 +103,32 @@ impl<'a> UdpSocket<'a> {
         })
     }
 
+    /// Wait until the socket becomes readable.
+    ///
+    /// A socket is readable when a packet has been received, or when there are queued packets in
+    /// the buffer.
+    pub async fn wait_recv_ready(&self) {
+        poll_fn(move |cx| self.poll_recv_ready(cx)).await
+    }
+
+    /// Wait until a datagram can be read.
+    ///
+    /// When no datagram is readable, this method will return `Poll::Pending` and
+    /// register the current task to be notified when a datagram is received.
+    ///
+    /// When a datagram is received, this method will return `Poll::Ready`.
+    pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_recv() {
+                Poll::Ready(())
+            } else {
+                // socket buffer is empty wait until at least one byte has arrived
+                s.register_recv_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     /// Receive a datagram.
     ///
     /// This method will wait until a datagram is received.
@@ -164,6 +190,33 @@ impl<'a> UdpSocket<'a> {
         .await
     }
 
+    /// Wait until the socket becomes writable.
+    ///
+    /// A socket becomes writable when there is space in the buffer, from initial memory or after
+    /// dispatching datagrams on a full buffer.
+    pub async fn wait_send_ready(&self) {
+        poll_fn(move |cx| self.poll_send_ready(cx)).await
+    }
+
+    /// Wait until a datagram can be sent.
+    ///
+    /// When no datagram can be sent (i.e. the buffer is full), this method will return
+    /// `Poll::Pending` and register the current task to be notified when
+    /// space is freed in the buffer after a datagram has been dispatched.
+    ///
+    /// When a datagram can be sent, this method will return `Poll::Ready`.
+    pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_send() {
+                Poll::Ready(())
+            } else {
+                // socket buffer is full wait until a datagram has been dispatched
+                s.register_send_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     /// Send a datagram to the specified remote endpoint.
     ///
     /// This method will wait until the datagram has been sent.