Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akshitj11 wants to merge 2 commits into
base: master
Choose a base branch
from
+133 −5
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
55b3728
fix(swarm): bound Connection::poll iterations for cooperative scheduling
akshitj11 Jul 8, 2026
1532071
fix(swarm): address review feedback on Connection::poll budget
akshitj11 Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions swarm/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
## 0.48.0

- Bound `Connection::poll` fixed-point loop iterations to improve cooperative scheduling under load.
See [Issue 6438](https://github.com/libp2p/rust-libp2p/issues/6438).

- Remove `wasm-bindgen` feature and make `wasm` support implicit.
See [PR 6102](https://github.com/libp2p/rust-libp2p/pull/6102)

Expand Down
135 changes: 130 additions & 5 deletions swarm/src/connection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,22 @@ 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;
Comment thread
akshitj11 marked this conversation as resolved.

fn consume_poll_budget(budget: &mut u32, cx: &Context<'_>) -> bool {
*budget -= 1;
if *budget == 0 {
cx.waker().wake_by_ref();
true
} else {
false
}
}

static NEXT_CONNECTION_ID: AtomicUsize = AtomicUsize::new(1);

/// Connection identifier.
Expand Down Expand Up @@ -271,16 +287,26 @@ 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(()))) => {
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some(Err(info))) => {
handler.on_connection_event(ConnectionEvent::DialUpgradeError(
DialUpgradeError {
info,
error: StreamUpgradeError::Timeout,
},
));
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(None) | Poll::Pending => {}
Expand All @@ -294,7 +320,10 @@ where
let (upgrade, user_data) = protocol.into_upgrade();

requested_substreams.push(SubstreamRequested::new(user_data, timeout, upgrade));
continue; // Poll handler until exhausted.
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event)) => {
return Poll::Ready(Ok(Event::Handler(event)));
Expand All @@ -308,6 +337,9 @@ where
handler.on_connection_event(ConnectionEvent::RemoteProtocolsChange(added));
remote_supported_protocols.extend(protocol_buffer.drain(..));
}
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(
Expand All @@ -321,6 +353,9 @@ where
handler
.on_connection_event(ConnectionEvent::RemoteProtocolsChange(removed));
}
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
}
Expand All @@ -333,12 +368,18 @@ where
handler.on_connection_event(ConnectionEvent::FullyNegotiatedOutbound(
FullyNegotiatedOutbound { protocol, info },
));
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some((info, Err(error)))) => {
handler.on_connection_event(ConnectionEvent::DialUpgradeError(
DialUpgradeError { info, error },
));
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
}
Expand All @@ -351,24 +392,39 @@ where
handler.on_connection_event(ConnectionEvent::FullyNegotiatedInbound(
FullyNegotiatedInbound { protocol, info },
));
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some((info, Err(StreamUpgradeError::Apply(error))))) => {
handler.on_connection_event(ConnectionEvent::ListenUpgradeError(
ListenUpgradeError { info, error },
));
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some((_, Err(StreamUpgradeError::Io(e))))) => {
tracing::debug!("failed to upgrade inbound stream: {e}");
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some((_, Err(StreamUpgradeError::NegotiationFailed)))) => {
tracing::debug!("no protocol could be agreed upon for inbound stream");
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
Poll::Ready(Some((_, Err(StreamUpgradeError::Timeout)))) => {
tracing::debug!("inbound stream upgrade timed out");
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
}
Expand Down Expand Up @@ -428,6 +484,9 @@ where

// Go back to the top,
// handler can potentially make progress again.
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
}
Expand All @@ -447,6 +506,9 @@ where

// Go back to the top,
// handler can potentially make progress again.
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}
}
Expand All @@ -463,6 +525,9 @@ where
handler.on_connection_event(ConnectionEvent::LocalProtocolsChange(change));
}
// Go back to the top, handler can potentially make progress again.
if consume_poll_budget(&mut poll_budget, cx) {
return Poll::Pending;
}
continue;
}

Expand Down Expand Up @@ -783,7 +848,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,
};

Expand Down Expand Up @@ -818,9 +887,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,
Expand All @@ -831,6 +900,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);
Expand Down