diff --git a/crates/net/p2p/src/full_block.rs b/crates/net/p2p/src/full_block.rs index 91b786e410ca..0116f1348810 100644 --- a/crates/net/p2p/src/full_block.rs +++ b/crates/net/p2p/src/full_block.rs @@ -177,6 +177,9 @@ where fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); + // preemptive yield point + let mut budget = 4; + loop { match ready!(this.request.poll(cx)) { ResponseResult::Header(res) => { @@ -232,6 +235,14 @@ where if let Some(res) = this.take_block() { return Poll::Ready(res) } + + // ensure we still have enough budget for another iteration + budget -= 1; + if budget == 0 { + // make sure we're woken up again + cx.waker().wake_by_ref(); + return Poll::Pending + } } } }