fix: prevent message loss in peer_connection_listener race condition (#2255)

Co-authored-by: Claude Opus 4.5 <[email protected]>

Author: sanity

crates/core/src/client_events/combinator.rs

@@ -138,32 +138,67 @@ impl<const N: usize> super::ClientEventsProxy for ClientEventsCombinator<N> {
     }
 }
 
+/// Handles bidirectional communication between a client and the host.
+///
+/// Uses `tokio::select! { biased; ... }` to ensure host responses are always
+/// processed before client disconnect errors, preventing lost responses.
+/// The `biased` modifier ensures futures are polled in declaration order,
+/// and if multiple are ready, the first one wins.
 async fn client_fn(
     mut client: BoxedClient,
     mut rx: Receiver<(ClientId, HostResult)>,
     tx_host: Sender<Result<OpenRequest<'static>, ClientError>>,
 ) {
     loop {
         tokio::select! {
-            host_msg = rx.recv() => {
-                if let Some((client_id, response)) = host_msg {
-                    if client.send(client_id, response).await.is_err() {
+            biased;
+
+            // Priority 1: Host responses (highest priority)
+            // Ensures responses are sent to client before handling disconnect
+            msg = rx.recv() => {
+                match msg {
+                    Some((client_id, response)) => {
+                        if client.send(client_id, response).await.is_err() {
+                            break;
+                        }
+                    }
+                    None => {
+                        tracing::debug!("disconnected host");
                         break;
                     }
-                } else {
-                    tracing::debug!("disconnected host");
-                    break;
                 }
             }
-            client_msg = client.recv() => {
-                match client_msg {
-                    Ok(OpenRequest { client_id, request_id, request, notification_channel, token, attested_contract }) => {
-                        tracing::debug!("received msg @ combinator from external id {client_id}, msg: {request}");
-                        if tx_host.send(Ok(OpenRequest { client_id, request_id, request, notification_channel, token, attested_contract })).await.is_err() {
+
+            // Priority 2: Client requests
+            result = client.recv() => {
+                match result {
+                    Ok(OpenRequest {
+                        client_id,
+                        request_id,
+                        request,
+                        notification_channel,
+                        token,
+                        attested_contract,
+                    }) => {
+                        tracing::debug!(
+                            "received msg @ combinator from external id {client_id}, msg: {request}"
+                        );
+                        if tx_host
+                            .send(Ok(OpenRequest {
+                                client_id,
+                                request_id,
+                                request,
+                                notification_channel,
+                                token,
+                                attested_contract,
+                            }))
+                            .await
+                            .is_err()
+                        {
                             break;
                         }
                     }
-                    Err(err) if matches!(err.kind(), ErrorKind::ChannelClosed) =>{
+                    Err(err) if matches!(err.kind(), ErrorKind::ChannelClosed) => {
                         tracing::debug!("disconnected client");
                         let _ = tx_host.send(Err(err)).await;
                         break;

crates/core/src/contract/executor.rs

@@ -6,9 +6,13 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::Display;
 use std::future::Future;
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::{Duration, Instant};
 
+use futures::Stream;
+
 use either::Either;
 use freenet_stdlib::client_api::{
     ClientError as WsClientError, ClientRequest, ContractError as StdContractError,
@@ -343,16 +347,11 @@ impl ExecutorToEventLoopChannel<ExecutorHalve> {
 impl ExecutorToEventLoopChannel<NetworkEventListenerHalve> {
     pub async fn transaction_from_executor(&mut self) -> anyhow::Result<Transaction> {
         tracing::trace!("Waiting to receive transaction from executor channel");
-        let tx = self
-            .end
-            .waiting_for_op_rx
-            .recv()
-            .await
-            .ok_or_else(|| {
-                tracing::error!("Executor channel closed - all senders have been dropped");
-                tracing::error!("This typically happens when: 1) The executor task panicked/exited, 2) Network timeout cascaded to channel closure, 3) Resource constraints in CI");
-                anyhow::anyhow!("channel closed")
-            })?;
+        let tx = self.end.waiting_for_op_rx.recv().await.ok_or_else(|| {
+            tracing::error!("Executor channel closed - all senders have been dropped");
+            tracing::error!("This typically happens when: 1) The executor task panicked/exited, 2) Network timeout cascaded to channel closure, 3) Resource constraints in CI");
+            anyhow::anyhow!("channel closed")
+        })?;
         tracing::trace!("Successfully received transaction from executor channel");
         Ok(tx)
     }
@@ -367,6 +366,14 @@ impl ExecutorToEventLoopChannel<NetworkEventListenerHalve> {
     }
 }
 
+impl Stream for ExecutorToEventLoopChannel<NetworkEventListenerHalve> {
+    type Item = Transaction;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.end.waiting_for_op_rx).poll_recv(cx)
+    }
+}
+
 impl ExecutorToEventLoopChannel<Callback> {
     pub async fn response(&mut self, result: OpEnum) {
         if self.end.response_for_tx.send(result).await.is_err() {

crates/core/src/contract/handler.rs

@@ -9,10 +9,14 @@
 use std::collections::BTreeMap;
 use std::future::Future;
 use std::hash::Hash;
+use std::pin::Pin;
 use std::sync::atomic::{AtomicU64, Ordering::SeqCst};
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::Duration;
 
+use futures::Stream;
+
 use freenet_stdlib::client_api::DelegateRequest;
 use freenet_stdlib::prelude::*;
 use serde::{Deserialize, Serialize};
@@ -281,6 +285,14 @@ impl ContractHandlerChannel<WaitingResolution> {
     }
 }
 
+impl Stream for ContractHandlerChannel<WaitingResolution> {
+    type Item = (ClientId, WaitingTransaction);
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.end.wait_for_res_rx).poll_recv(cx)
+    }
+}
+
 impl ContractHandlerChannel<SenderHalve> {
     // TODO: the timeout should be derived from whatever is the worst
     // case we are willing to accept for waiting out for an event;

crates/core/src/node/network_bridge/p2p_protoc.rs

@@ -2075,6 +2075,11 @@ impl P2pConnManager {
             tokio::spawn(async move {
                 peer_connection_listener(rx, connection, peer_addr, conn_events).await;
             });
+            // Yield to allow the spawned peer_connection_listener task to start.
+            // This is important because on some runtimes (especially in tests with boxed_local
+            // futures), spawned tasks may not be scheduled immediately, causing messages
+            // sent to the channel to pile up without being processed.
+            tokio::task::yield_now().await;
             newly_inserted = true;
         } else {
             tracing::debug!(
@@ -2669,42 +2674,47 @@ async fn notify_transport_closed(
     }
 }
 
+/// Listens for messages on a peer connection using drain-then-select pattern.
+///
+/// On each iteration, drains all pending outbound messages via `try_recv()` before
+/// waiting for either new outbound or inbound messages. This approach:
+/// 1. Ensures queued outbound messages are sent promptly
+/// 2. Avoids starving inbound (which would happen with biased select)
+/// 3. No messages are lost due to poll ordering
 async fn peer_connection_listener(
     mut rx: PeerConnChannelRecv,
     mut conn: PeerConnection,
     peer_addr: SocketAddr,
     conn_events: Sender<ConnEvent>,
 ) {
-    const MAX_IMMEDIATE_SENDS: usize = 32;
     let remote_addr = conn.remote_addr();
     tracing::debug!(
         to = %remote_addr,
         %peer_addr,
         "[CONN_LIFECYCLE] Starting peer_connection_listener task"
     );
+
     loop {
-        let mut drained = 0;
+        // Drain all pending outbound messages first before checking inbound.
+        // This ensures queued outbound messages are sent promptly without starving
+        // the inbound channel (which would happen with biased select).
         loop {
             match rx.try_recv() {
                 Ok(msg) => {
                     if let Err(error) = handle_peer_channel_message(&mut conn, msg).await {
                         tracing::debug!(
                             to = %remote_addr,
                             ?error,
-                            "[CONN_LIFECYCLE] Shutting down connection after send failure"
+                            "[CONN_LIFECYCLE] Connection closed after channel command"
                         );
                         notify_transport_closed(&conn_events, remote_addr, error).await;
                         return;
                     }
-                    drained += 1;
-                    if drained >= MAX_IMMEDIATE_SENDS {
-                        break;
-                    }
                 }
                 Err(TryRecvError::Empty) => break,
                 Err(TryRecvError::Disconnected) => {
                     tracing::warn!(
-                        to = %conn.remote_addr(),
+                        to = %remote_addr,
                         "[CONN_LIFECYCLE] peer_connection_listener channel closed without explicit DropConnection"
                     );
                     notify_transport_closed(
@@ -2718,74 +2728,80 @@ async fn peer_connection_listener(
             }
         }
 
+        // Now wait for either new outbound or inbound messages fairly
         tokio::select! {
             msg = rx.recv() => {
-                let Some(msg) = msg else {
-                    tracing::warn!(
-                        to = %conn.remote_addr(),
-                        "[CONN_LIFECYCLE] peer_connection_listener channel closed without explicit DropConnection"
-                    );
-                    notify_transport_closed(
-                        &conn_events,
-                        remote_addr,
-                        TransportError::ConnectionClosed(remote_addr),
-                    )
-                    .await;
-                    return;
-                };
-                if let Err(error) = handle_peer_channel_message(&mut conn, msg).await {
-                    tracing::debug!(
-                        to = %remote_addr,
-                        ?error,
-                        "[CONN_LIFECYCLE] Connection closed after channel command"
-                    );
-                    notify_transport_closed(&conn_events, remote_addr, error).await;
-                    return;
+                match msg {
+                    Some(msg) => {
+                        if let Err(error) = handle_peer_channel_message(&mut conn, msg).await {
+                            tracing::debug!(
+                                to = %remote_addr,
+                                ?error,
+                                "[CONN_LIFECYCLE] Connection closed after channel command"
+                            );
+                            notify_transport_closed(&conn_events, remote_addr, error).await;
+                            return;
+                        }
+                    }
+                    None => {
+                        tracing::warn!(
+                            to = %remote_addr,
+                            "[CONN_LIFECYCLE] peer_connection_listener channel closed without explicit DropConnection"
+                        );
+                        notify_transport_closed(
+                            &conn_events,
+                            remote_addr,
+                            TransportError::ConnectionClosed(remote_addr),
+                        )
+                        .await;
+                        return;
+                    }
                 }
             }
+
             msg = conn.recv() => {
                 match msg {
-                    Ok(msg) => {
-                        match decode_msg(&msg) {
-                            Ok(net_message) => {
-                                let tx = *net_message.id();
+                    Ok(msg) => match decode_msg(&msg) {
+                        Ok(net_message) => {
+                            let tx = *net_message.id();
+                            tracing::debug!(
+                                from = %remote_addr,
+                                %tx,
+                                tx_type = ?tx.transaction_type(),
+                                msg_type = %net_message,
+                                "[CONN_LIFECYCLE] Received inbound NetMessage from peer"
+                            );
+                            if conn_events
+                                .send(ConnEvent::InboundMessage(IncomingMessage::with_remote(
+                                    net_message,
+                                    remote_addr,
+                                )))
+                                .await
+                                .is_err()
+                            {
                                 tracing::debug!(
-                                    from = %conn.remote_addr(),
-                                    %tx,
-                                    tx_type = ?tx.transaction_type(),
-                                    msg_type = %net_message,
-                                    "[CONN_LIFECYCLE] Received inbound NetMessage from peer"
-                                );
-                                if conn_events.send(ConnEvent::InboundMessage(IncomingMessage::with_remote(net_message, remote_addr))).await.is_err() {
-                                    tracing::debug!(
-                                        from = %remote_addr,
-                                        "[CONN_LIFECYCLE] conn_events receiver dropped; stopping listener"
-                                    );
-                                    return;
-                                }
-                            }
-                            Err(error) => {
-                                tracing::error!(
-                                    from = %conn.remote_addr(),
-                                    ?error,
-                                    "[CONN_LIFECYCLE] Failed to deserialize inbound message; closing connection"
+                                    from = %remote_addr,
+                                    "[CONN_LIFECYCLE] conn_events receiver dropped; stopping listener"
                                 );
-                                let transport_error = TransportError::Other(anyhow!(
-                                    "Failed to deserialize inbound message from {remote_addr}: {error:?}"
-                                ));
-                                notify_transport_closed(
-                                    &conn_events,
-                                    remote_addr,
-                                    transport_error,
-                                )
-                                .await;
                                 return;
                             }
                         }
-                    }
+                        Err(error) => {
+                            tracing::error!(
+                                from = %remote_addr,
+                                ?error,
+                                "[CONN_LIFECYCLE] Failed to deserialize inbound message; closing connection"
+                            );
+                            let transport_error = TransportError::Other(anyhow!(
+                                "Failed to deserialize inbound message from {remote_addr}: {error:?}"
+                            ));
+                            notify_transport_closed(&conn_events, remote_addr, transport_error).await;
+                            return;
+                        }
+                    },
                     Err(error) => {
                         tracing::debug!(
-                            from = %conn.remote_addr(),
+                            from = %remote_addr,
                             ?error,
                             "[CONN_LIFECYCLE] peer_connection_listener terminating after recv error"
                         );