From e92fef24746c3ecfb5c0f5f23ba8bca93e8a5db7 Mon Sep 17 00:00:00 2001 From: Jan Malinowski <149345204+jancionear@users.noreply.github.com> Date: Tue, 7 Jan 2025 15:43:47 +0000 Subject: [PATCH] feat(bandwidth_scheduler) - make BandwidthSchedulerState versioned (#12694) Change this struct to a versioned enum to make future changes easier. --- core/primitives/src/bandwidth_scheduler.rs | 8 ++++++-- .../runtime/src/bandwidth_scheduler/mod.rs | 20 ++++++++++++------- .../src/bandwidth_scheduler/scheduler.rs | 18 ++++++++++++----- .../src/bandwidth_scheduler/simulator.rs | 6 +++--- .../res/protocol_schema.toml | 3 ++- 5 files changed, 37 insertions(+), 18 deletions(-) diff --git a/core/primitives/src/bandwidth_scheduler.rs b/core/primitives/src/bandwidth_scheduler.rs index db57cce6725..18df48b32ad 100644 --- a/core/primitives/src/bandwidth_scheduler.rs +++ b/core/primitives/src/bandwidth_scheduler.rs @@ -302,9 +302,13 @@ impl BlockBandwidthRequests { /// The state should be the same on all shards. All shards start with the same state /// and apply the same bandwidth scheduler algorithm at the same heights, so the resulting /// scheduler state stays the same. -/// TODO(bandwidth_scheduler) - make this struct versioned. #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq, Eq, ProtocolSchema)] -pub struct BandwidthSchedulerState { +pub enum BandwidthSchedulerState { + V1(BandwidthSchedulerStateV1), +} + +#[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq, Eq, ProtocolSchema)] +pub struct BandwidthSchedulerStateV1 { /// Allowance for every pair of (sender, receiver). Used in the scheduler algorithm. /// Bandwidth scheduler updates the allowances on every run. pub link_allowances: Vec, diff --git a/runtime/runtime/src/bandwidth_scheduler/mod.rs b/runtime/runtime/src/bandwidth_scheduler/mod.rs index 72ea9a4d7cf..6a33e91bf5f 100644 --- a/runtime/runtime/src/bandwidth_scheduler/mod.rs +++ b/runtime/runtime/src/bandwidth_scheduler/mod.rs @@ -1,7 +1,9 @@ use std::collections::BTreeMap; use std::num::NonZeroU64; -use near_primitives::bandwidth_scheduler::{BandwidthSchedulerParams, BandwidthSchedulerState}; +use near_primitives::bandwidth_scheduler::{ + BandwidthSchedulerParams, BandwidthSchedulerState, BandwidthSchedulerStateV1, +}; use near_primitives::congestion_info::CongestionControl; use near_primitives::errors::RuntimeError; use near_primitives::hash::{hash, CryptoHash}; @@ -48,10 +50,10 @@ pub fn run_bandwidth_scheduler( Some(prev_state) => prev_state, None => { tracing::debug!(target: "runtime", "Bandwidth scheduler state not found - initializing"); - BandwidthSchedulerState { + BandwidthSchedulerState::V1(BandwidthSchedulerStateV1 { link_allowances: Vec::new(), sanity_check_hash: CryptoHash::default(), - } + }) } }; @@ -103,10 +105,14 @@ pub fn run_bandwidth_scheduler( // This is a sanity check to make sure that all shards run the scheduler with the same inputs. // It would be a bit nicer to hash all inputs, but that could be slow and the serialization // format of the hashed structs would become part of the protocol. - let mut sanity_check_bytes = Vec::new(); - sanity_check_bytes.extend_from_slice(scheduler_state.sanity_check_hash.as_ref()); - sanity_check_bytes.extend_from_slice(CryptoHash::hash_borsh(&all_shards).as_ref()); - scheduler_state.sanity_check_hash = CryptoHash::hash_bytes(&sanity_check_bytes); + match &mut scheduler_state { + BandwidthSchedulerState::V1(scheduler_state) => { + let mut sanity_check_bytes = Vec::new(); + sanity_check_bytes.extend_from_slice(scheduler_state.sanity_check_hash.as_ref()); + sanity_check_bytes.extend_from_slice(CryptoHash::hash_borsh(&all_shards).as_ref()); + scheduler_state.sanity_check_hash = CryptoHash::hash_bytes(&sanity_check_bytes); + } + }; // Save the updated scheduler state to the trie. set_bandwidth_scheduler_state(state_update, &scheduler_state); diff --git a/runtime/runtime/src/bandwidth_scheduler/scheduler.rs b/runtime/runtime/src/bandwidth_scheduler/scheduler.rs index 19bfb3ef044..3ef961e4b24 100644 --- a/runtime/runtime/src/bandwidth_scheduler/scheduler.rs +++ b/runtime/runtime/src/bandwidth_scheduler/scheduler.rs @@ -144,7 +144,8 @@ use std::collections::{BTreeMap, VecDeque}; use near_primitives::bandwidth_scheduler::{ Bandwidth, BandwidthRequest, BandwidthRequestValues, BandwidthRequests, - BandwidthSchedulerParams, BandwidthSchedulerState, BlockBandwidthRequests, LinkAllowance, + BandwidthSchedulerParams, BandwidthSchedulerState, BandwidthSchedulerStateV1, + BlockBandwidthRequests, LinkAllowance, }; use near_primitives::shard_layout::ShardLayout; use near_primitives::types::{ShardId, ShardIndex}; @@ -209,6 +210,10 @@ impl BandwidthScheduler { return GrantedBandwidth { granted: BTreeMap::new() }; } + let state = match state { + BandwidthSchedulerState::V1(v1) => v1, + }; + // Convert link allowances to the internal representation. let mut link_allowances: ShardLinkMap = ShardLinkMap::new(&shard_layout); for link_allowance in &state.link_allowances { @@ -504,7 +509,7 @@ impl BandwidthScheduler { /// Update the persistent scheduler state after running the scheduler algorithm. /// This state is persisted in the trie between runs. - fn update_scheduler_state(&self, state: &mut BandwidthSchedulerState) { + fn update_scheduler_state(&self, state: &mut BandwidthSchedulerStateV1) { let mut new_state_allowances: Vec = Vec::new(); for link in self.iter_links() { @@ -696,7 +701,8 @@ mod tests { use near_primitives::bandwidth_scheduler::{ BandwidthRequest, BandwidthRequestBitmap, BandwidthRequests, BandwidthRequestsV1, - BandwidthSchedulerParams, BandwidthSchedulerState, BlockBandwidthRequests, LinkAllowance, + BandwidthSchedulerParams, BandwidthSchedulerState, BandwidthSchedulerStateV1, + BlockBandwidthRequests, LinkAllowance, }; use near_primitives::hash::CryptoHash; use near_primitives::shard_layout::ShardLayout; @@ -732,8 +738,10 @@ mod tests { next_allowance -= 1; } } - let mut scheduler_state = - BandwidthSchedulerState { link_allowances, sanity_check_hash: CryptoHash::default() }; + let mut scheduler_state = BandwidthSchedulerState::V1(BandwidthSchedulerStateV1 { + link_allowances, + sanity_check_hash: CryptoHash::default(), + }); // Shards are not congested, scheduler can grant as many requests as possible. let shards_status = shard_layout diff --git a/runtime/runtime/src/bandwidth_scheduler/simulator.rs b/runtime/runtime/src/bandwidth_scheduler/simulator.rs index 27466c3a221..4c6d19f248f 100644 --- a/runtime/runtime/src/bandwidth_scheduler/simulator.rs +++ b/runtime/runtime/src/bandwidth_scheduler/simulator.rs @@ -16,7 +16,7 @@ use borsh::BorshSerialize; use bytesize::ByteSize; use near_primitives::bandwidth_scheduler::{ BandwidthRequest, BandwidthRequests, BandwidthRequestsV1, BandwidthSchedulerParams, - BandwidthSchedulerState, BlockBandwidthRequests, + BandwidthSchedulerState, BandwidthSchedulerStateV1, BlockBandwidthRequests, }; use near_primitives::hash::CryptoHash; use near_primitives::shard_layout::ShardLayout; @@ -177,10 +177,10 @@ impl ChainSimulator { } let mut shard_state = shard_state_opt.unwrap_or_else(|| ShardState { buffered_outgoing_receipts: BTreeMap::new(), - scheduler_state: BandwidthSchedulerState { + scheduler_state: BandwidthSchedulerState::V1(BandwidthSchedulerStateV1 { link_allowances: Vec::new(), sanity_check_hash: CryptoHash::default(), - }, + }), }); let pre_state_hash = CryptoHash::hash_borsh(&shard_state); diff --git a/tools/protocol-schema-check/res/protocol_schema.toml b/tools/protocol-schema-check/res/protocol_schema.toml index a280f8e61dd..dcb30b62353 100644 --- a/tools/protocol-schema-check/res/protocol_schema.toml +++ b/tools/protocol-schema-check/res/protocol_schema.toml @@ -20,7 +20,8 @@ BandwidthRequest = 234831851 BandwidthRequestBitmap = 2138002689 BandwidthRequests = 984876287 BandwidthRequestsV1 = 3810915065 -BandwidthSchedulerState = 3401315484 +BandwidthSchedulerState = 2982803600 +BandwidthSchedulerStateV1 = 34546280 BitArray = 3709965115 Block = 3541579558 BlockBody = 206872245