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

@@ -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<LinkAllowance>,

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);

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<Bandwidth> = 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<LinkAllowance> = 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

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);
 

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