perf(network): pack block sync ranges by size hint

Author: evan-forbes

zebra-network/src/zakura/block_sync/peer_routine.rs

@@ -547,47 +547,43 @@ impl PeerRoutine {
             }
             // One contiguous chunk up to the peer's per-request count cap; the
             // outer loop fills the rest of the peer's slots.
-            let max_count = usize::try_from(
+            let local_peer_count_cap = usize::try_from(
                 self.max_blocks_per_response
                     .min(self.config.advertised_max_blocks_per_response())
                     .max(1),
             )
             .unwrap_or(usize::MAX);
             let (servable_low, servable_high) = (self.servable_low, self.servable_high);
 
-            // Compute this chunk's byte ceiling BEFORE taking any work, from the
-            // global byte budget and this peer's per-response cap. The reservation
-            // is worst-case per block (it only ever shrinks toward the actual size
-            // on receipt, so a valid body is never dropped for a full budget). If
-            // the ceiling cannot fund even one worst-case block, break and wait on
-            // the budget-capacity / work-added notifications instead of taking a
-            // chunk only to return it. Taking-then-returning would call
-            // `work.return_items` → `notify_waiters`, which re-wakes THIS routine's
-            // own enabled `work_added` notification (registered before the fill) and
-            // busy-loops the want-work arm (missed-wake's mirror image — a
-            // self-wake spin). Gating before the take keeps the routine parked on
-            // `capacity` until budget frees up.
-            let max_bytes = self
-                .budget
-                .available()
-                .min(u64::from(self.max_response_bytes.max(1)));
-            // Cap the chunk taken to what the byte ceiling can fund at worst case;
-            // break (without taking) when not even one block fits, so no take/return
-            // self-wake cycle can occur.
-            let byte_capped_count = max_bytes
+            // Compute this chunk's byte ceiling BEFORE taking any work. The
+            // count cap is still bounded by what the global budget can reserve at
+            // worst case, so advertised sizes never weaken the existing
+            // pre-send reservation. The estimate cap only decides how many of the
+            // already-affordable heights to pack into this one request.
+            let available_bytes = self.budget.available();
+            let max_count = available_bytes
                 .checked_div(worst)
-                .map(|count| usize::try_from(count).unwrap_or(usize::MAX).min(max_count))
-                .unwrap_or(max_count);
-            if byte_capped_count == 0 {
+                .map(|count| {
+                    usize::try_from(count)
+                        .unwrap_or(usize::MAX)
+                        .min(local_peer_count_cap)
+                })
+                .unwrap_or(local_peer_count_cap);
+            if max_count == 0 {
                 break;
             }
+            let max_estimated_bytes =
+                available_bytes.min(u64::from(self.max_response_bytes.max(1)));
 
             // Take work in this peer's servable range. `servable_high` is NOT
             // clamped to the floor: a peer fetches as far ahead of the committed
             // floor as its servable range and the byte budget allow.
-            let mut items = self
-                .work
-                .take_in_range(servable_low, servable_high, byte_capped_count);
+            let mut items = self.work.take_in_range_budgeted(
+                servable_low,
+                servable_high,
+                max_count,
+                max_estimated_bytes,
+            );
             if items.is_empty() {
                 break;
             }
@@ -619,8 +615,8 @@ impl PeerRoutine {
             }
             self.trace_work_taken(servable_low, servable_high, items.len());
 
-            // `take_in_range` already honoured the byte-capped count, so every
-            // taken item fits under `max_bytes`; nothing is returned here.
+            // `take_in_range_budgeted` already honoured the byte-capped count;
+            // every taken item has worst-case reservation capacity.
             let kept_count = items.len();
 
             let reserved_bytes = worst.saturating_mul(kept_count as u64);

zebra-network/src/zakura/block_sync/tests.rs

@@ -962,6 +962,95 @@ fn work_queue_take_respects_servable_range_contiguity_and_max() {
     );
 }
 
+#[test]
+fn work_queue_budgeted_take_respects_count_cap() {
+    let queue = work_queue_with(
+        0,
+        (1..=4).map(|height| needed(height, BlockSizeEstimate::Advertised(100))),
+    );
+
+    let taken = queue.take_in_range_budgeted(block::Height(1), block::Height(4), 2, u64::MAX);
+
+    assert_eq!(
+        taken.iter().map(|(height, _)| height.0).collect::<Vec<_>>(),
+        vec![1, 2]
+    );
+}
+
+#[test]
+fn work_queue_budgeted_take_respects_estimated_byte_cap() {
+    let queue = work_queue_with(
+        0,
+        [
+            needed(1, BlockSizeEstimate::Advertised(100)),
+            needed(2, BlockSizeEstimate::Advertised(150)),
+            needed(3, BlockSizeEstimate::Advertised(1)),
+        ],
+    );
+
+    let taken = queue.take_in_range_budgeted(block::Height(1), block::Height(3), 3, 250);
+
+    assert_eq!(
+        taken.iter().map(|(height, _)| height.0).collect::<Vec<_>>(),
+        vec![1, 2]
+    );
+    assert!(queue.pending_contains(block::Height(3)));
+}
+
+#[test]
+fn work_queue_budgeted_take_stops_at_gaps() {
+    let queue = work_queue_with(
+        0,
+        [
+            needed(10, BlockSizeEstimate::Advertised(100)),
+            needed(11, BlockSizeEstimate::Advertised(100)),
+            needed(13, BlockSizeEstimate::Advertised(100)),
+        ],
+    );
+
+    let taken = queue.take_in_range_budgeted(block::Height(10), block::Height(13), 3, u64::MAX);
+
+    assert_eq!(
+        taken.iter().map(|(height, _)| height.0).collect::<Vec<_>>(),
+        vec![10, 11]
+    );
+    assert!(queue.pending_contains(block::Height(13)));
+}
+
+#[test]
+fn work_queue_budgeted_take_takes_one_oversized_first_item_for_progress() {
+    let queue = work_queue_with(
+        0,
+        [
+            needed(1, BlockSizeEstimate::Advertised(500)),
+            needed(2, BlockSizeEstimate::Advertised(1)),
+        ],
+    );
+
+    let taken = queue.take_in_range_budgeted(block::Height(1), block::Height(2), 2, 100);
+
+    assert_eq!(
+        taken.iter().map(|(height, _)| height.0).collect::<Vec<_>>(),
+        vec![1]
+    );
+    assert_eq!(taken[0].1.estimated_bytes, 500);
+    assert!(queue.pending_contains(block::Height(2)));
+}
+
+#[test]
+fn work_queue_budgeted_take_preserves_estimates_through_take_and_return() {
+    let queue = work_queue_with(0, [needed(10, BlockSizeEstimate::Advertised(12_345))]);
+
+    let taken = queue.take_in_range_budgeted(block::Height(10), block::Height(10), 1, 1);
+    assert_eq!(taken.len(), 1);
+    assert_eq!(taken[0].1.estimated_bytes, 12_345);
+
+    queue.return_items([block::Height(10)]);
+    let retaken = queue.take_in_range_budgeted(block::Height(10), block::Height(10), 1, 1);
+    assert_eq!(retaken.len(), 1);
+    assert_eq!(retaken[0].1.estimated_bytes, 12_345);
+}
+
 #[test]
 fn work_queue_take_does_not_clamp_high_to_floor() {
     // The committed floor is NOT an upper bound on a take: a peer fetches as far
@@ -1221,7 +1310,7 @@ async fn reactor_suppresses_needed_block_query_when_work_already_covers_tip() {
             (1..=4)
                 .map(|height| BlockSyncBlockMeta {
                     height: block::Height(height),
-                    hash: block::Hash([height as u8; 32]),
+                    hash: block::Hash([u8::try_from(height).expect("test height fits u8"); 32]),
                     size: BlockSizeEstimate::Advertised(1_000),
                 })
                 .collect(),
@@ -4442,6 +4531,129 @@ async fn reactor_reserves_worst_case_per_block_not_size_hint() {
     reactor_task.abort();
 }
 
+#[tokio::test]
+async fn reactor_packs_small_estimates_under_peer_response_byte_cap() {
+    let config = ZakuraBlockSyncConfig {
+        max_inflight_block_bytes: 4 * BS_PER_BLOCK_WORST_CASE_BYTES,
+        max_blocks_per_response: 4,
+        ..ZakuraBlockSyncConfig::default()
+    };
+
+    let (_tip_tx, tip_rx) = watch::channel((block::Height(0), block::Hash([0; 32])));
+    let startup = BlockSyncStartup::new(
+        BlockSyncFrontiers {
+            finalized_height: block::Height(0),
+            verified_block_tip: block::Height(0),
+            verified_block_hash: block::Hash([0; 32]),
+        },
+        (block::Height(0), block::Hash([0; 32])),
+        tip_rx,
+        config.clone(),
+    );
+    let (handle, mut actions, reactor_task) = spawn_block_sync_reactor(startup);
+    let service = BlockSyncService::new_with_handle_for_test(config, handle.clone());
+
+    let one_worst_case_response =
+        u32::try_from(BS_PER_BLOCK_WORST_CASE_BYTES).expect("worst-case block size fits u32");
+    let (_peer_id, _inbound, mut outbound) = connect_peer_with_status_message(
+        &service,
+        &mut actions,
+        52,
+        BlockSyncStatus {
+            servable_low: block::Height(1),
+            servable_high: block::Height(4),
+            tip_hash: block::Hash([4; 32]),
+            max_blocks_per_response: 4,
+            max_inflight_requests: 1,
+            max_response_bytes: one_worst_case_response,
+        },
+    )
+    .await;
+
+    handle
+        .send(BlockSyncEvent::NeededBlocks(
+            (1..=4)
+                .map(|height| BlockSyncBlockMeta {
+                    height: block::Height(height),
+                    hash: block::Hash([height as u8; 32]),
+                    size: BlockSizeEstimate::Advertised(1_000),
+                })
+                .collect(),
+        ))
+        .await
+        .expect("needed metadata queues");
+
+    let (start_height, count) = wait_for_outbound_getblocks(&mut outbound).await;
+    assert_eq!(start_height, block::Height(1));
+    assert_eq!(
+        count, 4,
+        "small advertised estimates should pack more than the one worst-case block \
+         allowed by the peer response byte cap",
+    );
+
+    reactor_task.abort();
+}
+
+#[tokio::test]
+async fn reactor_tiny_estimates_do_not_exceed_one_worst_case_budget_block() {
+    let config = ZakuraBlockSyncConfig {
+        max_inflight_block_bytes: BS_PER_BLOCK_WORST_CASE_BYTES,
+        max_blocks_per_response: 4,
+        ..ZakuraBlockSyncConfig::default()
+    };
+
+    let (_tip_tx, tip_rx) = watch::channel((block::Height(0), block::Hash([0; 32])));
+    let startup = BlockSyncStartup::new(
+        BlockSyncFrontiers {
+            finalized_height: block::Height(0),
+            verified_block_tip: block::Height(0),
+            verified_block_hash: block::Hash([0; 32]),
+        },
+        (block::Height(0), block::Hash([0; 32])),
+        tip_rx,
+        config.clone(),
+    );
+    let (handle, mut actions, reactor_task) = spawn_block_sync_reactor(startup);
+    let service = BlockSyncService::new_with_handle_for_test(config, handle.clone());
+
+    let (_peer_id, _inbound, mut outbound) = connect_peer_with_status_message(
+        &service,
+        &mut actions,
+        53,
+        BlockSyncStatus {
+            servable_low: block::Height(1),
+            servable_high: block::Height(4),
+            tip_hash: block::Hash([4; 32]),
+            max_blocks_per_response: 4,
+            max_inflight_requests: 1,
+            max_response_bytes: MAX_BS_RESPONSE_BYTES,
+        },
+    )
+    .await;
+
+    handle
+        .send(BlockSyncEvent::NeededBlocks(
+            (1..=4)
+                .map(|height| BlockSyncBlockMeta {
+                    height: block::Height(height),
+                    hash: block::Hash([u8::try_from(height).expect("test height fits u8"); 32]),
+                    size: BlockSizeEstimate::Advertised(1),
+                })
+                .collect(),
+        ))
+        .await
+        .expect("needed metadata queues");
+
+    let (start_height, count) = wait_for_outbound_getblocks(&mut outbound).await;
+    assert_eq!(start_height, block::Height(1));
+    assert_eq!(
+        count, 1,
+        "only one worst-case block of global budget is available, regardless of tiny estimates",
+    );
+
+    reactor_task.abort();
+}
+
 #[tokio::test]
 async fn reactor_zero_pause_threshold_preserves_lag_one_downloads() {
     let config = immediate_body_download_config();

zebra-network/src/zakura/block_sync/work_queue.rs

@@ -190,6 +190,58 @@ impl WorkQueue {
         taken
     }
 
+    /// Move up to `max_count` contiguous-ascending `pending` heights within
+    /// `low..=high` from `pending` to `in_flight`, also stopping before the
+    /// sum of stored size estimates would exceed `max_estimated_bytes`.
+    ///
+    /// The estimate cap is scheduler input hygiene only. Callers must still
+    /// reserve their existing worst-case bytes before sending the request. To
+    /// guarantee progress, the first eligible item is always taken when
+    /// `max_count > 0`, even if its estimate alone exceeds the cap.
+    pub(super) fn take_in_range_budgeted(
+        &self,
+        low: block::Height,
+        high: block::Height,
+        max_count: usize,
+        max_estimated_bytes: u64,
+    ) -> Vec<(block::Height, WorkItem)> {
+        if max_count == 0 || low > high {
+            return Vec::new();
+        }
+        let mut inner = self.lock();
+        let mut taken: Vec<(block::Height, WorkItem)> = Vec::new();
+        let mut estimated_bytes = 0u64;
+        let mut next_expected: Option<block::Height> = None;
+        for (height, item) in inner.pending.range(low..=high) {
+            if let Some(expected) = next_expected {
+                if *height != expected {
+                    break;
+                }
+            }
+
+            let next_estimated_bytes = estimated_bytes.saturating_add(item.estimated_bytes);
+            if !taken.is_empty() && next_estimated_bytes > max_estimated_bytes {
+                break;
+            }
+
+            taken.push((*height, *item));
+            estimated_bytes = next_estimated_bytes;
+            if taken.len() >= max_count {
+                break;
+            }
+            // Stop the run at the end of the height space rather than overflowing.
+            match height.0.checked_add(1) {
+                Some(raw) => next_expected = Some(block::Height(raw)),
+                None => break,
+            }
+        }
+        for (height, item) in &taken {
+            inner.pending.remove(height);
+            inner.in_flight.insert(*height, *item);
+        }
+        taken
+    }
+
     /// Move each given height `in_flight → pending`, preserving its stored
     /// [`WorkItem`]. Heights not currently `in_flight` are skipped (idempotent).
     /// Wakes waiters if anything moved.

zebra-state/src/request.rs

@@ -1496,8 +1496,8 @@ pub enum ReadRequest {
 
     /// Returns scheduling-only body-size hints for a contiguous height range.
     ///
-    /// Confirmed committed block sizes win over untrusted advertised header
-    /// hints. Unknown advertised sizes are returned as `None`.
+    /// Only confirmed committed block sizes are returned. Unknown sizes are
+    /// returned as `None`.
     BlockSizeHints {
         /// First height to read.
         from: block::Height,