Skip to content

Commit

Permalink
feat(embassy-net): Implement wait_read_ready() + `wait_write_ready(…
Browse files Browse the repository at this point in the history
…)` for TcpSocket
  • Loading branch information
AnthonyGrondin committed Sep 24, 2024
1 parent 143be71 commit 45d32c5
Showing 1 changed file with 44 additions and 2 deletions.
46 changes: 44 additions & 2 deletions embassy-net/src/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -441,7 +461,7 @@ impl<'d> TcpIo<'d> {
})
}

fn with_mut<R>(&mut self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
fn with_mut<R>(&self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
self.stack.with_mut(|i| {
let socket = i.sockets.get_mut::<tcp::Socket>(self.handle);
let res = f(socket, &mut i.iface);
Expand All @@ -450,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<usize, Error> {
poll_fn(move |cx| {
// CAUTION: smoltcp semantics around EOF are different to what you'd expect
Expand Down Expand Up @@ -478,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<usize, Error> {
poll_fn(move |cx| {
self.with_mut(|s, _| match s.send_slice(buf) {
Expand Down

0 comments on commit 45d32c5

Please sign in to comment.