feat(congestion_control) - handling missing chunks (#11274)

When there are multiple missing chunks in a row in a shard we want to
consider that shard as congested. That is in order to prevent outgoing
receipts to that shard accumulating and then blowing up the state
witness size.

I decided to not embed the information about missing chunks in the
congestion info but rather I added a new struct called
ExtendedCongestionInfo. The Block now constructs congestion info,
extends it with the information about missing chunks and provides this
new struct to the runtime.

Since from now the congestion level cannot be calculated without the
missing chunks information I added `missing_chunks_count` argument to
all methods that rely on the congestion level. That is to make sure the
users of those structs do not forget about the missing chunks - compiler
will warn them about it. In the runtime the ExtendedCongestionInfo
struct acts as a helper to make it as convenient as it used to be.

The congestion level itself is now a maximum of 4 values - the first
three as before and a new one for missing chunks. In this PR I made it
so that 10 missed chunks in a row would lead to full congestion - that
number is to be adjusted based on data. Other changes can also be
considered such as adding the missing chunks congestion to the max of
the others. I'm open for suggestions here.


It's in draft because I still need to add tests for this.

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 {