diff --git a/embassy-net/src/tcp.rs b/embassy-net/src/tcp.rs index 043062e06..a00ced8f4 100644 --- a/embassy-net/src/tcp.rs +++ b/embassy-net/src/tcp.rs @@ -10,7 +10,7 @@ use core::future::poll_fn; use core::mem; -use core::task::Poll; +use core::task::{Context, Poll}; use embassy_time::Duration; use smoltcp::iface::{Interface, SocketHandle}; @@ -274,6 +274,16 @@ impl<'a> TcpSocket<'a> { .await } + /// Wait until the socket becomes readable. + /// + /// A socket becomes readable when the receive half of the full-duplex connection is open + /// (see [may_recv](#method.may_recv)), and there is some pending data in the receive buffer. + /// + /// This is the equivalent of [read](#method.read), without buffering any data. + pub async fn wait_read_ready(&self) { + poll_fn(move |cx| self.io.poll_read_ready(cx)).await + } + /// Read data from the socket. /// /// Returns how many bytes were read, or an error. If no data is available, it waits @@ -285,6 +295,16 @@ impl<'a> TcpSocket<'a> { self.io.read(buf).await } + /// Wait until the socket becomes writable. + /// + /// A socket becomes writable when the transmit half of the full-duplex connection is open + /// (see [may_send](#method.may_send)), and the transmit buffer is not full. + /// + /// This is the equivalent of [write](#method.write), without sending any data. + pub async fn wait_write_ready(&self) { + poll_fn(move |cx| self.io.poll_write_ready(cx)).await + } + /// Write data to the socket. /// /// Returns how many bytes were written, or an error. If the socket is not ready to @@ -376,11 +396,25 @@ impl<'a> TcpSocket<'a> { self.io.with_mut(|s, _| s.abort()) } - /// Get whether the socket is ready to send data, i.e. whether there is space in the send buffer. + /// Return whether the transmit half of the full-duplex connection is open. + /// + /// This function returns true if it's possible to send data and have it arrive + /// to the remote endpoint. However, it does not make any guarantees about the state + /// of the transmit buffer, and even if it returns true, [write](#method.write) may + /// not be able to enqueue any octets. + /// + /// In terms of the TCP state machine, the socket must be in the `ESTABLISHED` or + /// `CLOSE-WAIT` state. pub fn may_send(&self) -> bool { self.io.with(|s, _| s.may_send()) } + /// Check whether the transmit half of the full-duplex connection is open + /// (see [may_send](#method.may_send)), and the transmit buffer is not full. + pub fn can_send(&self) -> bool { + self.io.with(|s, _| s.can_send()) + } + /// return whether the receive half of the full-duplex connection is open. /// This function returns true if it’s possible to receive data from the remote endpoint. /// It will return true while there is data in the receive buffer, and if there isn’t, @@ -427,7 +461,7 @@ impl<'d> TcpIo<'d> { }) } - fn with_mut(&mut self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R { + fn with_mut(&self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R { self.stack.with_mut(|i| { let socket = i.sockets.get_mut::(self.handle); let res = f(socket, &mut i.iface); @@ -436,6 +470,17 @@ impl<'d> TcpIo<'d> { }) } + fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<()> { + self.with_mut(|s, _| { + if s.can_recv() { + Poll::Ready(()) + } else { + s.register_recv_waker(cx.waker()); + Poll::Pending + } + }) + } + async fn read(&mut self, buf: &mut [u8]) -> Result { poll_fn(move |cx| { // CAUTION: smoltcp semantics around EOF are different to what you'd expect @@ -464,6 +509,17 @@ impl<'d> TcpIo<'d> { .await } + fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<()> { + self.with_mut(|s, _| { + if s.can_send() { + Poll::Ready(()) + } else { + s.register_send_waker(cx.waker()); + Poll::Pending + } + }) + } + async fn write(&mut self, buf: &[u8]) -> Result { poll_fn(move |cx| { self.with_mut(|s, _| match s.send_slice(buf) { 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.