diff --git a/zebra-network/src/zakura/block_sync/peer_routine.rs b/zebra-network/src/zakura/block_sync/peer_routine.rs index 7b187c58941..3162a5f622d 100644 --- a/zebra-network/src/zakura/block_sync/peer_routine.rs +++ b/zebra-network/src/zakura/block_sync/peer_routine.rs @@ -547,7 +547,7 @@ 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), @@ -555,39 +555,35 @@ impl PeerRoutine { .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); diff --git a/zebra-network/src/zakura/block_sync/tests.rs b/zebra-network/src/zakura/block_sync/tests.rs index e1df312d1b0..21d28fcb057 100644 --- a/zebra-network/src/zakura/block_sync/tests.rs +++ b/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![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![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![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![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(); diff --git a/zebra-network/src/zakura/block_sync/work_queue.rs b/zebra-network/src/zakura/block_sync/work_queue.rs index 3041f2cb677..186228d996b 100644 --- a/zebra-network/src/zakura/block_sync/work_queue.rs +++ b/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 = 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. diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs index 751e35e7f87..140b9bcda88 100644 --- a/zebra-state/src/service.rs +++ b/zebra-state/src/service.rs @@ -1517,7 +1517,7 @@ where finalized_state::serve_block_roots(db, height..=height) .into_iter() .next() - // If the height is in the chain, serve the roots from the chain + // If the height is in the chain, serve the roots from the chain } else if let Some(chain) = chain .as_ref() .map(|chain| chain.as_ref()) diff --git a/zebra-state/src/service/finalized_state/zebra_db/block.rs b/zebra-state/src/service/finalized_state/zebra_db/block.rs index d8a9fd7a270..75598a299c1 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/block.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/block.rs @@ -2079,10 +2079,21 @@ impl DiskWriteBatch { for (index, (height, hash, header, body_size)) in validated_headers.into_iter().enumerate() { + let same_header = zebra_db.zakura_header_hash(height) == Some(hash); + let advertised_body_size = match ( + same_header, + zebra_db.advertised_body_size(height), + AdvertisedBodySize::new(body_size).map(AdvertisedBodySize::get), + ) { + (true, existing, Some(new)) => Some(existing.unwrap_or(0).max(new)), + (true, existing, None) => existing, + (false, _existing, new) => new, + }; + self.zs_insert(&header_by_height, height, header); self.zs_insert(&hash_by_height, height, hash); self.zs_insert(&height_by_hash, hash, height); - if let Some(body_size) = AdvertisedBodySize::new(body_size) { + if let Some(body_size) = advertised_body_size.and_then(AdvertisedBodySize::new) { self.zs_insert(&body_size_by_height, height, body_size); } else { self.zs_delete(&body_size_by_height, height); diff --git a/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs b/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs index 289c76d837c..2c412f9488e 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs @@ -155,6 +155,87 @@ fn header_range_commit_stores_advertised_body_sizes_with_zero_as_unknown() { assert_eq!(state.advertised_body_size(Height(2)), None); } +#[test] +fn header_range_commit_merges_same_header_advertised_body_size_by_max() { + let _init_guard = zebra_test::init(); + let (state, genesis, block1) = mainnet_state_with_genesis(); + + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch( + &state, + genesis.hash(), + std::slice::from_ref(&block1.header), + &[123_456], + ) + .expect("block 1 header links to genesis and has valid context"); + state.write_batch(batch).expect("header batch writes"); + + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch( + &state, + genesis.hash(), + std::slice::from_ref(&block1.header), + &[0], + ) + .expect("same block 1 header can refresh its advertised body size"); + state.write_batch(batch).expect("header batch writes"); + assert_eq!(state.advertised_body_size(Height(1)), Some(123_456)); + + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch( + &state, + genesis.hash(), + std::slice::from_ref(&block1.header), + &[999_999], + ) + .expect("same block 1 header can refresh its advertised body size"); + state.write_batch(batch).expect("header batch writes"); + assert_eq!(state.advertised_body_size(Height(1)), Some(999_999)); + + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch( + &state, + genesis.hash(), + std::slice::from_ref(&block1.header), + &[100], + ) + .expect("same block 1 header can refresh its advertised body size"); + state.write_batch(batch).expect("header batch writes"); + assert_eq!(state.advertised_body_size(Height(1)), Some(999_999)); +} + +#[test] +fn header_range_reorg_resets_advertised_body_sizes() { + let _init_guard = zebra_test::init(); + let genesis = mainnet_block(0); + let network = no_extra_checkpoint_test_network(genesis.hash()); + let state = state_with_genesis(&network, genesis.clone()); + + let original = synthetic_headers_from_state(&state, Height(0), genesis.hash(), 2, 1); + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch(&state, genesis.hash(), &original, &[111, 222]) + .expect("original synthetic headers are valid"); + state.write_batch(batch).expect("header batch writes"); + assert_eq!(state.advertised_body_size(Height(1)), Some(111)); + assert_eq!(state.advertised_body_size(Height(2)), Some(222)); + + let replacement = synthetic_headers_from_state(&state, Height(0), genesis.hash(), 3, 9); + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch(&state, genesis.hash(), &replacement, &[0, 0, 333]) + .expect("higher-work replacement synthetic headers are valid"); + state.write_batch(batch).expect("header batch writes"); + + assert_eq!(state.advertised_body_size(Height(1)), None); + assert_eq!(state.advertised_body_size(Height(2)), None); + assert_eq!(state.advertised_body_size(Height(3)), Some(333)); +} + #[test] fn block_size_hints_prefer_confirmed_block_info_over_advertised_hint() { let _init_guard = zebra_test::init(); @@ -198,6 +279,35 @@ fn block_size_hints_prefer_confirmed_block_info_over_advertised_hint() { ); } +#[test] +fn block_size_hints_use_advertised_hints_when_unconfirmed() { + let _init_guard = zebra_test::init(); + let (state, genesis, block1) = mainnet_state_with_genesis(); + + let mut batch = DiskWriteBatch::new(); + batch + .prepare_header_range_batch( + &state, + genesis.hash(), + std::slice::from_ref(&block1.header), + &[999_999], + ) + .expect("block 1 header links to genesis and has valid context"); + state + .write_batch(batch) + .expect("header range batch writes successfully"); + + assert_eq!( + crate::service::read::block_size_hints( + None::>, + &state, + Height(1), + 1, + ), + vec![(Height(1), Some(999_999))], + ); +} + #[test] fn header_range_read_is_contiguous_capped_and_stops_at_first_gap() { let _init_guard = zebra_test::init(); diff --git a/zebra-state/src/service/read/block.rs b/zebra-state/src/service/read/block.rs index 11b4fe5defe..4d7f9bd91a8 100644 --- a/zebra-state/src/service/read/block.rs +++ b/zebra-state/src/service/read/block.rs @@ -387,14 +387,14 @@ where let Some(height) = from.0.checked_add(offset).map(Height) else { break; }; - let confirmed_size = chain + let size_hint = chain .as_ref() .and_then(|chain| chain.as_ref().block_info(height.into())) .or_else(|| db.block_info(height.into())) - .map(|info| info.size()); - let size = confirmed_size.or_else(|| db.advertised_body_size(height)); + .map(|info| info.size()) + .or_else(|| db.advertised_body_size(height)); - hints.push((height, size)); + hints.push((height, size_hint)); } hints