diff --git a/src/channels.rs b/src/channels.rs index ed8904d..76dee0d 100644 --- a/src/channels.rs +++ b/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> + 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. diff --git a/src/mqueue.rs b/src/mqueue.rs index e732491..1bd912c 100644 --- a/src/mqueue.rs +++ b/src/mqueue.rs @@ -94,7 +94,6 @@ 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(())) => {} @@ -102,7 +101,6 @@ impl MessageIo { } 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>>> { - 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 >::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 >::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, + } } } } diff --git a/src/protocol.rs b/src/protocol.rs index 74c3d2e..38743ab 100644 --- a/src/protocol.rs +++ b/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> + 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); 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> + 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> + use<> { + self.send(Command::Open(key)) } /// Close a protocol channel.