fix(swarm): bound Connection::poll iterations for cooperative scheduling

Cargo.toml

@@ -103,7 +103,7 @@ libp2p-rendezvous = { version = "0.18.0", path = "protocols/rendezvous" }
 libp2p-request-response = { version = "0.30.0", path = "protocols/request-response" }
 libp2p-server = { version = "0.13.0", path = "misc/server" }
 libp2p-stream = { version = "0.5.0-alpha", path = "protocols/stream" }
-libp2p-swarm = { version = "0.48.0", path = "swarm" }
+libp2p-swarm = { version = "0.48.1", path = "swarm" }
 libp2p-swarm-derive = { version = "=0.36.0", path = "swarm-derive" } # `libp2p-swarm-derive` may not be compatible with different `libp2p-swarm` non-breaking releases. E.g. `libp2p-swarm` might introduce a new enum variant `FromSwarm` (which is `#[non-exhaustive]`) in a non-breaking release. Older versions of `libp2p-swarm-derive` would not forward this enum variant within the `NetworkBehaviour` hierarchy. Thus the version pinning is required.
 libp2p-swarm-test = { version = "0.7.0", path = "swarm-test" }
 libp2p-tcp = { version = "0.45.0", path = "transports/tcp" }

swarm/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.48.1
+
+- Bound `Connection::poll` fixed-point loop iterations to improve cooperative scheduling under load.
+  See [Issue 6438](https://github.com/libp2p/rust-libp2p/issues/6438).
+
 ## 0.48.0
 
 - Remove `wasm-bindgen` feature and make `wasm` support implicit.

swarm/Cargo.toml

@@ -3,7 +3,7 @@ name = "libp2p-swarm"
 edition.workspace = true
 rust-version = { workspace = true }
 description = "The libp2p swarm"
-version = "0.48.0"
+version = "0.48.1"
 authors = ["Parity Technologies <admin@parity.io>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"

swarm/src/connection.rs

@@ -63,6 +63,23 @@ use crate::{
     upgrade::{InboundUpgradeSend, OutboundUpgradeSend},
 };
 
+/// Maximum number of internal progress iterations per [`Connection::poll`] call.
+///
+/// When exhausted, the connection yields with a waker notification so the executor
+/// can schedule other tasks. Aligned with Tokio's default cooperative budget (128).
+const CONNECTION_POLL_ITERATION_BUDGET: u32 = 128;
+
+macro_rules! continue_after_poll_progress {
+    ($budget:expr, $cx:expr) => {{
+        $budget -= 1;
+        if $budget == 0 {
+            $cx.waker().wake_by_ref();
+            return Poll::Pending;
+        }
+        continue;
+    }};
+}
+
 static NEXT_CONNECTION_ID: AtomicUsize = AtomicUsize::new(1);
 
 /// Connection identifier.
@@ -271,17 +288,19 @@ where
             ..
         } = self.get_mut();
 
+        let mut poll_budget = CONNECTION_POLL_ITERATION_BUDGET;
+
         loop {
             match requested_substreams.poll_next_unpin(cx) {
-                Poll::Ready(Some(Ok(()))) => continue,
+                Poll::Ready(Some(Ok(()))) => continue_after_poll_progress!(poll_budget, cx),
                 Poll::Ready(Some(Err(info))) => {
                     handler.on_connection_event(ConnectionEvent::DialUpgradeError(
                         DialUpgradeError {
                             info,
                             error: StreamUpgradeError::Timeout,
                         },
                     ));
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(None) | Poll::Pending => {}
             }
@@ -294,7 +313,7 @@ where
                     let (upgrade, user_data) = protocol.into_upgrade();
 
                     requested_substreams.push(SubstreamRequested::new(user_data, timeout, upgrade));
-                    continue; // Poll handler until exhausted.
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event)) => {
                     return Poll::Ready(Ok(Event::Handler(event)));
@@ -308,7 +327,7 @@ where
                         handler.on_connection_event(ConnectionEvent::RemoteProtocolsChange(added));
                         remote_supported_protocols.extend(protocol_buffer.drain(..));
                     }
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(
                     ProtocolSupport::Removed(protocols),
@@ -321,7 +340,7 @@ where
                         handler
                             .on_connection_event(ConnectionEvent::RemoteProtocolsChange(removed));
                     }
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
             }
 
@@ -333,13 +352,13 @@ where
                     handler.on_connection_event(ConnectionEvent::FullyNegotiatedOutbound(
                         FullyNegotiatedOutbound { protocol, info },
                     ));
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(Some((info, Err(error)))) => {
                     handler.on_connection_event(ConnectionEvent::DialUpgradeError(
                         DialUpgradeError { info, error },
                     ));
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
             }
 
@@ -351,25 +370,25 @@ where
                     handler.on_connection_event(ConnectionEvent::FullyNegotiatedInbound(
                         FullyNegotiatedInbound { protocol, info },
                     ));
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(Some((info, Err(StreamUpgradeError::Apply(error))))) => {
                     handler.on_connection_event(ConnectionEvent::ListenUpgradeError(
                         ListenUpgradeError { info, error },
                     ));
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(Some((_, Err(StreamUpgradeError::Io(e))))) => {
                     tracing::debug!("failed to upgrade inbound stream: {e}");
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(Some((_, Err(StreamUpgradeError::NegotiationFailed)))) => {
                     tracing::debug!("no protocol could be agreed upon for inbound stream");
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
                 Poll::Ready(Some((_, Err(StreamUpgradeError::Timeout)))) => {
                     tracing::debug!("inbound stream upgrade timed out");
-                    continue;
+                    continue_after_poll_progress!(poll_budget, cx);
                 }
             }
 
@@ -428,7 +447,7 @@ where
 
                         // Go back to the top,
                         // handler can potentially make progress again.
-                        continue;
+                        continue_after_poll_progress!(poll_budget, cx);
                     }
                 }
             }
@@ -447,7 +466,7 @@ where
 
                         // Go back to the top,
                         // handler can potentially make progress again.
-                        continue;
+                        continue_after_poll_progress!(poll_budget, cx);
                     }
                 }
             }
@@ -463,7 +482,7 @@ where
                     handler.on_connection_event(ConnectionEvent::LocalProtocolsChange(change));
                 }
                 // Go back to the top, handler can potentially make progress again.
-                continue;
+                continue_after_poll_progress!(poll_budget, cx);
             }
 
             // Nothing can make progress, return `Pending`.
@@ -783,7 +802,11 @@ impl<T: AsRef<str>> std::hash::Hash for AsStrHashEq<T> {
 mod tests {
     use std::{
         convert::Infallible,
-        sync::{Arc, Weak},
+        sync::{
+            Arc, Weak,
+            atomic::{AtomicUsize, Ordering},
+        },
+        task::{Context, RawWaker, RawWakerVTable, Waker},
         time::Instant,
     };
 
@@ -818,9 +841,9 @@ mod tests {
                 Duration::ZERO,
             );
 
-            let result = connection.poll_noop_waker();
+            while connection.poll_noop_waker().is_ready() {}
 
-            assert!(result.is_pending());
+            assert!(connection.poll_noop_waker().is_pending());
             assert_eq!(
                 Arc::weak_count(&alive_substream_counter),
                 max_negotiating_inbound_streams,
@@ -831,6 +854,62 @@ mod tests {
         QuickCheck::new().quickcheck(prop as fn(_));
     }
 
+    #[test]
+    fn poll_budget_yields_and_wakes() {
+        static WAKE_COUNT: AtomicUsize = AtomicUsize::new(0);
+
+        fn counting_waker() -> Waker {
+            unsafe fn clone(_: *const ()) -> RawWaker {
+                RawWaker::new(std::ptr::null(), &VTABLE)
+            }
+            unsafe fn wake(_: *const ()) {
+                WAKE_COUNT.fetch_add(1, Ordering::SeqCst);
+            }
+            unsafe fn wake_by_ref(_: *const ()) {
+                WAKE_COUNT.fetch_add(1, Ordering::SeqCst);
+            }
+            unsafe fn drop(_: *const ()) {}
+
+            static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
+
+            unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) }
+        }
+
+        WAKE_COUNT.store(0, Ordering::SeqCst);
+
+        let max_negotiating_inbound_streams = (CONNECTION_POLL_ITERATION_BUDGET + 64) as usize;
+        let alive_substream_counter = Arc::new(());
+        let mut connection = Connection::new(
+            StreamMuxerBox::new(DummyStreamMuxer {
+                counter: alive_substream_counter.clone(),
+            }),
+            MockConnectionHandler::new(Duration::from_secs(10)),
+            None,
+            max_negotiating_inbound_streams,
+            Duration::ZERO,
+        );
+
+        let waker = counting_waker();
+        let mut cx = Context::from_waker(&waker);
+
+        assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
+        assert!(
+            WAKE_COUNT.load(Ordering::SeqCst) >= 1,
+            "expected waker notification when poll budget is exhausted"
+        );
+        assert!(
+            Arc::weak_count(&alive_substream_counter) <= CONNECTION_POLL_ITERATION_BUDGET as usize,
+        );
+
+        while Pin::new(&mut connection).poll(&mut cx).is_ready() {}
+
+        assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
+        assert_eq!(
+            Arc::weak_count(&alive_substream_counter),
+            max_negotiating_inbound_streams,
+        );
+    }
+
     #[test]
     fn outbound_stream_timeout_starts_on_request() {
         let upgrade_timeout = Duration::from_secs(1);