diff --git a/iroh/src/magicsock/relay_actor.rs b/iroh/src/magicsock/relay_actor.rs index f624d2804a..28d2331901 100644 --- a/iroh/src/magicsock/relay_actor.rs +++ b/iroh/src/magicsock/relay_actor.rs @@ -182,7 +182,7 @@ impl ActiveRelayActor { // When this future has an inner, it is a future which is currently sending // something to the relay server. Nothing else can be sent to the relay server at // the same time. - let mut relay_send_fut = MaybeFuture::none(); + let mut relay_send_fut = std::pin::pin!(MaybeFuture::none()); loop { // If a read error occurred on the connection it might have been lost. But we @@ -204,7 +204,7 @@ impl ActiveRelayActor { } // Only poll relay_send_fut if it is sending to the relay. _ = &mut relay_send_fut, if relay_send_fut.is_some() => { - relay_send_fut = MaybeFuture::none(); + relay_send_fut.as_mut().set_none(); } // Only poll for new datagrams if relay_send_fut is not busy. Some(msg) = self.relay_datagrams_send.recv(), if relay_send_fut.is_none() => { @@ -212,7 +212,7 @@ impl ActiveRelayActor { let fut = async move { relay_client.send(msg.node_id, msg.packet).await }; - relay_send_fut = MaybeFuture::with_future(Box::pin(fut)); + relay_send_fut.as_mut().set_future(fut); self.last_write = Instant::now(); } diff --git a/iroh/src/util.rs b/iroh/src/util.rs index ee777da8fb..4c80f4c551 100644 --- a/iroh/src/util.rs +++ b/iroh/src/util.rs @@ -6,54 +6,58 @@ use std::{ task::{Context, Poll}, }; +use pin_project::pin_project; + /// A future which may not be present. /// /// This is a single type which may optionally contain a future. If there is no inner /// future polling will always return [`Poll::Pending`]. /// /// The [`Default`] impl will create a [`MaybeFuture`] without an inner. -#[derive(Debug)] -pub(crate) struct MaybeFuture { +#[derive(Default, Debug)] +#[pin_project(project = MaybeFutureProj, project_replace = MaybeFutureProjReplace)] +pub(crate) enum MaybeFuture { /// Future to be polled. - pub inner: Option, + Some(#[pin] T), + #[default] + None, } impl MaybeFuture { /// Creates a [`MaybeFuture`] without an inner future. pub(crate) fn none() -> Self { - Self { inner: None } + Self::default() + } + + /// Clears the value + pub(crate) fn set_none(mut self: Pin<&mut Self>) { + self.as_mut().project_replace(Self::None); } - /// Creates a [`MaybeFuture`] with an inner future. - pub(crate) fn with_future(fut: T) -> Self { - Self { inner: Some(fut) } + /// Sets a new future. + pub(crate) fn set_future(mut self: Pin<&mut Self>, fut: T) { + self.as_mut().project_replace(Self::Some(fut)); } /// Returns `true` if the inner is empty. pub(crate) fn is_none(&self) -> bool { - self.inner.is_none() + matches!(self, Self::None) } /// Returns `true` if the inner contains a future. pub(crate) fn is_some(&self) -> bool { - self.inner.is_some() - } -} - -// NOTE: explicit implementation to bypass derive unnecessary bounds -impl Default for MaybeFuture { - fn default() -> Self { - Self::none() + matches!(self, Self::Some(_)) } } -impl Future for MaybeFuture { +impl Future for MaybeFuture { type Output = T::Output; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match self.inner { - Some(ref mut t) => Pin::new(t).poll(cx), - None => Poll::Pending, + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + match this { + MaybeFutureProj::Some(t) => t.poll(cx), + MaybeFutureProj::None => Poll::Pending, } } }