diff --git a/sdk/rust/nym-sdk/examples/parallel_sending_and_receiving.rs b/sdk/rust/nym-sdk/examples/parallel_sending_and_receiving.rs index e3690ce9e0..4c048e27ae 100644 --- a/sdk/rust/nym-sdk/examples/parallel_sending_and_receiving.rs +++ b/sdk/rust/nym-sdk/examples/parallel_sending_and_receiving.rs @@ -28,9 +28,7 @@ async fn main() { // receiving task tokio::spawn(async move { if let Some(received) = client.next().await { - for r in received { - println!("Received: {}", String::from_utf8_lossy(&r.message)); - } + println!("Received: {}", String::from_utf8_lossy(&received.message)); } client.disconnect().await; diff --git a/sdk/rust/nym-sdk/src/mixnet/client.rs b/sdk/rust/nym-sdk/src/mixnet/client.rs index 33185ccd88..4340b79174 100644 --- a/sdk/rust/nym-sdk/src/mixnet/client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/client.rs @@ -496,15 +496,15 @@ where let reconstructed_receiver = client_output.register_receiver()?; - Ok(MixnetClient { + Ok(MixnetClient::new( nym_address, client_input, client_output, client_state, reconstructed_receiver, - task_manager: started_client.task_manager, - packet_type: None, - }) + started_client.task_manager, + None, + )) } } diff --git a/sdk/rust/nym-sdk/src/mixnet/native_client.rs b/sdk/rust/nym-sdk/src/mixnet/native_client.rs index e03d585ebb..3266b37506 100644 --- a/sdk/rust/nym-sdk/src/mixnet/native_client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/native_client.rs @@ -2,7 +2,8 @@ use crate::mixnet::client::MixnetClientBuilder; use crate::mixnet::traits::MixnetMessageSender; use crate::{Error, Result}; use async_trait::async_trait; -use futures::{Stream, StreamExt}; +use futures::{ready, Stream, StreamExt}; +use log::error; use nym_client_core::client::{ base_client::{ClientInput, ClientOutput, ClientState}, inbound_messages::InputMessage, @@ -42,9 +43,33 @@ pub struct MixnetClient { /// The task manager that controlls all the spawned tasks that the clients uses to do it's job. pub(crate) task_manager: TaskManager, pub(crate) packet_type: Option, + + // internal state used for the `Stream` implementation + _buffered: Vec, } impl MixnetClient { + pub(crate) fn new( + nym_address: Recipient, + client_input: ClientInput, + client_output: ClientOutput, + client_state: ClientState, + reconstructed_receiver: ReconstructedMessagesReceiver, + task_manager: TaskManager, + packet_type: Option, + ) -> Self { + Self { + nym_address, + client_input, + client_output, + client_state, + reconstructed_receiver, + task_manager, + packet_type, + _buffered: Vec::new(), + } + } + /// Create a new client and connect to the mixnet using ephemeral in-memory keys that are /// discarded at application close. /// @@ -145,10 +170,32 @@ pub struct MixnetClientSender { } impl Stream for MixnetClient { - type Item = Vec; + type Item = ReconstructedMessage; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.reconstructed_receiver).poll_next(cx) + if let Some(next) = self._buffered.pop() { + cx.waker().wake_by_ref(); + return Poll::Ready(Some(next)); + } + match ready!(Pin::new(&mut self.reconstructed_receiver).poll_next(cx)) { + None => Poll::Ready(None), + Some(mut msgs) => { + // the vector itself should never be empty + if let Some(next) = msgs.pop() { + // there's more than a single message - buffer them and wake the waker + // to get polled again immediately + if !msgs.is_empty() { + self._buffered = msgs; + cx.waker().wake_by_ref(); + } + Poll::Ready(Some(next)) + } else { + error!("the reconstructed messages vector is empty - please let the developers know if you see this message"); + cx.waker().wake_by_ref(); + Poll::Pending + } + } + } } }