feat: wiring for bandwidth scheduler

chain/chain-primitives/src/error.rs

@@ -185,6 +185,9 @@ pub enum Error {
     /// Invalid Congestion Info
     #[error("Invalid Congestion Info: {0}")]
     InvalidCongestionInfo(String),
+    /// Invalid bandwidth requests
+    #[error("Invalid bandwidth requests - chunk extra doesn't match chunk header: {0}")]
+    InvalidBandwidthRequests(String),
     /// Invalid shard id
     #[error("Shard id {0} does not exist")]
     InvalidShardId(ShardId),
@@ -315,6 +318,7 @@ impl Error {
             | Error::InvalidGasUsed
             | Error::InvalidBalanceBurnt
             | Error::InvalidCongestionInfo(_)
+            | Error::InvalidBandwidthRequests(_)
             | Error::InvalidShardId(_)
             | Error::InvalidStateRequest(_)
             | Error::InvalidRandomnessBeaconOutput
@@ -393,6 +397,7 @@ impl Error {
             Error::InvalidGasUsed => "invalid_gas_used",
             Error::InvalidBalanceBurnt => "invalid_balance_burnt",
             Error::InvalidCongestionInfo(_) => "invalid_congestion_info",
+            Error::InvalidBandwidthRequests(_) => "invalid_bandwidth_requests",
             Error::InvalidShardId(_) => "invalid_shard_id",
             Error::InvalidStateRequest(_) => "invalid_state_request",
             Error::InvalidRandomnessBeaconOutput => "invalid_randomness_beacon_output",

chain/chain/src/chain.rs

@@ -47,6 +47,7 @@ use near_chain_configs::{MutableConfigValue, MutableValidatorSigner};
 use near_chain_primitives::error::{BlockKnownError, Error, LogTransientStorageError};
 use near_epoch_manager::shard_tracker::ShardTracker;
 use near_epoch_manager::EpochManagerAdapter;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::{genesis_chunks, Block, BlockValidityError, Tip};
 use near_primitives::block_header::BlockHeader;
 use near_primitives::challenge::{
@@ -621,6 +622,7 @@ impl Chain {
             gas_limit,
             0,
             congestion_info,
+            BandwidthRequests::default_for_protocol_version(genesis_protocol_version),
         )
     }
 
@@ -3209,7 +3211,12 @@ impl Chain {
         };
         let congestion_info = block.block_congestion_info();
 
-        Ok(ApplyChunkBlockContext::from_header(block_header, gas_price, congestion_info))
+        Ok(ApplyChunkBlockContext::from_header(
+            block_header,
+            gas_price,
+            congestion_info,
+            block.block_bandwidth_requests(),
+        ))
     }
 
     fn block_catch_up_postprocess(

chain/chain/src/chain_update.rs

@@ -85,6 +85,7 @@ impl<'a> ChainUpdate<'a> {
         should_save_state_transition_data: bool,
     ) -> Result<(), Error> {
         let _span = tracing::debug_span!(target: "chain", "apply_chunk_postprocessing", height=block.header().height()).entered();
+        Self::bandwidth_scheduler_state_sanity_check(&apply_results);
         for result in apply_results {
             self.process_apply_chunk_result(block, result, should_save_state_transition_data)?;
         }
@@ -123,6 +124,7 @@ impl<'a> ChainUpdate<'a> {
                         gas_limit,
                         apply_result.total_balance_burnt,
                         apply_result.congestion_info,
+                        apply_result.bandwidth_requests,
                     ),
                 );
 
@@ -193,6 +195,24 @@ impl<'a> ChainUpdate<'a> {
         Ok(())
     }
 
+    /// Extra sanity check for bandwdith scheduler - the scheduler state should be the same on all shards.
+    fn bandwidth_scheduler_state_sanity_check(apply_results: &[ShardUpdateResult]) {
+        let state_hashes: Vec<CryptoHash> = apply_results
+            .iter()
+            .map(|r| match r {
+                ShardUpdateResult::NewChunk(new_res) => {
+                    new_res.apply_result.bandwidth_scheduler_state_hash
+                }
+                ShardUpdateResult::OldChunk(old_res) => {
+                    old_res.apply_result.bandwidth_scheduler_state_hash
+                }
+            })
+            .collect();
+        for hash in &state_hashes {
+            assert_eq!(*hash, state_hashes[0]);
+        }
+    }
+
     /// This is the last step of process_block_single, where we take the preprocess block info
     /// apply chunk results and store the results on chain.
     #[tracing::instrument(
@@ -527,6 +547,7 @@ impl<'a> ChainUpdate<'a> {
                 challenges_result: block_header.challenges_result().clone(),
                 random_seed: *block_header.random_value(),
                 congestion_info: block.block_congestion_info(),
+                bandwidth_requests: block.block_bandwidth_requests(),
             },
             &receipts,
             chunk.transactions(),
@@ -562,6 +583,7 @@ impl<'a> ChainUpdate<'a> {
             gas_limit,
             apply_result.total_balance_burnt,
             apply_result.congestion_info,
+            apply_result.bandwidth_requests,
         );
         self.chain_store_update.save_chunk_extra(block_header.hash(), &shard_uid, chunk_extra);
 
@@ -633,6 +655,7 @@ impl<'a> ChainUpdate<'a> {
                 &block_header,
                 prev_block_header.next_gas_price(),
                 block.block_congestion_info(),
+                block.block_bandwidth_requests(),
             ),
             &[],
             &[],

chain/chain/src/runtime/mod.rs

@@ -269,6 +269,7 @@ impl NightshadeRuntime {
             challenges_result,
             random_seed,
             congestion_info,
+            bandwidth_requests,
         } = block;
         let ApplyChunkShardContext {
             shard_id,
@@ -390,6 +391,7 @@ impl NightshadeRuntime {
                 is_first_block_with_chunk_of_version,
             },
             congestion_info,
+            bandwidth_requests,
         };
 
         let instant = Instant::now();
@@ -467,6 +469,8 @@ impl NightshadeRuntime {
             applied_receipts_hash: hash(&borsh::to_vec(receipts).unwrap()),
             congestion_info: apply_result.congestion_info,
             contract_accesses: apply_result.contract_accesses,
+            bandwidth_requests: apply_result.bandwidth_requests,
+            bandwidth_scheduler_state_hash: apply_result.bandwidth_scheduler_state_hash,
         };
 
         Ok(result)

chain/chain/src/runtime/tests.rs

@@ -10,6 +10,7 @@ use near_pool::{
 };
 use near_primitives::action::FunctionCallAction;
 use near_primitives::apply::ApplyChunkReason;
+use near_primitives::bandwidth_scheduler::BlockBandwidthRequests;
 use near_primitives::congestion_info::{BlockCongestionInfo, ExtendedCongestionInfo};
 use near_primitives::epoch_block_info::BlockInfo;
 use near_primitives::receipt::{ActionReceipt, ReceiptV1};
@@ -214,6 +215,7 @@ impl TestEnv {
     ) -> ApplyChunkResult {
         // TODO(congestion_control): pass down prev block info and read congestion info from there
         // For now, just use default.
+        // TODO(bandwidth_scheduler) - pass bandwidth requests from prev_block
         let prev_block_hash = self.head.last_block_hash;
         let epoch_id = self.epoch_manager.get_epoch_id_from_prev_block(&prev_block_hash).unwrap();
         let shard_layout = self.epoch_manager.get_shard_layout(&epoch_id).unwrap();
@@ -256,6 +258,7 @@ impl TestEnv {
                     challenges_result,
                     random_seed: CryptoHash::default(),
                     congestion_info,
+                    bandwidth_requests: BlockBandwidthRequests::empty(),
                 },
                 receipts,
                 transactions,

chain/chain/src/stateless_validation/chunk_validation.rs

@@ -587,6 +587,7 @@ pub fn apply_result_to_chunk_extra(
         chunk.gas_limit(),
         apply_result.total_balance_burnt,
         apply_result.congestion_info,
+        apply_result.bandwidth_requests,
     )
 }
 

chain/chain/src/test_utils/kv_runtime.rs

@@ -16,6 +16,7 @@ use near_parameters::RuntimeConfig;
 use near_pool::types::TransactionGroupIterator;
 use near_primitives::account::{AccessKey, Account};
 use near_primitives::apply::ApplyChunkReason;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::Tip;
 use near_primitives::block_header::{Approval, ApprovalInner};
 use near_primitives::congestion_info::{CongestionInfo, ExtendedCongestionInfo};
@@ -1382,6 +1383,8 @@ impl RuntimeAdapter for KeyValueRuntime {
             congestion_info: Self::get_congestion_info(PROTOCOL_VERSION),
             // Since all actions are transfer actions, there is no contracts accessed.
             contract_accesses: vec![],
+            bandwidth_requests: BandwidthRequests::default_for_protocol_version(PROTOCOL_VERSION),
+            bandwidth_scheduler_state_hash: CryptoHash::default(),
         })
     }
 

chain/chain/src/tests/simple_chain.rs

@@ -32,7 +32,7 @@ fn build_chain() {
     //     cargo insta test --accept -p near-chain --features nightly -- tests::simple_chain::build_chain
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_snapshot!(hash, @"Hc3bWEd7ikHf9BAe2SknvH2jAAakEtBRU1FBu6Udocm3");
+        insta::assert_snapshot!(hash, @"2V6auEJDpFWUadSMYwAUg68tn9KmoMmHw2JnHizYicwc");
     } else {
         insta::assert_snapshot!(hash, @"GHZFAFiMdGzAfnWTcS9u9wqFvxMrgFpyEr6Use7jk2Lo");
     }
@@ -50,7 +50,7 @@ fn build_chain() {
 
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_snapshot!(hash, @"39R6bDFXkPfwdYs4crV3RyCde85ecycqP5DBwdtwyjcJ");
+        insta::assert_snapshot!(hash, @"8yW4usbwYcRDKmKmDkHTFpzggDZBGu6avKLu1iTf4Lr6");
     } else {
         insta::assert_snapshot!(hash, @"3Pdm44L71Bk8EokPHF1pxakHojsriNadBdZZSpcoDv9q");
     }

chain/chain/src/types.rs

@@ -9,6 +9,8 @@ pub use near_epoch_manager::EpochManagerAdapter;
 use near_parameters::RuntimeConfig;
 use near_pool::types::TransactionGroupIterator;
 use near_primitives::apply::ApplyChunkReason;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
+use near_primitives::bandwidth_scheduler::BlockBandwidthRequests;
 pub use near_primitives::block::{Block, BlockHeader, Tip};
 use near_primitives::challenge::{ChallengesResult, PartialState};
 use near_primitives::checked_feature;
@@ -104,6 +106,11 @@ pub struct ApplyChunkResult {
     pub congestion_info: Option<CongestionInfo>,
     /// Hashes of the contracts accessed while applying the chunk.
     pub contract_accesses: Vec<CodeHash>,
+    /// Requests for bandwidth to send receipts to other shards.
+    /// Will be None for protocol versions that don't have the BandwidthScheduler feature enabled.
+    pub bandwidth_requests: Option<BandwidthRequests>,
+    /// Used only for a sanity check.
+    pub bandwidth_scheduler_state_hash: CryptoHash,
 }
 
 impl ApplyChunkResult {
@@ -298,13 +305,15 @@ pub struct ApplyChunkBlockContext {
     pub challenges_result: ChallengesResult,
     pub random_seed: CryptoHash,
     pub congestion_info: BlockCongestionInfo,
+    pub bandwidth_requests: BlockBandwidthRequests,
 }
 
 impl ApplyChunkBlockContext {
     pub fn from_header(
         header: &BlockHeader,
         gas_price: Balance,
         congestion_info: BlockCongestionInfo,
+        bandwidth_requests: BlockBandwidthRequests,
     ) -> Self {
         Self {
             height: header.height(),
@@ -315,6 +324,7 @@ impl ApplyChunkBlockContext {
             challenges_result: header.challenges_result().clone(),
             random_seed: *header.random_value(),
             congestion_info,
+            bandwidth_requests,
         }
     }
 }

chain/chain/src/validate.rs

@@ -4,6 +4,7 @@ use borsh::BorshDeserialize;
 
 use near_crypto::PublicKey;
 use near_epoch_manager::EpochManagerAdapter;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::{Block, BlockHeader};
 use near_primitives::challenge::{
     BlockDoubleSign, Challenge, ChallengeBody, ChunkProofs, ChunkState, MaybeEncodedShardChunk,
@@ -178,6 +179,10 @@ pub fn validate_chunk_with_chunk_extra_and_receipts_root(
     }
 
     validate_congestion_info(&prev_chunk_extra.congestion_info(), &chunk_header.congestion_info())?;
+    validate_bandwidth_requests(
+        prev_chunk_extra.bandwidth_requests(),
+        chunk_header.bandwidth_requests(),
+    )?;
 
     Ok(())
 }
@@ -211,6 +216,41 @@ fn validate_congestion_info(
     }
 }
 
+fn validate_bandwidth_requests(
+    extra_bandwidth_requests: Option<&BandwidthRequests>,
+    header_bandwidth_requests: Option<&BandwidthRequests>,
+) -> Result<(), Error> {
+    if extra_bandwidth_requests.is_none()
+        && header_bandwidth_requests == Some(&BandwidthRequests::empty())
+    {
+        // This corner case happens for the first chunk that has the BandwidthScheduler feature enabled.
+        // The previous chunk was applied with a protocol version which doesn't have bandwidth scheduler
+        // enabled and because of that the bandwidth requests in ChunkExtra are None.
+        // The header was produced in the new protocol version, and the newer version of header always
+        // has some bandwidth requests, it's not an `Option`. Because of that the header requests are `Some(BandwidthRequests::empty())`.
+        return Ok(());
+    }
+
+    if extra_bandwidth_requests != header_bandwidth_requests {
+        fn requests_len(requests_opt: Option<&BandwidthRequests>) -> usize {
+            match requests_opt {
+                Some(BandwidthRequests::V1(requests_v1)) => requests_v1.requests.len(),
+                None => 0,
+            }
+        }
+        let error_info_str = format!(
+            "chunk extra: (is_some: {}, len: {}) chunk header: (is_some: {}, len: {})",
+            extra_bandwidth_requests.is_some(),
+            requests_len(extra_bandwidth_requests),
+            header_bandwidth_requests.is_some(),
+            requests_len(header_bandwidth_requests)
+        );
+        return Err(Error::InvalidBandwidthRequests(error_info_str));
+    }
+
+    Ok(())
+}
+
 /// Validates a double sign challenge.
 /// Only valid if ancestors of both blocks are present in the chain.
 fn validate_double_sign(

chain/chunks/src/shards_manager_actor.rs

@@ -107,6 +107,7 @@ use near_network::types::{
 };
 use near_network::types::{NetworkRequests, PeerManagerMessageRequest};
 use near_performance_metrics_macros::perf;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::Tip;
 use near_primitives::congestion_info::CongestionInfo;
 use near_primitives::errors::EpochError;
@@ -1994,6 +1995,7 @@ impl ShardsManagerActor {
         prev_outgoing_receipts_root: CryptoHash,
         tx_root: CryptoHash,
         congestion_info: Option<CongestionInfo>,
+        bandwidth_requests: Option<BandwidthRequests>,
         signer: &ValidatorSigner,
         rs: &ReedSolomon,
         protocol_version: ProtocolVersion,
@@ -2014,6 +2016,7 @@ impl ShardsManagerActor {
             prev_outgoing_receipts,
             prev_outgoing_receipts_root,
             congestion_info,
+            bandwidth_requests,
             signer,
             protocol_version,
         )

chain/chunks/src/test_utils.rs

@@ -6,6 +6,7 @@ use near_epoch_manager::test_utils::setup_epoch_manager_with_block_and_chunk_pro
 use near_epoch_manager::EpochManagerHandle;
 use near_network::shards_manager::ShardsManagerRequestFromNetwork;
 use near_network::test_utils::MockPeerManagerAdapter;
+use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::congestion_info::CongestionInfo;
 use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::{self, MerklePath};
@@ -153,6 +154,7 @@ impl ChunkTestFixture {
             receipts_root,
             MerkleHash::default(),
             congestion_info,
+            BandwidthRequests::default_for_protocol_version(PROTOCOL_VERSION),
             &signer,
             &rs,
             PROTOCOL_VERSION,

chain/client/src/chunk_distribution_network.rs

@@ -220,11 +220,12 @@ mod tests {
         time::Clock,
     };
     use near_primitives::{
+        bandwidth_scheduler::BandwidthRequests,
         congestion_info::CongestionInfo,
         hash::hash,
         sharding::{
-            PartialEncodedChunkV2, ShardChunkHeaderInner, ShardChunkHeaderInnerV3,
-            ShardChunkHeaderV3,
+            shard_chunk_header_inner::ShardChunkHeaderInnerV4, PartialEncodedChunkV2,
+            ShardChunkHeaderInner, ShardChunkHeaderV3,
         },
         validator_signer::EmptyValidatorSigner,
     };
@@ -398,7 +399,7 @@ mod tests {
         let mut mock_hashes = MockHashes::new(prev_block_hash);
 
         let signer = EmptyValidatorSigner::default().into();
-        let header_inner = ShardChunkHeaderInner::V3(ShardChunkHeaderInnerV3 {
+        let header_inner = ShardChunkHeaderInner::V4(ShardChunkHeaderInnerV4 {
             prev_block_hash,
             prev_state_root: mock_hashes.next().unwrap(),
             prev_outcome_root: mock_hashes.next().unwrap(),
@@ -413,6 +414,7 @@ mod tests {
             tx_root: mock_hashes.next().unwrap(),
             prev_validator_proposals: Vec::new(),
             congestion_info: CongestionInfo::default(),
+            bandwidth_requests: BandwidthRequests::empty(),
         });
         let header = ShardChunkHeaderV3::from_inner(header_inner, &signer);
         PartialEncodedChunk::V2(PartialEncodedChunkV2 {

chain/client/src/client.rs

@@ -924,6 +924,7 @@ impl Client {
             outgoing_receipts_root,
             tx_root,
             congestion_info,
+            chunk_extra.bandwidth_requests().cloned(),
             &*validator_signer,
             &mut self.rs_for_chunk_production,
             protocol_version,