net: Add flush for UDP and Raw sockets.

2 files changed, 34 insertions, 0 deletions

diff --git a/embassy-net/src/raw.rs b/embassy-net/src/raw.rs
index 1098dc208..ace325a46 100644
--- a/embassy-net/src/raw.rs
+++ b/embassy-net/src/raw.rs
@@ -108,6 +108,23 @@ impl<'a> RawSocket<'a> {
             }
         })
     }
+
+    /// Flush the socket.
+    ///
+    /// This method will wait until the socket is flushed.
+    pub async fn flush(&mut self) {
+        poll_fn(move |cx| {
+            self.with_mut(|s, _| {
+                if s.send_queue() == 0 {
+                    Poll::Ready(())
+                } else {
+                    s.register_send_waker(cx.waker());
+                    Poll::Pending
+                }
+            })
+        })
+        .await
+    }
 }
 
 impl Drop for RawSocket<'_> {
diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index 3eb6e2f83..ac650364e 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -241,6 +241,23 @@ impl<'a> UdpSocket<'a> {
         .await
     }
 
+    /// Flush the socket.
+    ///
+    /// This method will wait until the socket is flushed.
+    pub async fn flush(&mut self) {
+        poll_fn(move |cx| {
+            self.with_mut(|s, _| {
+                if s.send_queue() == 0 {
+                    Poll::Ready(())
+                } else {
+                    s.register_send_waker(cx.waker());
+                    Poll::Pending
+                }
+            })
+        })
+        .await
+    }
+
     /// Returns the local endpoint of the socket.
     pub fn endpoint(&self) -> IpListenEndpoint {
         self.with(|s, _| s.endpoint())