Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[stateless_validation] Add testing for handling state witness and chunk endorsements #10504

Merged
merged 2 commits into from
Jan 26, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 70 additions & 61 deletions chain/client/src/test_utils/test_env.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,15 @@ use near_o11y::testonly::TracingCapture;
use near_parameters::RuntimeConfig;
use near_primitives::action::delegate::{DelegateAction, NonDelegateAction, SignedDelegateAction};
use near_primitives::block::Block;
use near_primitives::chunk_validation::ChunkStateWitness;
use near_primitives::epoch_manager::RngSeed;
use near_primitives::errors::InvalidTxError;
use near_primitives::hash::CryptoHash;
use near_primitives::network::PeerId;
use near_primitives::sharding::{ChunkHash, PartialEncodedChunk};
use near_primitives::test_utils::create_test_signer;
use near_primitives::transaction::{Action, FunctionCallAction, SignedTransaction};
use near_primitives::types::{AccountId, Balance, BlockHeight, EpochId, NumSeats, ShardId};
use near_primitives::types::{AccountId, Balance, BlockHeight, EpochId, NumSeats};
use near_primitives::utils::MaybeValidated;
use near_primitives::version::ProtocolVersion;
use near_primitives::views::{
Expand All @@ -38,7 +39,7 @@ use super::setup::{setup_client_with_runtime, ShardsManagerAdapterForTest};
use super::test_env_builder::TestEnvBuilder;
use super::TEST_SEED;

const CHUNK_ENDORSEMENTS_TIMEOUT: Duration = Duration::from_secs(60);
const CHUNK_ENDORSEMENTS_TIMEOUT: Duration = Duration::from_secs(10);

/// An environment for writing integration tests with multiple clients.
/// This environment can simulate near nodes without network and it can be configured to use different runtimes.
Expand All @@ -62,6 +63,9 @@ pub struct StateWitnessPropagationOutput {
/// Whether some propagated state witness includes two different post state
/// roots.
pub found_differing_post_state_root_due_to_state_transitions: bool,

/// All the account_ids that we've sent the chunk witness to.
pub chunk_hash_to_account_ids: HashMap<ChunkHash, HashSet<AccountId>>,
}

impl TestEnv {
Expand Down Expand Up @@ -270,87 +274,92 @@ impl TestEnv {
}
}

fn found_differing_post_state_root_due_to_state_transitions(
chunk_state_witness: &ChunkStateWitness,
) -> bool {
let chunk_state_witness_inner = &chunk_state_witness.inner;
let mut post_state_roots =
HashSet::from([chunk_state_witness_inner.main_state_transition.post_state_root]);
post_state_roots.extend(
chunk_state_witness_inner.implicit_transitions.iter().map(|t| t.post_state_root),
);
post_state_roots.len() >= 2
}

/// Triggers processing of all chunk state witnesses received by network.
pub fn propagate_chunk_state_witnesses(&mut self) -> StateWitnessPropagationOutput {
let mut found_differing_post_state_root_due_to_state_transitions = false;
for idx in 0..self.clients.len() {
let _span =
tracing::debug_span!(target: "test", "propagate_chunk_state_witnesses", client=idx)
.entered();

self.network_adapters[idx].handle_filtered(|msg| {
if let PeerManagerMessageRequest::NetworkRequests(
NetworkRequests::ChunkStateWitness(accounts, chunk_state_witness),
) = msg
{
let chunk_state_witness_inner = &chunk_state_witness.inner;
let mut post_state_roots = HashSet::from([chunk_state_witness_inner
.main_state_transition
.post_state_root]);
post_state_roots.extend(
chunk_state_witness_inner
.implicit_transitions
.iter()
.map(|t| t.post_state_root),
);
found_differing_post_state_root_due_to_state_transitions |=
post_state_roots.len() >= 2;
for account in accounts {
let sender_public_key = self.clients[idx]
.validator_signer
.as_ref()
.unwrap()
.public_key()
.clone();
self.account_indices
.lookup_mut(&mut self.clients, &account)
.process_chunk_state_witness(
chunk_state_witness.clone(),
PeerId::new(sender_public_key),
)
let mut output = StateWitnessPropagationOutput {
found_differing_post_state_root_due_to_state_transitions: false,
chunk_hash_to_account_ids: HashMap::new(),
};
let network_adapters = self.network_adapters.clone();
for network_adapter in network_adapters {
network_adapter.handle_filtered(|request| match request {
PeerManagerMessageRequest::NetworkRequests(NetworkRequests::ChunkStateWitness(
account_ids,
state_witness,
)) => {
// Process chunk state witness for each client.
for account_id in account_ids.iter() {
self.client(account_id)
.process_chunk_state_witness(state_witness.clone(), PeerId::random())
.unwrap();
}

// Update output.
output.found_differing_post_state_root_due_to_state_transitions |=
Self::found_differing_post_state_root_due_to_state_transitions(
&state_witness,
);

output.chunk_hash_to_account_ids.insert(
state_witness.inner.chunk_header.chunk_hash(),
account_ids.into_iter().collect(),
);
None
} else {
Some(msg)
}
_ => Some(request),
});
}
StateWitnessPropagationOutput { found_differing_post_state_root_due_to_state_transitions }
output
}

/// Waits for `CHUNK_ENDORSEMENTS_TIMEOUT` to receive at least one chunk
/// endorsement for the given chunk hashes.
/// Waits for `CHUNK_ENDORSEMENTS_TIMEOUT` to receive chunk endorsement for the given chunk hashes.
/// Panics if it doesn't happen.
/// `chunk_hashes` maps hashes to pair of height at which chunk is included
/// and shard id for better debugging.
pub fn wait_for_chunk_endorsements(
/// `chunk_hash_to_account_ids` maps hashes to the set of account ids that are expected to endorse the chunk.
/// Note that we need to wait here as the chunk state witness is processed asynchronously.
pub fn wait_to_propagate_chunk_endorsements(
&mut self,
mut chunk_hashes: HashMap<ChunkHash, (BlockHeight, ShardId)>,
chunk_hash_to_account_ids: &mut HashMap<ChunkHash, HashSet<AccountId>>,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you need this to be a &mut? Could you take the whole object by ownership?

) {
let _span = tracing::debug_span!(target: "test", "get_all_chunk_endorsements").entered();
let network_adapters = self.network_adapters.clone();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we actually need to clone this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes as otherwise we would hold a mut reference to self within the loop. rustc wasn't happy

let timer = Instant::now();
loop {
tracing::debug!(target: "test", "remaining endorsements: {:?}", chunk_hashes);
for idx in 0..self.clients.len() {
self.network_adapters[idx].handle_filtered(|msg| {
if let PeerManagerMessageRequest::NetworkRequests(
NetworkRequests::ChunkEndorsement(_, endorsement),
) = msg
{
chunk_hashes.remove(endorsement.chunk_hash());
for network_adapter in network_adapters.iter() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: you can replace network_adapters.iter() with &network_adapters

network_adapter.handle_filtered(|request| match request {
PeerManagerMessageRequest::NetworkRequests(
NetworkRequests::ChunkEndorsement(account_id, endorsement),
) => {
// Remove endorsement.account_id on receiving endorsement.
chunk_hash_to_account_ids
.get_mut(endorsement.chunk_hash())
.map(|entry| entry.remove(&endorsement.account_id));

self.client(&account_id).process_chunk_endorsement(endorsement).unwrap();

None
} else {
Some(msg)
}
_ => Some(request),
});
}

if chunk_hashes.is_empty() {
break;
// Check if we received all endorsements.
chunk_hash_to_account_ids.retain(|_, v| !v.is_empty());
if chunk_hash_to_account_ids.is_empty() {
return;
}
if timer.elapsed() > CHUNK_ENDORSEMENTS_TIMEOUT {
panic!("Missing chunk endorsements: {:?}", chunk_hashes);
panic!("Missing chunk endorsements: {:?}", chunk_hash_to_account_ids);
}
std::thread::sleep(Duration::from_micros(100));
}
Expand Down
44 changes: 12 additions & 32 deletions integration-tests/src/tests/client/features/chunk_validation.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::collections::HashSet;

use near_chain::{ChainGenesis, Provenance};
use near_chain_configs::{Genesis, GenesisConfig, GenesisRecords};
use near_client::test_utils::TestEnv;
Expand All @@ -18,9 +22,6 @@ use near_primitives_core::hash::CryptoHash;
use near_primitives_core::types::{AccountId, NumSeats};
use near_primitives_core::version::PROTOCOL_VERSION;
use nearcore::test_utils::TestEnvNightshadeSetupExt;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::collections::{HashMap, HashSet};

const ONE_NEAR: u128 = 1_000_000_000_000_000_000_000_000;

Expand Down Expand Up @@ -126,7 +127,6 @@ fn run_chunk_validation_test(seed: u64, prob_missing_chunk: f64) {
let mut tx_hashes = vec![];

let mut rng: StdRng = SeedableRng::seed_from_u64(seed);
let mut expected_chunks = HashMap::new();
let mut found_differing_post_state_root_due_to_state_transitions = false;
for round in 0..blocks_to_produce {
let heads = env
Expand Down Expand Up @@ -185,9 +185,15 @@ fn run_chunk_validation_test(seed: u64, prob_missing_chunk: f64) {
for j in 0..env.clients.len() {
env.process_shards_manager_responses_and_finish_processing_blocks(j);
}
let result = env.propagate_chunk_state_witnesses();

// First propagate chunk state witnesses, then chunk endorsements.
// Note that validation of chunk state witness is done asynchonously
// which is why it's important to pass the expected set of chunk endorsements to wait for.
let mut output = env.propagate_chunk_state_witnesses();
env.wait_to_propagate_chunk_endorsements(&mut output.chunk_hash_to_account_ids);

found_differing_post_state_root_due_to_state_transitions |=
result.found_differing_post_state_root_due_to_state_transitions;
output.found_differing_post_state_root_due_to_state_transitions;
}

// Check that at least one tx was fully executed, ensuring that executing
Expand All @@ -211,32 +217,6 @@ fn run_chunk_validation_test(seed: u64, prob_missing_chunk: f64) {
if prob_missing_chunk >= 0.8 {
assert!(found_differing_post_state_root_due_to_state_transitions);
}

// Collect chunk hashes which have to be endorsed and check that it indeed
// happens.
let mut block = env.clients[0]
.chain
.get_block(&env.clients[0].chain.head().unwrap().last_block_hash)
.unwrap();
loop {
let prev_hash = *block.header().prev_hash();
if prev_hash == CryptoHash::default() {
break;
}
let prev_block = env.clients[0].chain.get_block(&prev_hash).unwrap();
for (chunk, prev_chunk) in block.chunks().iter().zip(prev_block.chunks().iter()) {
if chunk.is_new_chunk(block.header().height()) {
// First chunk after genesis doesn't have to be endorsed.
if prev_chunk.prev_block_hash() != &CryptoHash::default() {
expected_chunks
.insert(chunk.chunk_hash(), (block.header().height(), chunk.shard_id()));
}
}
}
block = prev_block;
}

env.wait_for_chunk_endorsements(expected_chunks);
}

#[test]
Expand Down
Loading