Skip to content

Commit

Permalink
feat(congestion_control) - handling missing chunks
Browse files Browse the repository at this point in the history
  • Loading branch information
wacban committed May 9, 2024
1 parent 152c2d9 commit 61e4aa7
Show file tree
Hide file tree
Showing 12 changed files with 191 additions and 90 deletions.
2 changes: 0 additions & 2 deletions chain/chain/src/chain_update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions chain/chain/src/runtime/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
5 changes: 3 additions & 2 deletions chain/chain/src/runtime/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions chain/chain/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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 {
Expand Down
25 changes: 16 additions & 9 deletions core/primitives/src/block.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -591,15 +591,22 @@ 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 - height_included;

let extended_congestion_info =
ExtendedCongestionInfo::new(congestion_info, missed_chunks_count);
result.insert(shard_id, extended_congestion_info);
}
result
}

pub fn hash(&self) -> &CryptoHash {
Expand Down

0 comments on commit 61e4aa7

Please sign in to comment.