feat(congestion_control) - handling missing chunks

chain/chain/src/chain_update.rs

@@ -396,7 +396,6 @@ impl<'a> ChainUpdate<'a> {
                 let old_extra = self.chain_store_update.get_chunk_extra(prev_hash, &shard_uid)?;
                 let mut new_extra = ChunkExtra::clone(&old_extra);
                 *new_extra.state_root_mut() = apply_result.new_root;
-                // TODO(congestion_control) handle missing chunks congestion info #11039
 
                 let flat_storage_manager = self.runtime_adapter.get_flat_storage_manager();
                 let store_update = flat_storage_manager.save_flat_state_changes(
@@ -906,7 +905,6 @@ impl<'a> ChainUpdate<'a> {
         // extra and apply changes to it.
         let mut new_chunk_extra = ChunkExtra::clone(&chunk_extra);
         *new_chunk_extra.state_root_mut() = apply_result.new_root;
-        // TODO(congestion_control) handle missing chunks congestion info #11039
 
         self.chain_store_update.save_chunk_extra(block_header.hash(), &shard_uid, new_chunk_extra);
         Ok(true)

chain/chain/src/runtime/mod.rs

@@ -828,10 +828,10 @@ impl RuntimeAdapter for NightshadeRuntime {
                             &tx.transaction.receiver_id,
                             &epoch_id,
                         )?;
-                        if let Some(shard_congestion) =
+                        if let Some(congestion_info) =
                             prev_block.congestion_info.get(&receiving_shard)
                         {
-                            if !shard_congestion.shard_accepts_transactions() {
+                            if !congestion_info.shard_accepts_transactions() {
                                 tracing::trace!(target: "runtime", tx=?tx.get_hash(), "discarding transaction due to congestion");
                                 continue;
                             }

chain/chain/src/runtime/tests.rs

@@ -11,6 +11,7 @@ use near_pool::{
 };
 use near_primitives::apply::ApplyChunkReason;
 use near_primitives::checked_feature;
+use near_primitives::congestion_info::ExtendedCongestionInfo;
 use near_primitives::test_utils::create_test_signer;
 use near_primitives::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
 use near_primitives::version::PROTOCOL_VERSION;
@@ -83,14 +84,14 @@ impl NightshadeRuntime {
         let epoch_id =
             self.epoch_manager.get_epoch_id_from_prev_block(block_hash).unwrap_or_default();
         let protocol_version = self.epoch_manager.get_epoch_protocol_version(&epoch_id).unwrap();
-        let congestion_info_map: HashMap<ShardId, CongestionInfo> =
+        let congestion_info_map: HashMap<ShardId, ExtendedCongestionInfo> =
             if !ProtocolFeature::CongestionControl.enabled(protocol_version) {
                 HashMap::new()
             } else {
                 let shard_ids = self.epoch_manager.shard_ids(&epoch_id).unwrap();
                 shard_ids
                     .into_iter()
-                    .map(|shard_id| (shard_id, CongestionInfo::default()))
+                    .map(|shard_id| (shard_id, ExtendedCongestionInfo::default()))
                     .collect()
             };
         let mut result = self

chain/chain/src/types.rs

@@ -12,6 +12,7 @@ pub use near_primitives::block::{Block, BlockHeader, Tip};
 use near_primitives::challenge::{ChallengesResult, PartialState};
 use near_primitives::checked_feature;
 use near_primitives::congestion_info::CongestionInfo;
+use near_primitives::congestion_info::ExtendedCongestionInfo;
 use near_primitives::errors::InvalidTxError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::{merklize, MerklePath};
@@ -291,14 +292,14 @@ pub struct ApplyChunkBlockContext {
     pub gas_price: Balance,
     pub challenges_result: ChallengesResult,
     pub random_seed: CryptoHash,
-    pub congestion_info: HashMap<ShardId, CongestionInfo>,
+    pub congestion_info: HashMap<ShardId, ExtendedCongestionInfo>,
 }
 
 impl ApplyChunkBlockContext {
     pub fn from_header(
         header: &BlockHeader,
         gas_price: Balance,
-        congestion_info: HashMap<ShardId, CongestionInfo>,
+        congestion_info: HashMap<ShardId, ExtendedCongestionInfo>,
     ) -> Self {
         Self {
             height: header.height(),
@@ -348,7 +349,7 @@ pub struct PrepareTransactionsBlockContext {
     pub next_gas_price: Balance,
     pub height: BlockHeight,
     pub block_hash: CryptoHash,
-    pub congestion_info: HashMap<ShardId, CongestionInfo>,
+    pub congestion_info: HashMap<ShardId, ExtendedCongestionInfo>,
 }
 
 impl From<&Block> for PrepareTransactionsBlockContext {

core/primitives/src/block.rs

@@ -6,7 +6,7 @@ use crate::block_body::{BlockBody, BlockBodyV1, ChunkEndorsementSignatures};
 pub use crate::block_header::*;
 use crate::challenge::{Challenges, ChallengesResult};
 use crate::checked_feature;
-use crate::congestion_info::CongestionInfo;
+use crate::congestion_info::{CongestionInfo, ExtendedCongestionInfo};
 use crate::hash::{hash, CryptoHash};
 use crate::merkle::{merklize, verify_path, MerklePath};
 use crate::num_rational::Rational32;
@@ -591,15 +591,24 @@ impl Block {
         }
     }
 
-    pub fn shards_congestion_info(&self) -> HashMap<ShardId, CongestionInfo> {
-        self.chunks()
-            .iter()
-            .enumerate()
+    pub fn shards_congestion_info(&self) -> HashMap<ShardId, ExtendedCongestionInfo> {
+        let mut result = HashMap::new();
+
+        for chunk in self.chunks().iter() {
+            let shard_id = chunk.shard_id();
             // TODO(congestion_control): default is not always appropriate!
-            .map(|(i, chunk_header)| {
-                (i as ShardId, chunk_header.congestion_info().unwrap_or_default())
-            })
-            .collect()
+            let congestion_info = chunk.congestion_info().unwrap_or_default();
+            let height_included = chunk.height_included();
+            let height_current = self.header().height();
+            let missed_chunks_count = height_current.checked_sub(height_included);
+            let missed_chunks_count = missed_chunks_count
+                .expect("The chunk height included must be less or equal than block height!");
+
+            let extended_congestion_info =
+                ExtendedCongestionInfo::new(congestion_info, missed_chunks_count);
+            result.insert(shard_id, extended_congestion_info);
+        }
+        result
     }
 
     pub fn hash(&self) -> &CryptoHash {