From af5a3fb64d5c70f1603378184f9fb3609a32912b Mon Sep 17 00:00:00 2001 From: Waclaw Banasik Date: Tue, 10 Dec 2024 18:28:41 +0000 Subject: [PATCH] fix(congestion_control) - proper protocol upgrade handling for congestion seed (#12594) Fixing unhandled protocol upgrade introduced in #12555. --- runtime/runtime/src/lib.rs | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/runtime/runtime/src/lib.rs b/runtime/runtime/src/lib.rs index 5b68843d614..66c9b95290d 100644 --- a/runtime/runtime/src/lib.rs +++ b/runtime/runtime/src/lib.rs @@ -2083,18 +2083,27 @@ impl Runtime { let mut own_congestion_info = receipt_sink.own_congestion_info(); if let Some(congestion_info) = &mut own_congestion_info { delayed_receipts.apply_congestion_changes(congestion_info)?; - let shard_layout = epoch_info_provider.shard_layout(&apply_state.epoch_id)?; - let shard_ids = shard_layout.shard_ids().collect_vec(); - let shard_index = shard_layout - .get_shard_index(apply_state.shard_id) - .map_err(Into::::into)? - .try_into() - .expect("Shard Index must fit within u64"); - - let congestion_seed = apply_state.block_height.wrapping_add(shard_index); + let protocol_version = apply_state.current_protocol_version; + + let (all_shards, shard_seed) = + if ProtocolFeature::SimpleNightshadeV4.enabled(protocol_version) { + let shard_layout = epoch_info_provider.shard_layout(&apply_state.epoch_id)?; + let shard_ids = shard_layout.shard_ids().collect_vec(); + let shard_index = shard_layout + .get_shard_index(apply_state.shard_id) + .map_err(Into::::into)? + .try_into() + .expect("Shard Index must fit within u64"); + + (shard_ids, shard_index) + } else { + (apply_state.congestion_info.all_shards(), apply_state.shard_id.into()) + }; + + let congestion_seed = apply_state.block_height.wrapping_add(shard_seed); congestion_info.finalize_allowed_shard( apply_state.shard_id, - shard_ids.as_slice(), + &all_shards, congestion_seed, ); }