diff --git a/src/rpc.rs b/src/rpc.rs index e979b8f..cbdf737 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -527,11 +527,41 @@ impl NoqSenderInner { } } -#[derive(Default)] enum NoqSenderState { Open(NoqSenderInner), - #[default] - Closed, + Closed(CloseReason), +} + +#[derive(Debug, Clone, Copy)] +enum CloseReason { + ReceiverClosed, + Other, +} + +impl CloseReason { + fn error(self) -> SendError { + match self { + Self::ReceiverClosed => e!(SendError::ReceiverClosed), + Self::Other => io::Error::from(io::ErrorKind::BrokenPipe).into(), + } + } +} + +impl Default for NoqSenderState { + fn default() -> Self { + Self::Closed(CloseReason::Other) + } +} + +impl NoqSenderState { + /// Returns the state after a send, so that later sends fail with the same kind of error. + fn after_send(sender: NoqSenderInner, res: &Result) -> Self { + match res { + Ok(_) => Self::Open(sender), + Err(SendError::ReceiverClosed { .. }) => Self::Closed(CloseReason::ReceiverClosed), + Err(_) => Self::Closed(CloseReason::Other), + } + } } struct NoqSender(tokio::sync::Mutex>); @@ -550,12 +580,13 @@ impl DynSender for NoqSender { match sender { NoqSenderState::Open(mut sender) => { let res = sender.send(value).await; - if res.is_ok() { - *guard = NoqSenderState::Open(sender); - } + *guard = NoqSenderState::after_send(sender, &res); res } - NoqSenderState::Closed => Err(io::Error::from(io::ErrorKind::BrokenPipe).into()), + NoqSenderState::Closed(reason) => { + *guard = NoqSenderState::Closed(reason); + Err(reason.error()) + } } }) } @@ -570,12 +601,13 @@ impl DynSender for NoqSender { match sender { NoqSenderState::Open(mut sender) => { let res = sender.try_send(value).await; - if res.is_ok() { - *guard = NoqSenderState::Open(sender); - } + *guard = NoqSenderState::after_send(sender, &res); res } - NoqSenderState::Closed => Err(io::Error::from(io::ErrorKind::BrokenPipe).into()), + NoqSenderState::Closed(reason) => { + *guard = NoqSenderState::Closed(reason); + Err(reason.error()) + } } }) } @@ -585,7 +617,7 @@ impl DynSender for NoqSender { let mut guard = self.0.lock().await; match guard.deref_mut() { NoqSenderState::Open(sender) => sender.closed().await, - NoqSenderState::Closed => {} + NoqSenderState::Closed(_) => {} } }) } diff --git a/tests/mpsc_channel.rs b/tests/mpsc_channel.rs index f22416a..4a47288 100644 --- a/tests/mpsc_channel.rs +++ b/tests/mpsc_channel.rs @@ -152,6 +152,54 @@ async fn mpsc_receiver_dropped() -> TestResult<()> { Ok(()) } +/// Checks that after the remote receiver is dropped, every send on every clone +/// of a `Sender` gets `ReceiverClosed`, also on clones that did not see the stop. +#[tokio::test] +async fn mpsc_receiver_dropped_clones() -> TestResult<()> { + let (server, client, server_addr) = create_connected_endpoints()?; + let server = tokio::spawn(async move { + let conn = server.accept().await.unwrap().await?; + let (_, recv) = conn.accept_bi().await?; + drop(Receiver::>::from(recv)); + // keep the connection, so that the stop reaches the client + conn.closed().await; + TestResult::Ok(()) + }); + let conn = client.connect(server_addr, "localhost")?.await?; + let (send, _) = conn.open_bi().await?; + let send1 = mpsc::Sender::>::from(send); + let send2 = send1.clone(); + // send on the first clone until the stop arrives + let err = timeout(Duration::from_secs(5), async { + loop { + if let Err(err) = send1.send(vec![1, 2, 3]).await { + break err; + } + } + }) + .await?; + assert!( + matches!(err, SendError::ReceiverClosed { .. }), + "expected ReceiverClosed, got {err:?}" + ); + // later sends on the same clone, and sends on a clone that never saw the stop + assert!(matches!( + send1.send(vec![1]).await, + Err(SendError::ReceiverClosed { .. }) + )); + assert!(matches!( + send2.send(vec![1]).await, + Err(SendError::ReceiverClosed { .. }) + )); + assert!(matches!( + send2.try_send(vec![1]).await, + Err(SendError::ReceiverClosed { .. }) + )); + conn.close(0u32.into(), b""); + server.await??; + Ok(()) +} + async fn vec_receiver(server: Endpoint) -> Result<(), RecvError> { let conn = server .accept()