Rewrite async methods to return owned futures

src/channels.rs

@@ -105,7 +105,7 @@ impl Channel {
     }
 
     /// Send a batch of messages over the channel.
-    pub async fn send_batch(&self, messages: &[Message]) -> Result<()> {
+    pub fn send_batch(&self, messages: &[Message]) -> impl Future<Output = Result<()>> + use<> {
         // In javascript this is cork()/uncork(), e.g.:
         //
         // https://github.com/holepunchto/hypercore/blob/c338b9aaa4442d35bc9d283d2c242b86a46de6d4/lib/replicator.js#L402-L418
@@ -116,22 +116,30 @@ impl Channel {
         // https://github.com/holepunchto/protomux/blob/d3d6f8f55e52c2fbe5cd56f5d067ac43ca13c27d/index.js#L368-L389
         //
         // Batching messages across channels like protomux is capable of doing is not (yet) implemented.
-        if self.closed() {
-            return Err(Error::new(
-                ErrorKind::ConnectionAborted,
-                "Channel is closed",
-            ));
-        }
 
-        let messages = messages
-            .iter()
-            .map(|message| ChannelMessage::new(self.local_id as u64, message.clone()))
-            .collect();
+        let closed = self.closed();
 
-        self.outbound_tx
-            .send(messages)
-            .await
-            .map_err(map_channel_err)
+        // we do this to avoid having the future capture &[Messages]
+        let messages = if !closed {
+            messages
+                .iter()
+                .map(|message| ChannelMessage::new(self.local_id as u64, message.clone()))
+                .collect()
+        } else {
+            vec![]
+        };
+
+        let outbound_tx = self.outbound_tx.clone();
+        async move {
+            if closed {
+                return Err(Error::new(
+                    ErrorKind::ConnectionAborted,
+                    "Channel is closed",
+                ));
+            }
+
+            outbound_tx.send(messages).await.map_err(map_channel_err)
+        }
     }
 
     /// Take the receiving part out of the channel.

src/mqueue.rs

@@ -94,15 +94,13 @@ impl MessageIo {
             match Sink::poll_flush(Pin::new(&mut self.stream), cx) {
                 Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                 Poll::Pending => {
-                    cx.waker().wake_by_ref();
                     return Poll::Pending;
                 }
                 Poll::Ready(Ok(())) => {}
             }
         }
 
         if pending {
-            cx.waker().wake_by_ref();
             Poll::Pending
         } else {
             Poll::Ready(Ok(()))
@@ -115,24 +113,29 @@ impl MessageIo {
         &mut self,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Vec<ChannelMessage>>>> {
-        match Pin::new(&mut self.stream).poll_next(cx) {
-            Poll::Ready(Some(event)) => match event {
-                hypercore_handshake::CipherEvent::HandshakePayload(_x) => Poll::Pending,
-                hypercore_handshake::CipherEvent::Message(msg) => {
-                    match <Vec<ChannelMessage>>::decode(&msg) {
-                        Ok((messages, _rest)) => {
-                            for m in messages.iter() {
-                                trace!("RX ChannelMessage::{m}");
+        loop {
+            match Pin::new(&mut self.stream).poll_next(cx) {
+                Poll::Ready(Some(event)) => match event {
+                    // Skip handshake payloads: loop so the next poll registers a waker.
+                    hypercore_handshake::CipherEvent::HandshakePayload(_x) => {}
+                    hypercore_handshake::CipherEvent::Message(msg) => {
+                        return match <Vec<ChannelMessage>>::decode(&msg) {
+                            Ok((messages, _rest)) => {
+                                for m in messages.iter() {
+                                    trace!("RX ChannelMessage::{m}");
+                                }
+                                Poll::Ready(Some(Ok(messages)))
                             }
-                            Poll::Ready(Some(Ok(messages)))
-                        }
-                        Err(e) => Poll::Ready(Some(Err(e.into()))),
+                            Err(e) => Poll::Ready(Some(Err(e.into()))),
+                        };
                     }
-                }
-                hypercore_handshake::CipherEvent::ErrStuff(e) => Poll::Ready(Some(Err(e))),
-            },
-            Poll::Ready(None) => Poll::Ready(None),
-            Poll::Pending => Poll::Pending,
+                    hypercore_handshake::CipherEvent::ErrStuff(e) => {
+                        return Poll::Ready(Some(Err(e)));
+                    }
+                },
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Pending => return Poll::Pending,
+            }
         }
     }
 }

src/protocol.rs

@@ -171,8 +171,8 @@ impl Protocol {
     ///
     /// Once the other side proofed that it also knows the `key`, the channel is emitted as
     /// `Event::Channel` on the protocol event stream.
-    pub async fn open(&self, key: Key) -> Result<()> {
-        self.command_tx.open(key).await
+    pub fn open(&self, key: Key) -> impl Future<Output = Result<()>> + use<> {
+        self.command_tx.open(key)
     }
 
     /// Iterator of all currently opened channels.
@@ -189,7 +189,6 @@ impl Protocol {
             return_error!(this.poll_outbound_write(cx));
             return_error!(this.poll_inbound_read(cx));
             if this.io.handshake_hash().is_none() {
-                cx.waker().wake_by_ref();
                 return Poll::Pending;
             }
         }
@@ -198,8 +197,6 @@ impl Protocol {
             if let Some(remote_pubkey) = this.io.remote_public_key() {
                 this.handshake_emitted = true;
                 return Poll::Ready(Ok(Event::Handshake(remote_pubkey)));
-            } else {
-                cx.waker().wake_by_ref();
             }
         }
 
@@ -478,14 +475,15 @@ pub struct CommandTx(Sender<Command>);
 
 impl CommandTx {
     /// Send a protocol command
-    pub async fn send(&self, command: Command) -> Result<()> {
-        self.0.send(command).await.map_err(map_channel_err)
+    pub fn send(&self, command: Command) -> impl Future<Output = Result<()>> + use<> {
+        let sender = self.0.clone();
+        async move { sender.send(command).await.map_err(map_channel_err) }
     }
     /// Open a protocol channel.
     ///
     /// The channel will be emitted on the main protocol.
-    pub async fn open(&self, key: Key) -> Result<()> {
-        self.send(Command::Open(key)).await
+    pub fn open(&self, key: Key) -> impl Future<Output = Result<()>> + use<> {
+        self.send(Command::Open(key))
     }
 
     /// Close a protocol channel.