Add poll functions on UdpSocket.

This commit is contained in:
Roy Buitenhuis 2023-07-12 11:32:02 +02:00
parent c6e2f4a90b
commit f54e1cea90

View File

@ -102,37 +102,42 @@ impl<'a> UdpSocket<'a> {
/// ///
/// Returns the number of bytes received and the remote endpoint. /// Returns the number of bytes received and the remote endpoint.
pub async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, IpEndpoint), Error> { pub async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, IpEndpoint), Error> {
poll_fn(move |cx| { poll_fn(move |cx| self.poll_recv_from(buf, cx)).await
self.with_mut(|s, _| match s.recv_slice(buf) { }
Ok((n, meta)) => Poll::Ready(Ok((n, meta.endpoint))),
// No data ready pub fn poll_recv_from(&self, buf: &mut[u8], cx: Context) -> Poll<Result<(usize, IpEndpoint), Error>> {
Err(udp::RecvError::Exhausted) => { self.with_mut(|s, _| match s.recv_slice(buf) {
s.register_recv_waker(cx.waker()); Ok((n, meta)) => Poll::Ready(Ok((n, meta.endpoint))),
Poll::Pending // No data ready
} Err(udp::RecvError::Exhausted) => {
}) s.register_recv_waker(cx.waker());
Poll::Pending
}
}) })
.await
} }
/// Send a datagram to the specified remote endpoint. /// Send a datagram to the specified remote endpoint.
pub async fn send_to<T>(&self, buf: &[u8], remote_endpoint: T) -> Result<(), Error> pub async fn send_to<T>(&self, buf: &[u8], remote_endpoint: T) -> Result<(), Error>
where
T: Into<IpEndpoint>,
{
poll_fn(move |cx| self.poll_send_to(buf, remote_endpoint, cx)).await
}
pub fn poll_send_to<T>(&self, buf: &[u8], remote_endpoint: T, cx: Context) -> Poll<Result<(), Error>>
where where
T: Into<IpEndpoint>, T: Into<IpEndpoint>,
{ {
let remote_endpoint = remote_endpoint.into(); let remote_endpoint = remote_endpoint.into();
poll_fn(move |cx| { self.with_mut(|s, _| match s.send_slice(buf, remote_endpoint) {
self.with_mut(|s, _| match s.send_slice(buf, remote_endpoint) { // Entire datagram has been sent
// Entire datagram has been sent Ok(()) => Poll::Ready(Ok(())),
Ok(()) => Poll::Ready(Ok(())), Err(udp::SendError::BufferFull) => {
Err(udp::SendError::BufferFull) => { s.register_send_waker(cx.waker());
s.register_send_waker(cx.waker()); Poll::Pending
Poll::Pending }
} Err(udp::SendError::Unaddressable) => Poll::Ready(Err(Error::NoRoute)),
Err(udp::SendError::Unaddressable) => Poll::Ready(Err(Error::NoRoute)),
})
}) })
.await
} }
/// Returns the local endpoint of the socket. /// Returns the local endpoint of the socket.