net: Add flush for UDP and Raw sockets.

1 files changed, 17 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<'_> {