fix: later sends get ReceiverClosed too after the remote receiver is dropped

src/rpc.rs

@@ -527,11 +527,41 @@ impl<T: RpcMessage> NoqSenderInner<T> {
     }
 }
 
-#[derive(Default)]
 enum NoqSenderState<T> {
     Open(NoqSenderInner<T>),
-    #[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<T> Default for NoqSenderState<T> {
+    fn default() -> Self {
+        Self::Closed(CloseReason::Other)
+    }
+}
+
+impl<T> NoqSenderState<T> {
+    /// Returns the state after a send, so that later sends fail with the same kind of error.
+    fn after_send<R>(sender: NoqSenderInner<T>, res: &Result<R, SendError>) -> Self {
+        match res {
+            Ok(_) => Self::Open(sender),
+            Err(SendError::ReceiverClosed { .. }) => Self::Closed(CloseReason::ReceiverClosed),
+            Err(_) => Self::Closed(CloseReason::Other),
+        }
+    }
 }
 
 struct NoqSender<T>(tokio::sync::Mutex<NoqSenderState<T>>);
@@ -550,12 +580,13 @@ impl<T: RpcMessage> DynSender<T> for NoqSender<T> {
             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<T: RpcMessage> DynSender<T> for NoqSender<T> {
             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<T: RpcMessage> DynSender<T> for NoqSender<T> {
             let mut guard = self.0.lock().await;
             match guard.deref_mut() {
                 NoqSenderState::Open(sender) => sender.closed().await,
-                NoqSenderState::Closed => {}
+                NoqSenderState::Closed(_) => {}
             }
         })
     }

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::<Vec<u8>>::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::<Vec<u8>>::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()