From 250785f833246ba7c9a1e48b47cd0c2a4606ed09 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Thu, 10 Oct 2024 19:15:06 +0200
Subject: [PATCH] chore: preempt single block downloading (#11647)

---
 crates/net/p2p/src/full_block.rs | 11 +++++++++++
 1 file changed, 11 insertions(+)

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<Self::Output> {
         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
+            }
         }
     }
 }