From 8319833a03231105f14f25b713fed52bb785cb01 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 2 Aug 2023 12:53:26 +0000
Subject: [PATCH 01/15] devlog

---
 core/o11y/src/testonly.rs        | 2 +-
 core/primitives-core/src/hash.rs | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/core/o11y/src/testonly.rs b/core/o11y/src/testonly.rs
index 2dddbc9f86c..2fff88ca204 100644
--- a/core/o11y/src/testonly.rs
+++ b/core/o11y/src/testonly.rs
@@ -27,7 +27,7 @@ fn setup_subscriber_from_filter(mut env_filter: EnvFilter) {
 
     let _ = fmt::Subscriber::builder()
         .with_ansi(use_color_auto())
-        .with_span_events(fmt::format::FmtSpan::CLOSE)
+        .with_span_events(fmt::format::FmtSpan::NONE)
         .with_env_filter(env_filter)
         .with_writer(fmt::TestWriter::new())
         .with_timer(TestUptime::default())
diff --git a/core/primitives-core/src/hash.rs b/core/primitives-core/src/hash.rs
index 7b703e9cd2d..8e77e8724ee 100644
--- a/core/primitives-core/src/hash.rs
+++ b/core/primitives-core/src/hash.rs
@@ -207,7 +207,8 @@ impl fmt::Debug for CryptoHash {
 
 impl fmt::Display for CryptoHash {
     fn fmt(&self, fmtr: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.to_base58_impl(|encoded| fmtr.write_str(encoded))
+        // TODO remove me debugging only
+        self.to_base58_impl(|encoded| fmtr.write_str(&encoded[..4]))
     }
 }
 

From 7998eda90fb8b6e589bcdef06275a542dcb3dc5a Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Thu, 3 Aug 2023 08:15:05 +0000
Subject: [PATCH 02/15] test script

---
 .gitignore |  2 ++
 r.sh       | 30 ++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100755 r.sh

diff --git a/.gitignore b/.gitignore
index f4106d16321..edd9464ff9a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -62,3 +62,5 @@ rusty-tags.vi
 costs-*.txt
 names-to-stats.txt
 data_dump_*.bin
+
+stdout-1
diff --git a/r.sh b/r.sh
new file mode 100755
index 00000000000..b84c1f5f691
--- /dev/null
+++ b/r.sh
@@ -0,0 +1,30 @@
+rm stdout-*
+
+ for i in {1..1}
+ do
+     echo "test $i"
+     OUT=stdout-$i
+
+     # RUST_LOG=info \
+     #RUST_BACKTRACE=all \
+     #RUST_LOG=info,catchup=trace,store=trace,client=debug,store=debug,test=debug,resharding=trace \
+     RUST_LOG=debug,resharding=trace \
+     cargo nextest run -p integration-tests \
+         --no-capture \
+         --features nightly \
+		test_shard_layout_upgrade_cross_contract_calls \
+         > $OUT
+         # | tee $OUT
+         # | egrep -v -i "FlatStorage is not ready|Add delta for flat storage creation|epoch_manager: all proposals" \
+         # test_shard_layout_upgrade_simple_v1 \
+
+     sed -E -i 's/ed25519:(.{4})(.{40})/ed25519:\1/g' $OUT
+
+     sed -E -i 's/([0-9]*)([0-9]{30})/\1e30/g' $OUT
+     sed -E -i 's/([0-9]*)([0-9]{25})/\1e25/g' $OUT
+     sed -E -i 's/([0-9]*)([0-9]{20})/\1e20/g' $OUT
+     sed -E -i 's/AccountId/AId/g' $OUT
+
+     cat $OUT | egrep -a -i error
+
+ done

From 0f6edbbd5b6d71dbcf381fbdedc188681f83a58b Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Thu, 3 Aug 2023 10:03:01 +0000
Subject: [PATCH 03/15] pass resharding type

---
 .../src/tests/client/sharding_upgrade.rs      | 61 +++++++++++--------
 1 file changed, 37 insertions(+), 24 deletions(-)

diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index c9fa17285a8..ea618625736 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -137,23 +137,6 @@ impl TestShardUpgradeEnv {
         }
     }
 
-    fn new(
-        epoch_length: u64,
-        num_validators: usize,
-        num_clients: usize,
-        num_init_accounts: usize,
-        gas_limit: Option<u64>,
-    ) -> Self {
-        Self::new_with_protocol_version(
-            epoch_length,
-            num_validators,
-            num_clients,
-            num_init_accounts,
-            gas_limit,
-            SIMPLE_NIGHTSHADE_PROTOCOL_VERSION - 1,
-        )
-    }
-
     /// `init_txs` are added before any block is produced
     fn set_init_tx(&mut self, init_txs: Vec<SignedTransaction>) {
         self.init_txs = init_txs;
@@ -868,8 +851,17 @@ fn gen_cross_contract_transaction(
 /// Return test_env and a map from tx hash to the new account that will be added by this transaction
 fn setup_test_env_with_cross_contract_txs(
     epoch_length: u64,
+    genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
-    let mut test_env = TestShardUpgradeEnv::new(epoch_length, 4, 4, 100, Some(100_000_000_000_000));
+    // let mut test_env = TestShardUpgradeEnv::new(epoch_length, 4, 4, 100, Some(100_000_000_000_000));
+    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
+        epoch_length,
+        4,
+        4,
+        100,
+        Some(100_000_000_000_000),
+        genesis_protocol_version,
+    );
     let mut rng = thread_rng();
 
     let genesis_hash = *test_env.env.clients[0].chain.genesis_block().hash();
@@ -963,37 +955,58 @@ fn setup_test_env_with_cross_contract_txs(
 
 // Test cross contract calls
 // This test case tests postponed receipts and delayed receipts
-#[test]
-fn test_shard_layout_upgrade_cross_contract_calls() {
+fn test_shard_layout_upgrade_cross_contract_calls_impl(resharding_type: ReshardingType) {
     init_test_logger();
 
     // setup
     let epoch_length = 5;
+    let genesis_protocol_version = get_genesis_protocol_version(&resharding_type);
+    let target_protocol_version = get_target_protocol_version(&resharding_type);
 
-    let (mut test_env, new_accounts) = setup_test_env_with_cross_contract_txs(epoch_length);
+    let (mut test_env, new_accounts) =
+        setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 
     for _ in 1..5 * epoch_length {
-        test_env.step(0.);
+        test_env.step_impl(0., target_protocol_version, &resharding_type);
         test_env.check_receipt_id_to_shard_id();
     }
 
     let successful_txs = test_env.check_tx_outcomes(false, vec![2 * epoch_length + 1]);
-    let new_accounts: Vec<_> =
+    let new_accounts =
         successful_txs.iter().flat_map(|tx_hash| new_accounts.get(tx_hash)).collect();
+
     test_env.check_accounts(new_accounts);
 
     test_env.check_split_states_artifacts();
 }
 
+// Test cross contract calls
+// This test case tests postponed receipts and delayed receipts
+#[test]
+fn test_shard_layout_upgrade_cross_contract_calls_v1() {
+    test_shard_layout_upgrade_cross_contract_calls_impl(ReshardingType::V1);
+}
+
+// Test cross contract calls
+// This test case tests postponed receipts and delayed receipts
+#[test]
+fn test_shard_layout_upgrade_cross_contract_calls_v2() {
+    test_shard_layout_upgrade_cross_contract_calls_impl(ReshardingType::V2);
+}
+
 // Test cross contract calls
 // This test case tests when there are missing chunks in the produced blocks
 // This is to test that all the chunk management logic in sharding split is correct
 fn test_shard_layout_upgrade_missing_chunks(p_missing: f64) {
     init_test_logger();
 
+    let resharding_type = ReshardingType::V1;
+    let genesis_protocol_versino = get_genesis_protocol_version(&resharding_type);
+
     // setup
     let epoch_length = 10;
-    let (mut test_env, new_accounts) = setup_test_env_with_cross_contract_txs(epoch_length);
+    let (mut test_env, new_accounts) =
+        setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_versino);
 
     // randomly dropping chunks at the first few epochs when sharding splits happens
     // make sure initial txs (deploy smart contracts) are processed succesfully

From 368ba9df42650abcdf64b95994898a1c5a338700 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Fri, 4 Aug 2023 11:05:06 +0000
Subject: [PATCH 04/15] debug logs

---
 chain/chain/src/chain.rs                      |   9 +
 .../src/tests/client/sharding_upgrade.rs      | 170 +++++++++---------
 r.sh                                          |   4 +-
 3 files changed, 98 insertions(+), 85 deletions(-)

diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index f8cb3acbf54..8cda355a90a 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -1642,6 +1642,7 @@ impl Chain {
                 let chunk_hash = chunk_header.chunk_hash();
 
                 if let Err(_) = self.store.get_partial_chunk(&chunk_header.chunk_hash()) {
+                    tracing::debug!(target: "waclaw", "get partial chunk failed");
                     missing.push(chunk_header.clone());
                 } else if self.shard_tracker.care_about_shard(
                     me.as_ref(),
@@ -1655,6 +1656,7 @@ impl Chain {
                     true,
                 ) {
                     if let Err(_) = self.store.get_chunk(&chunk_hash) {
+                        tracing::debug!(target: "waclaw", "get chunk failed");
                         missing.push(chunk_header.clone());
                     }
                 }
@@ -2105,6 +2107,7 @@ impl Chain {
             block_received_time,
             state_patch,
         );
+        tracing::debug!(target: "waclaw", preprocess_res=preprocess_res.is_ok(), preprocess_res_err=?preprocess_res.as_ref().err(), "preprocess_res");
         let preprocess_res = match preprocess_res {
             Ok(preprocess_res) => {
                 preprocess_timer.observe_duration();
@@ -2482,6 +2485,7 @@ impl Chain {
         debug!(target: "chain", block_hash = ?header.hash(), me=?me, is_caught_up=is_caught_up, "Process block");
 
         // Check the header is valid before we proceed with the full block.
+        tracing::debug!(target: "waclaw", "validate header");
         self.validate_header(header, provenance, challenges)?;
 
         self.epoch_manager.verify_block_vrf(
@@ -2496,6 +2500,7 @@ impl Chain {
             return Err(Error::InvalidRandomnessBeaconOutput);
         }
 
+        tracing::debug!(target: "waclaw", "validate with");
         let res = block.validate_with(|block| {
             Chain::validate_block_impl(self.epoch_manager.as_ref(), &self.genesis, block)
                 .map(|_| true)
@@ -2531,9 +2536,13 @@ impl Chain {
 
         let prev_block = self.get_block(&prev_hash)?;
 
+        tracing::debug!(target: "waclaw", "validate chunk headers");
         self.validate_chunk_headers(&block, &prev_block)?;
 
+        tracing::debug!(target: "waclaw", "ping missing chunks");
         self.ping_missing_chunks(me, prev_hash, block)?;
+
+        tracing::debug!(target: "waclaw", "collect incoming receipts from block");
         let incoming_receipts = self.collect_incoming_receipts_from_block(me, block)?;
 
         // Check if block can be finalized and drop it otherwise.
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index ea618625736..3eed913542d 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -46,7 +46,7 @@ const SIMPLE_NIGHTSHADE_V2_PROTOCOL_VERSION: ProtocolVersion =
 #[cfg(not(feature = "protocol_feature_simple_nightshade_v2"))]
 const SIMPLE_NIGHTSHADE_V2_PROTOCOL_VERSION: ProtocolVersion = PROTOCOL_VERSION + 1;
 
-const P_CATCHUP: f64 = 0.2;
+// const P_CATCHUP: f64 = 0.2;
 
 enum ReshardingType {
     // In the V0->V1 resharding outgoing receipts are reassigned to receiver.
@@ -94,7 +94,6 @@ struct TestShardUpgradeEnv {
     init_txs: Vec<SignedTransaction>,
     txs_by_height: BTreeMap<u64, Vec<SignedTransaction>>,
     epoch_length: u64,
-    num_validators: usize,
     num_clients: usize,
 }
 
@@ -126,11 +125,11 @@ impl TestShardUpgradeEnv {
             .real_epoch_managers(&genesis.config)
             .nightshade_runtimes(&genesis)
             .build();
+        assert_eq!(env.validators.len(), num_validators);
         Self {
             env,
             initial_accounts,
             epoch_length,
-            num_validators,
             num_clients,
             init_txs: vec![],
             txs_by_height: BTreeMap::new(),
@@ -181,12 +180,7 @@ impl TestShardUpgradeEnv {
         // add transactions for the next block
         if height == 1 {
             for tx in self.init_txs.iter() {
-                for j in 0..self.num_validators {
-                    assert_eq!(
-                        env.clients[j].process_tx(tx.clone(), false, false),
-                        ProcessTxResponse::ValidTx
-                    );
-                }
+                Self::process_tx(env, tx);
             }
         }
 
@@ -196,47 +190,39 @@ impl TestShardUpgradeEnv {
         // it when we are producing the block at `height`
         if let Some(txs) = self.txs_by_height.get(&(height + 1)) {
             for tx in txs {
-                let mut response_valid_count = 0;
-                let mut response_routed_count = 0;
-                for j in 0..self.num_validators {
-                    let response = env.clients[j].process_tx(tx.clone(), false, false);
-                    tracing::debug!(target: "test", client=j, tx=?tx.get_hash(), ?response, "process tx");
-                    match response {
-                        ProcessTxResponse::ValidTx => response_valid_count += 1,
-                        ProcessTxResponse::RequestRouted => response_routed_count += 1,
-                        response => {
-                            panic!("invalid tx response {response:?}");
-                        }
-                    }
-                }
-                assert_ne!(response_valid_count, 0);
-                assert_eq!(response_valid_count + response_routed_count, self.num_validators);
+                Self::process_tx(env, tx);
             }
         }
 
         // produce block
-        let block_producer = {
-            let epoch_id = env.clients[0]
-                .epoch_manager
-                .get_epoch_id_from_prev_block(&head.last_block_hash)
-                .unwrap();
-            env.clients[0].epoch_manager.get_block_producer(&epoch_id, height).unwrap()
+        let block = {
+            let client = &env.clients[0];
+            let epoch_id =
+                client.epoch_manager.get_epoch_id_from_prev_block(&head.last_block_hash).unwrap();
+            let block_producer =
+                client.epoch_manager.get_block_producer(&epoch_id, height).unwrap();
+            let _span = tracing::debug_span!(target: "test", "", client=?block_producer);
+            let _span = _span.entered();
+            let block_producer_client = env.client(&block_producer);
+            let mut block = block_producer_client.produce_block(height).unwrap().unwrap();
+            set_block_protocol_version(&mut block, block_producer.clone(), protocol_version);
+            block
         };
-        let block_producer_client = env.client(&block_producer);
-        let mut block = block_producer_client.produce_block(height).unwrap().unwrap();
-        set_block_protocol_version(&mut block, block_producer.clone(), protocol_version);
         // Make sure that catchup is done before the end of each epoch, but when it is done is
         // by chance. This simulates when catchup takes a long time to be done
         // Note: if the catchup happens only at the last block of an epoch then
         // client will fail to produce the chunks in the first block of the next epoch.
-        let should_catchup = rng.gen_bool(P_CATCHUP) || height % self.epoch_length == 0;
+        // let should_catchup = rng.gen_bool(P_CATCHUP) || height % self.epoch_length == 0;
+        let should_catchup = (height + 1) % self.epoch_length == 0;
         // process block, this also triggers chunk producers for the next block to produce chunks
         for j in 0..self.num_clients {
+            let client = &mut env.clients[j];
+            let _span = tracing::debug_span!(target: "test", "process block", client=j).entered();
             let produce_chunks = !rng.gen_bool(p_drop_chunk);
             // Here we don't just call self.clients[i].process_block_sync_with_produce_chunk_options
             // because we want to call run_catchup before finish processing this block. This simulates
             // that catchup and block processing run in parallel.
-            env.clients[j]
+            client
                 .start_process_block(
                     MaybeValidated::from(block.clone()),
                     Provenance::NONE,
@@ -244,11 +230,10 @@ impl TestShardUpgradeEnv {
                 )
                 .unwrap();
             if should_catchup {
-                run_catchup(&mut env.clients[j], &[]).unwrap();
+                run_catchup(client, &[]).unwrap();
             }
-            while wait_for_all_blocks_in_processing(&mut env.clients[j].chain) {
-                let (_, errors) =
-                    env.clients[j].postprocess_ready_blocks(Arc::new(|_| {}), produce_chunks);
+            while wait_for_all_blocks_in_processing(&mut client.chain) {
+                let (_, errors) = client.postprocess_ready_blocks(Arc::new(|_| {}), produce_chunks);
                 assert!(errors.is_empty(), "unexpected errors: {:?}", errors);
             }
             if should_catchup {
@@ -276,6 +261,24 @@ impl TestShardUpgradeEnv {
         }
     }
 
+    fn process_tx(env: &mut TestEnv, tx: &SignedTransaction) {
+        let mut response_valid_count = 0;
+        let mut response_routed_count = 0;
+        for j in 0..env.validators.len() {
+            let response = env.clients[j].process_tx(tx.clone(), false, false);
+            tracing::debug!(target: "test", client=j, tx=?tx.get_hash(), ?response, "process tx");
+            match response {
+                ProcessTxResponse::ValidTx => response_valid_count += 1,
+                ProcessTxResponse::RequestRouted => response_routed_count += 1,
+                response => {
+                    panic!("invalid tx response {response:?} {tx:?}");
+                }
+            }
+        }
+        assert_ne!(response_valid_count, 0);
+        assert_eq!(response_valid_count + response_routed_count, env.validators.len());
+    }
+
     /// check that all accounts in `accounts` exist in the current state
     fn check_accounts(&mut self, accounts: Vec<&AccountId>) {
         tracing::debug!(target: "test", "checking accounts");
@@ -640,7 +643,7 @@ fn setup_genesis(
     // shards so that we can avoid that problem for now. It would be nice to
     // set it up properly and test both state sync and resharding - once the
     // integration is fully supported.
-    genesis.config.minimum_validators_per_shard = 2;
+    genesis.config.minimum_validators_per_shard = num_validators;
 
     genesis
 }
@@ -874,33 +877,33 @@ fn setup_test_env_with_cross_contract_txs(
         test_env.initial_accounts[indices[0]].clone(),
         test_env.initial_accounts[indices[1]].clone(),
     ];
-    test_env.set_init_tx(
-        contract_accounts
-            .iter()
-            .map(|account_id| {
-                let signer = InMemorySigner::from_seed(
-                    account_id.clone(),
-                    KeyType::ED25519,
-                    &account_id.to_string(),
-                );
-                SignedTransaction::from_actions(
-                    1,
-                    account_id.clone(),
-                    account_id.clone(),
-                    &signer,
-                    vec![Action::DeployContract(DeployContractAction {
-                        code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
-                    })],
-                    genesis_hash,
-                )
-            })
-            .collect(),
-    );
+    // test_env.set_init_tx(
+    //     contract_accounts
+    //         .iter()
+    //         .map(|account_id| {
+    //             let signer = InMemorySigner::from_seed(
+    //                 account_id.clone(),
+    //                 KeyType::ED25519,
+    //                 &account_id.to_string(),
+    //             );
+    //             SignedTransaction::from_actions(
+    //                 1,
+    //                 account_id.clone(),
+    //                 account_id.clone(),
+    //                 &signer,
+    //                 vec![Action::DeployContract(DeployContractAction {
+    //                     code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
+    //                 })],
+    //                 genesis_hash,
+    //             )
+    //         })
+    //         .collect(),
+    // );
 
     let mut nonce = 100;
     let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
     let mut new_accounts = HashMap::new();
-    let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
+    let _generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
         &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
             let mut rng = thread_rng();
             let size = rng.gen_range(min_size..max_size + 1);
@@ -929,26 +932,26 @@ fn setup_test_env_with_cross_contract_txs(
             .collect()
         };
 
-    // add a bunch of transactions before the two epoch boundaries
-    for height in vec![
-        epoch_length - 2,
-        epoch_length - 1,
-        epoch_length,
-        2 * epoch_length - 2,
-        2 * epoch_length - 1,
-        2 * epoch_length,
-    ] {
-        test_env.set_tx_at_height(height, generate_txs(5, 8));
-    }
-
-    // adds some transactions after sharding change finishes
-    // but do not add too many because I want all transactions to
-    // finish processing before epoch 5
-    for height in 2 * epoch_length + 1..3 * epoch_length {
-        if rng.gen_bool(0.3) {
-            test_env.set_tx_at_height(height, generate_txs(5, 8));
-        }
-    }
+    // // add a bunch of transactions before the two epoch boundaries
+    // for height in vec![
+    //     epoch_length - 2,
+    //     epoch_length - 1,
+    //     epoch_length,
+    //     2 * epoch_length - 2,
+    //     2 * epoch_length - 1,
+    //     2 * epoch_length,
+    // ] {
+    //     test_env.set_tx_at_height(height, generate_txs(5, 8));
+    // }
+
+    // // adds some transactions after sharding change finishes
+    // // but do not add too many because I want all transactions to
+    // // finish processing before epoch 5
+    // for height in 2 * epoch_length + 1..3 * epoch_length {
+    //     if rng.gen_bool(0.3) {
+    //         test_env.set_tx_at_height(height, generate_txs(5, 8));
+    //     }
+    // }
 
     (test_env, new_accounts)
 }
@@ -967,6 +970,7 @@ fn test_shard_layout_upgrade_cross_contract_calls_impl(resharding_type: Reshardi
         setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 
     for _ in 1..5 * epoch_length {
+        std::thread::sleep(std::time::Duration::from_millis(1000));
         test_env.step_impl(0., target_protocol_version, &resharding_type);
         test_env.check_receipt_id_to_shard_id();
     }
diff --git a/r.sh b/r.sh
index b84c1f5f691..8dd7c4d65f0 100755
--- a/r.sh
+++ b/r.sh
@@ -6,13 +6,13 @@ rm stdout-*
      OUT=stdout-$i
 
      # RUST_LOG=info \
-     #RUST_BACKTRACE=all \
      #RUST_LOG=info,catchup=trace,store=trace,client=debug,store=debug,test=debug,resharding=trace \
+     RUST_BACKTRACE=all \
      RUST_LOG=debug,resharding=trace \
      cargo nextest run -p integration-tests \
          --no-capture \
          --features nightly \
-		test_shard_layout_upgrade_cross_contract_calls \
+		test_shard_layout_upgrade_cross_contract_calls_v2 \
          > $OUT
          # | tee $OUT
          # | egrep -v -i "FlatStorage is not ready|Add delta for flat storage creation|epoch_manager: all proposals" \

From 51f81a9bb7d1b46e494606c321316941b62079c5 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Mon, 7 Aug 2023 15:14:30 +0000
Subject: [PATCH 05/15] it works

---
 Cargo.lock                                    |  1 +
 chain/chain/src/chain.rs                      | 11 ++-
 chain/chunks/Cargo.toml                       |  1 +
 chain/chunks/src/lib.rs                       | 67 +++++++++++--------
 chain/chunks/src/logic.rs                     |  2 +-
 chain/chunks/src/test_loop.rs                 |  2 +-
 chain/client/src/client.rs                    |  3 +
 chain/client/src/test_utils.rs                | 61 ++++++++++++-----
 chain/network/src/test_utils.rs               |  2 +
 .../src/tests/client/sharding_upgrade.rs      | 43 +++++++-----
 r.sh                                          | 15 +++--
 11 files changed, 142 insertions(+), 66 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index dede42fa941..8d916f43382 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3542,6 +3542,7 @@ dependencies = [
  "derive-enum-from-into",
  "derive_more",
  "futures",
+ "itertools",
  "lru",
  "near-async",
  "near-chain",
diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index 8cda355a90a..4e6c7e5d802 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -1120,6 +1120,7 @@ impl Chain {
             &prev_hash,
         )?;
         let prev_block = self.get_block(&prev_hash)?;
+        tracing::trace!(target: "waclaw", ?shards_to_state_sync, "get_state_sync_info");
 
         if prev_block.chunks().len() != block.chunks().len() && !shards_to_state_sync.is_empty() {
             // Currently, the state sync algorithm assumes that the number of chunks do not change
@@ -2720,8 +2721,14 @@ impl Chain {
         parent_hash: &CryptoHash,
         shard_id: ShardId,
     ) -> bool {
-        let will_shard_layout_change =
-            epoch_manager.will_shard_layout_change(parent_hash).unwrap_or(false);
+        let result = epoch_manager.will_shard_layout_change(parent_hash);
+        let will_shard_layout_change = match result {
+            Ok(will_shard_layout_change) => will_shard_layout_change,
+            Err(err) => {
+                tracing::error!(target: "waclaw", ?err, "failed to check if shard layout will change");
+                return false;
+            }
+        };
         // if shard layout will change the next epoch, we should catch up the shard regardless
         // whether we already have the shard's state this epoch, because we need to generate
         // new states for shards split from the current shard for the next epoch
diff --git a/chain/chunks/Cargo.toml b/chain/chunks/Cargo.toml
index ade8f3c146d..e1322ed71a0 100644
--- a/chain/chunks/Cargo.toml
+++ b/chain/chunks/Cargo.toml
@@ -21,6 +21,7 @@ rand.workspace = true
 reed-solomon-erasure.workspace = true
 time.workspace = true
 tracing.workspace = true
+itertools.workspace = true
 
 near-async.workspace = true
 near-chain-configs.workspace = true
diff --git a/chain/chunks/src/lib.rs b/chain/chunks/src/lib.rs
index b734563e322..b36dd1e7c3b 100644
--- a/chain/chunks/src/lib.rs
+++ b/chain/chunks/src/lib.rs
@@ -137,6 +137,7 @@ pub mod metrics;
 pub mod shards_manager_actor;
 pub mod test_loop;
 pub mod test_utils;
+use itertools::Itertools;
 
 pub const CHUNK_REQUEST_RETRY: time::Duration = time::Duration::milliseconds(100);
 pub const CHUNK_REQUEST_SWITCH_TO_OTHERS: time::Duration = time::Duration::milliseconds(400);
@@ -606,17 +607,21 @@ impl ShardsManager {
             },
         );
 
-        if !mark_only {
-            let fetch_from_archival = chunk_needs_to_be_fetched_from_archival(
+        if mark_only {
+            debug!(target: "chunks", height, shard_id, ?chunk_hash, "Marked the chunk as being requested but did not send the request yet.");
+            return;
+        }
+
+        let fetch_from_archival = chunk_needs_to_be_fetched_from_archival(
                 &ancestor_hash, &self.chain_header_head.last_block_hash,
             self.epoch_manager.as_ref()).unwrap_or_else(|err| {
                 error!(target: "chunks", "Error during requesting partial encoded chunk. Cannot determine whether to request from an archival node, defaulting to not: {}", err);
                 false
             });
-            let old_block = self.chain_header_head.last_block_hash != prev_block_hash
-                && self.chain_header_head.prev_block_hash != prev_block_hash;
+        let old_block = self.chain_header_head.last_block_hash != prev_block_hash
+            && self.chain_header_head.prev_block_hash != prev_block_hash;
 
-            let should_wait_for_chunk_forwarding =
+        let should_wait_for_chunk_forwarding =
                 self.should_wait_for_chunk_forwarding(&ancestor_hash, chunk_header.shard_id(), chunk_header.height_created()+1).unwrap_or_else(|_| {
                     // ancestor_hash must be accepted because we don't request missing chunks through this
                     // this function for orphans
@@ -625,30 +630,27 @@ impl ShardsManager {
                     false
                 });
 
-            // If chunks forwarding is enabled,
-            // we purposely do not send chunk request messages right away for new blocks. Such requests
-            // will eventually be sent because of the `resend_chunk_requests` loop. However,
-            // we want to give some time for any `PartialEncodedChunkForward` messages to arrive
-            // before we send requests.
-            if !should_wait_for_chunk_forwarding || fetch_from_archival || old_block {
-                debug!(target: "chunks", height, shard_id, ?chunk_hash, "Requesting.");
-                let request_result = self.request_partial_encoded_chunk(
-                    height,
-                    &ancestor_hash,
-                    shard_id,
-                    &chunk_hash,
-                    false,
-                    old_block,
-                    fetch_from_archival,
-                );
-                if let Err(err) = request_result {
-                    error!(target: "chunks", "Error during requesting partial encoded chunk: {}", err);
-                }
-            } else {
-                debug!(target: "chunks",should_wait_for_chunk_forwarding, fetch_from_archival, old_block,  "Delaying the chunk request.");
+        // If chunks forwarding is enabled,
+        // we purposely do not send chunk request messages right away for new blocks. Such requests
+        // will eventually be sent because of the `resend_chunk_requests` loop. However,
+        // we want to give some time for any `PartialEncodedChunkForward` messages to arrive
+        // before we send requests.
+        if !should_wait_for_chunk_forwarding || fetch_from_archival || old_block {
+            debug!(target: "chunks", height, shard_id, ?chunk_hash, "Requesting.");
+            let request_result = self.request_partial_encoded_chunk(
+                height,
+                &ancestor_hash,
+                shard_id,
+                &chunk_hash,
+                false,
+                old_block,
+                fetch_from_archival,
+            );
+            if let Err(err) = request_result {
+                error!(target: "chunks", "Error during requesting partial encoded chunk: {}", err);
             }
         } else {
-            debug!(target: "chunks", height, shard_id, ?chunk_hash, "Marked the chunk as being requested but did not send the request yet.");
+            debug!(target: "chunks",should_wait_for_chunk_forwarding, fetch_from_archival, old_block,  "Delaying the chunk request.");
         }
     }
 
@@ -702,6 +704,7 @@ impl ShardsManager {
         // Process chunk one part requests.
         let requests = self.requested_partial_encoded_chunks.fetch(self.clock.now().into());
         for (chunk_hash, chunk_request) in requests {
+            tracing::trace!(target: "client", ?chunk_request, "resending chunk request");
             let fetch_from_archival =
                 chunk_needs_to_be_fetched_from_archival(&chunk_request.ancestor_hash, &self.chain_header_head.last_block_hash,
                 self.epoch_manager.as_ref()).unwrap_or_else(|err| {
@@ -1169,6 +1172,7 @@ impl ShardsManager {
         &mut self,
         forward: PartialEncodedChunkForwardMsg,
     ) -> Result<(), Error> {
+        tracing::debug!(target: "chunks", shard_id=forward.shard_id, parts_len=forward.parts.len(), "process_partial_encoded_chunk_forward");
         let maybe_header = self
             .validate_partial_encoded_chunk_forward(&forward)
             .and_then(|_| self.get_partial_encoded_chunk_header(&forward.chunk_hash));
@@ -1302,11 +1306,13 @@ impl ShardsManager {
         &mut self,
         header: &ShardChunkHeader,
     ) -> bool {
+        tracing::debug!(target: "waclaw", "insert_header_if_not_exists_and_process_cached_chunk_forwards");
         let header_known_before = self.encoded_chunks.get(&header.chunk_hash()).is_some();
         if self.encoded_chunks.get_or_insert_from_header(header).complete {
             return false;
         }
         if let Some(parts) = self.chunk_forwards_cache.pop(&header.chunk_hash()) {
+            tracing::debug!(target: "waclaw", parts_len=parts.len(), "insert_header_if_not_exists_and_process_cached_chunk_forwards");
             // Note that we don't need any further validation for the forwarded part.
             // The forwarded part was earlier validated via validate_partial_encoded_chunk_forward,
             // which checks the part against the merkle root in the forward message, and the merkle
@@ -1357,6 +1363,7 @@ impl ShardsManager {
         let chunk_hash = header.chunk_hash();
         debug!(
             target: "chunks",
+            me = ?self.me,
             ?chunk_hash,
             height = header.height_created(),
             shard_id = header.shard_id(),
@@ -1434,7 +1441,7 @@ impl ShardsManager {
             }
         }
 
-        // 2. Consider it valid; mergeparts and receipts included in the partial encoded chunk
+        // 2. Consider it valid; merge parts and receipts included in the partial encoded chunk
         // into chunk cache
         let new_part_ords =
             self.encoded_chunks.merge_in_partial_encoded_chunk(partial_encoded_chunk);
@@ -1684,6 +1691,7 @@ impl ShardsManager {
         epoch_id: &EpochId,
         lastest_block_hash: &CryptoHash,
     ) -> Result<(), Error> {
+        tracing::debug!(target: "client", me=?self.me, parts=?partial_encoded_chunk.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "send_partial_encoded_chunk_to_chunk_trackers");
         let me = match self.me.as_ref() {
             Some(me) => me,
             None => return Ok(()),
@@ -1700,6 +1708,7 @@ impl ShardsManager {
             })
             .cloned()
             .collect();
+        tracing::debug!(target: "client", ?me, owned_parts_len=owned_parts.len(), "send_partial_encoded_chunk_to_chunk_trackers");
 
         if owned_parts.is_empty() {
             return Ok(());
@@ -1728,6 +1737,7 @@ impl ShardsManager {
             }
             next_chunk_producers.remove(&bp_account_id);
 
+            tracing::debug!(target: "waclaw", ?bp_account_id, shard_id=partial_encoded_chunk.header.shard_id(), parts=?forward.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "sending forward");
             // Technically, here we should check if the block producer actually cares about the shard.
             // We don't because with the current implementation, we force all validators to track all
             // shards by making their config tracking all shards.
@@ -1900,6 +1910,7 @@ impl ShardsManager {
                 );
 
             if Some(&to_whom) != self.me.as_ref() {
+                tracing::trace!(target: "waclaw", parts_len=partial_encoded_chunk.parts.len(), "sending partial encoded chunk");
                 self.peer_manager_adapter.send(PeerManagerMessageRequest::NetworkRequests(
                     NetworkRequests::PartialEncodedChunkMessage {
                         account_id: to_whom.clone(),
diff --git a/chain/chunks/src/logic.rs b/chain/chunks/src/logic.rs
index 8586d3cc810..394fe36d79a 100644
--- a/chain/chunks/src/logic.rs
+++ b/chain/chunks/src/logic.rs
@@ -154,7 +154,7 @@ pub fn decode_encoded_chunk(
             Ok(shard_chunk)
         })
     {
-        debug!(target: "chunks", "Reconstructed and decoded chunk {}, encoded length was {}, num txs: {}, I'm {:?}", chunk_hash.0, encoded_chunk.encoded_length(), shard_chunk.transactions().len(), me);
+        debug!(target: "chunks", "Reconstructed and decoded chunk {}, shard id {}, encoded length was {}, num txs: {}, I'm {:?}", chunk_hash.0, encoded_chunk.shard_id(), encoded_chunk.encoded_length(), shard_chunk.transactions().len(), me);
 
         let partial_chunk = create_partial_chunk(
             encoded_chunk,
diff --git a/chain/chunks/src/test_loop.rs b/chain/chunks/src/test_loop.rs
index 32ab8f7ed10..12b38979faa 100644
--- a/chain/chunks/src/test_loop.rs
+++ b/chain/chunks/src/test_loop.rs
@@ -347,7 +347,7 @@ impl TestChunkEncoder {
     }
 
     pub fn part_ords(&self) -> Vec<u64> {
-        self.full_partial_chunk.parts().iter().map(|part| part.part_ord).collect()
+        self.full_partial_chunk.parts().iter().take(3).map(|part| part.part_ord).collect()
     }
 
     pub fn make_partial_encoded_chunk(
diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs
index 5545ca92543..e9f42e7b4fc 100644
--- a/chain/client/src/client.rs
+++ b/chain/client/src/client.rs
@@ -607,6 +607,8 @@ impl Client {
             panic!("The client protocol version is older than the protocol version of the network. Please update nearcore. Client protocol version:{}, network protocol version {}", PROTOCOL_VERSION, protocol_version);
         }
 
+        tracing::debug!(target: "client", ?next_height, ?epoch_id, ?protocol_version, "producing block");
+
         let approvals = self
             .epoch_manager
             .get_epoch_block_approvers_ordered(&prev_hash)?
@@ -2134,6 +2136,7 @@ impl Client {
         state_parts_arbiter_handle: &ArbiterHandle,
     ) -> Result<(), Error> {
         let me = &self.validator_signer.as_ref().map(|x| x.validator_id().clone());
+        tracing::trace!(target: "waclaw", client=?me, "run_catchup");
         for (sync_hash, state_sync_info) in self.chain.store().iterate_state_sync_infos()? {
             assert_eq!(sync_hash, state_sync_info.epoch_tail_hash);
             let network_adapter = self.network_adapter.clone();
diff --git a/chain/client/src/test_utils.rs b/chain/client/src/test_utils.rs
index 2cdd16fa76d..936da6e8af1 100644
--- a/chain/client/src/test_utils.rs
+++ b/chain/client/src/test_utils.rs
@@ -9,6 +9,7 @@ use actix::{Actor, Addr, AsyncContext, Context};
 use actix_rt::{Arbiter, System};
 use chrono::DateTime;
 use futures::{future, FutureExt};
+use itertools::Itertools;
 use near_async::actix::AddrWithAutoSpanContextExt;
 use near_async::messaging::{CanSend, IntoSender, LateBoundSender, Sender};
 use near_async::time;
@@ -832,6 +833,7 @@ pub fn setup_mock_all_validators(
                             );
                         }
                         NetworkRequests::PartialEncodedChunkForward { account_id, forward } => {
+                            tracing::debug!(target: "test", "just checking, is this called?");
                             send_chunks(
                                 connectors1,
                                 validators_clone2.iter().cloned().enumerate(),
@@ -1874,21 +1876,49 @@ impl TestEnv {
 
     pub fn process_partial_encoded_chunks(&mut self) {
         let network_adapters = self.network_adapters.clone();
-        for network_adapter in network_adapters {
-            // process partial encoded chunks
-            while let Some(request) = network_adapter.pop() {
-                if let PeerManagerMessageRequest::NetworkRequests(
-                    NetworkRequests::PartialEncodedChunkMessage {
-                        account_id,
-                        partial_encoded_chunk,
-                    },
-                ) = request
-                {
-                    self.shards_manager(&account_id).send(
-                        ShardsManagerRequestFromNetwork::ProcessPartialEncodedChunk(
-                            PartialEncodedChunk::from(partial_encoded_chunk),
-                        ),
-                    );
+
+        let mut keep_going = true;
+        while keep_going {
+            for (i, network_adapter) in network_adapters.iter().enumerate() {
+                keep_going = false;
+                // process partial encoded chunks
+                while let Some(request) = network_adapter.pop() {
+                    // if there are any requests in any of the adapters reset
+                    // keep going to true as processing of any message may
+                    // trigger more messages to be processed in other clients
+                    // it's a bit sad and it would be much nicer if all messages
+                    // were forwarded to a single queue
+                    keep_going = true;
+                    match request {
+                        PeerManagerMessageRequest::NetworkRequests(
+                            NetworkRequests::PartialEncodedChunkMessage {
+                                account_id,
+                                partial_encoded_chunk,
+                            },
+                        ) => {
+                            tracing::debug!(target: "test", client=i, ?account_id, shard_id=partial_encoded_chunk.header.shard_id(), parts=?partial_encoded_chunk.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "handling partial encoded chunk");
+                            let partial_encoded_chunk =
+                                PartialEncodedChunk::from(partial_encoded_chunk);
+                            let message =
+                                ShardsManagerRequestFromNetwork::ProcessPartialEncodedChunk(
+                                    partial_encoded_chunk,
+                                );
+                            self.shards_manager(&account_id).send(message);
+                        }
+                        PeerManagerMessageRequest::NetworkRequests(
+                            NetworkRequests::PartialEncodedChunkForward { account_id, forward },
+                        ) => {
+                            tracing::debug!(target: "test", client=i, ?account_id, shard_id=forward.shard_id, parts=?forward.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "handling forward");
+                            let message =
+                                ShardsManagerRequestFromNetwork::ProcessPartialEncodedChunkForward(
+                                    forward,
+                                );
+                            self.shards_manager(&account_id).send(message);
+                        }
+                        _ => {
+                            tracing::debug!(target: "test", ?request, "skipping unsupported request type");
+                        }
+                    }
                 }
             }
         }
@@ -2375,6 +2405,7 @@ pub fn run_catchup(
     client: &mut Client,
     highest_height_peers: &[HighestHeightPeerInfo],
 ) -> Result<(), Error> {
+    tracing::trace!(target: "waclaw", "run_catchup");
     let f = |_| {};
     let block_messages = Arc::new(RwLock::new(vec![]));
     let block_inside_messages = block_messages.clone();
diff --git a/chain/network/src/test_utils.rs b/chain/network/src/test_utils.rs
index c2ba09351cb..b83c59c96cf 100644
--- a/chain/network/src/test_utils.rs
+++ b/chain/network/src/test_utils.rs
@@ -241,6 +241,7 @@ impl CanSendAsync<PeerManagerMessageRequest, Result<PeerManagerMessageResponse,
         &self,
         message: PeerManagerMessageRequest,
     ) -> BoxFuture<'static, Result<PeerManagerMessageResponse, ()>> {
+        tracing::trace!(target: "waclaw", "send async");
         self.requests.write().unwrap().push_back(message);
         self.notify.notify_one();
         async { Ok(PeerManagerMessageResponse::NetworkResponses(NetworkResponses::NoResponse)) }
@@ -250,6 +251,7 @@ impl CanSendAsync<PeerManagerMessageRequest, Result<PeerManagerMessageResponse,
 
 impl CanSend<PeerManagerMessageRequest> for MockPeerManagerAdapter {
     fn send(&self, msg: PeerManagerMessageRequest) {
+        tracing::trace!(target: "waclaw", "send sync");
         self.requests.write().unwrap().push_back(msg);
         self.notify.notify_one();
     }
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index 3eed913542d..bede550658c 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -17,7 +17,10 @@ use near_primitives::hash::CryptoHash;
 use near_primitives::serialize::to_base64;
 use near_primitives::shard_layout::{account_id_to_shard_id, account_id_to_shard_uid};
 use near_primitives::transaction::{
-    Action, DeployContractAction, FunctionCallAction, SignedTransaction,
+    Action,
+    FunctionCallAction,
+    SignedTransaction,
+    // Action, DeployContractAction, FunctionCallAction, SignedTransaction,
 };
 use near_primitives::types::{BlockHeight, NumShards, ProtocolVersion, ShardId};
 use near_primitives::utils::MaybeValidated;
@@ -124,6 +127,7 @@ impl TestShardUpgradeEnv {
             .validator_seats(num_validators)
             .real_epoch_managers(&genesis.config)
             .nightshade_runtimes(&genesis)
+            .track_all_shards()
             .build();
         assert_eq!(env.validators.len(), num_validators);
         Self {
@@ -206,6 +210,7 @@ impl TestShardUpgradeEnv {
             let block_producer_client = env.client(&block_producer);
             let mut block = block_producer_client.produce_block(height).unwrap().unwrap();
             set_block_protocol_version(&mut block, block_producer.clone(), protocol_version);
+            // std::thread::sleep(std::time::Duration::from_millis(2000));
             block
         };
         // Make sure that catchup is done before the end of each epoch, but when it is done is
@@ -222,13 +227,13 @@ impl TestShardUpgradeEnv {
             // Here we don't just call self.clients[i].process_block_sync_with_produce_chunk_options
             // because we want to call run_catchup before finish processing this block. This simulates
             // that catchup and block processing run in parallel.
-            client
-                .start_process_block(
-                    MaybeValidated::from(block.clone()),
-                    Provenance::NONE,
-                    Arc::new(|_| {}),
-                )
-                .unwrap();
+            let result = client.start_process_block(
+                MaybeValidated::from(block.clone()),
+                Provenance::NONE,
+                Arc::new(|_| {}),
+            );
+            tracing::debug!(target: "test", start_process_block_result=?result);
+            result.unwrap();
             if should_catchup {
                 run_catchup(client, &[]).unwrap();
             }
@@ -247,11 +252,16 @@ impl TestShardUpgradeEnv {
                 .get_shard_layout_from_prev_block(block.hash())
                 .unwrap()
                 .num_shards();
-            assert_eq!(num_shards, expected_num_shards);
+            tracing::trace!(target: "waclaw", ?num_shards, ?expected_num_shards, "boom");
+            // assert_eq!(num_shards, expected_num_shards);
         }
 
         env.process_partial_encoded_chunks();
         for j in 0..self.num_clients {
+            let _span =
+                tracing::debug_span!(target: "test", "process shard manager things", client=j);
+            let _span = _span.entered();
+
             env.process_shards_manager_responses_and_finish_processing_blocks(j);
         }
 
@@ -639,11 +649,11 @@ fn setup_genesis(
     // TODO(resharding)
     // In the current simple test setup each validator may track different
     // shards. Unfortunately something is broken in how state sync is triggered
-    // from this test. The config below forces both validators to track all
+    // from this test. The config below forces all validators to track all
     // shards so that we can avoid that problem for now. It would be nice to
     // set it up properly and test both state sync and resharding - once the
     // integration is fully supported.
-    genesis.config.minimum_validators_per_shard = num_validators;
+    // genesis.config.minimum_validators_per_shard = num_validators;
 
     genesis
 }
@@ -856,11 +866,12 @@ fn setup_test_env_with_cross_contract_txs(
     epoch_length: u64,
     genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
+    let num = 4;
     // let mut test_env = TestShardUpgradeEnv::new(epoch_length, 4, 4, 100, Some(100_000_000_000_000));
-    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
+    let test_env = TestShardUpgradeEnv::new_with_protocol_version(
         epoch_length,
-        4,
-        4,
+        num,
+        num,
         100,
         Some(100_000_000_000_000),
         genesis_protocol_version,
@@ -966,11 +977,13 @@ fn test_shard_layout_upgrade_cross_contract_calls_impl(resharding_type: Reshardi
     let genesis_protocol_version = get_genesis_protocol_version(&resharding_type);
     let target_protocol_version = get_target_protocol_version(&resharding_type);
 
+    tracing::trace!(target: "waclaw",genesis_protocol_version, target_protocol_version, "protocol version");
+
     let (mut test_env, new_accounts) =
         setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 
     for _ in 1..5 * epoch_length {
-        std::thread::sleep(std::time::Duration::from_millis(1000));
+        // std::thread::sleep(std::time::Duration::from_millis(2000));
         test_env.step_impl(0., target_protocol_version, &resharding_type);
         test_env.check_receipt_id_to_shard_id();
     }
diff --git a/r.sh b/r.sh
index 8dd7c4d65f0..b27ae18b39c 100755
--- a/r.sh
+++ b/r.sh
@@ -1,6 +1,13 @@
+
+clear
+tmux clear-history
+
 rm stdout-*
 
- for i in {1..1}
+TEST=test_shard_layout_upgrade_simple_v2
+TEST=test_shard_layout_upgrade_cross_contract_calls_v2
+
+for i in {1..1}
  do
      echo "test $i"
      OUT=stdout-$i
@@ -8,15 +15,15 @@ rm stdout-*
      # RUST_LOG=info \
      #RUST_LOG=info,catchup=trace,store=trace,client=debug,store=debug,test=debug,resharding=trace \
      RUST_BACKTRACE=all \
-     RUST_LOG=debug,resharding=trace \
+     RUST_LOG=debug,resharding=trace,waclaw=trace,catchup=trace,resharding=trace \
      cargo nextest run -p integration-tests \
          --no-capture \
          --features nightly \
-		test_shard_layout_upgrade_cross_contract_calls_v2 \
+		 $TEST \
+         | egrep -v prev_prev_stake_change \
          > $OUT
          # | tee $OUT
          # | egrep -v -i "FlatStorage is not ready|Add delta for flat storage creation|epoch_manager: all proposals" \
-         # test_shard_layout_upgrade_simple_v1 \
 
      sed -E -i 's/ed25519:(.{4})(.{40})/ed25519:\1/g' $OUT
 

From 0302b063fbd7cf2c61dcb3851a2f319e901f41c8 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Tue, 8 Aug 2023 11:13:05 +0000
Subject: [PATCH 06/15] cleanup some debug logs

---
 chain/chain/src/chain.rs                      | 19 ++--
 chain/chunks/src/lib.rs                       |  5 -
 chain/client/src/client.rs                    |  1 -
 chain/client/src/test_utils.rs                |  1 -
 chain/network/src/test_utils.rs               |  2 -
 .../src/tests/client/sharding_upgrade.rs      | 95 +++++++++----------
 p.sh                                          | 11 +++
 r.sh                                          |  3 +-
 8 files changed, 66 insertions(+), 71 deletions(-)
 create mode 100755 p.sh

diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index 4e6c7e5d802..aa43b205ae6 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -1120,8 +1120,6 @@ impl Chain {
             &prev_hash,
         )?;
         let prev_block = self.get_block(&prev_hash)?;
-        tracing::trace!(target: "waclaw", ?shards_to_state_sync, "get_state_sync_info");
-
         if prev_block.chunks().len() != block.chunks().len() && !shards_to_state_sync.is_empty() {
             // Currently, the state sync algorithm assumes that the number of chunks do not change
             // between the epoch being synced to and the last epoch.
@@ -1643,7 +1641,6 @@ impl Chain {
                 let chunk_hash = chunk_header.chunk_hash();
 
                 if let Err(_) = self.store.get_partial_chunk(&chunk_header.chunk_hash()) {
-                    tracing::debug!(target: "waclaw", "get partial chunk failed");
                     missing.push(chunk_header.clone());
                 } else if self.shard_tracker.care_about_shard(
                     me.as_ref(),
@@ -1657,7 +1654,6 @@ impl Chain {
                     true,
                 ) {
                     if let Err(_) = self.store.get_chunk(&chunk_hash) {
-                        tracing::debug!(target: "waclaw", "get chunk failed");
                         missing.push(chunk_header.clone());
                     }
                 }
@@ -2108,7 +2104,6 @@ impl Chain {
             block_received_time,
             state_patch,
         );
-        tracing::debug!(target: "waclaw", preprocess_res=preprocess_res.is_ok(), preprocess_res_err=?preprocess_res.as_ref().err(), "preprocess_res");
         let preprocess_res = match preprocess_res {
             Ok(preprocess_res) => {
                 preprocess_timer.observe_duration();
@@ -2486,7 +2481,6 @@ impl Chain {
         debug!(target: "chain", block_hash = ?header.hash(), me=?me, is_caught_up=is_caught_up, "Process block");
 
         // Check the header is valid before we proceed with the full block.
-        tracing::debug!(target: "waclaw", "validate header");
         self.validate_header(header, provenance, challenges)?;
 
         self.epoch_manager.verify_block_vrf(
@@ -2501,7 +2495,6 @@ impl Chain {
             return Err(Error::InvalidRandomnessBeaconOutput);
         }
 
-        tracing::debug!(target: "waclaw", "validate with");
         let res = block.validate_with(|block| {
             Chain::validate_block_impl(self.epoch_manager.as_ref(), &self.genesis, block)
                 .map(|_| true)
@@ -2537,13 +2530,10 @@ impl Chain {
 
         let prev_block = self.get_block(&prev_hash)?;
 
-        tracing::debug!(target: "waclaw", "validate chunk headers");
         self.validate_chunk_headers(&block, &prev_block)?;
 
-        tracing::debug!(target: "waclaw", "ping missing chunks");
         self.ping_missing_chunks(me, prev_hash, block)?;
 
-        tracing::debug!(target: "waclaw", "collect incoming receipts from block");
         let incoming_receipts = self.collect_incoming_receipts_from_block(me, block)?;
 
         // Check if block can be finalized and drop it otherwise.
@@ -2725,8 +2715,13 @@ impl Chain {
         let will_shard_layout_change = match result {
             Ok(will_shard_layout_change) => will_shard_layout_change,
             Err(err) => {
-                tracing::error!(target: "waclaw", ?err, "failed to check if shard layout will change");
-                return false;
+                // This is a problem if happens continuously. If a node fails to
+                // reshard it will fall behind the network once it switches over
+                // to the new shard layout. Luckily this method is called
+                // periodically so the node has a chance to recover if it only
+                // happens a few times.
+                tracing::error!(target: "chain", ?err, "failed to check if shard layout will change");
+                false
             }
         };
         // if shard layout will change the next epoch, we should catch up the shard regardless
diff --git a/chain/chunks/src/lib.rs b/chain/chunks/src/lib.rs
index b36dd1e7c3b..3481a5078c8 100644
--- a/chain/chunks/src/lib.rs
+++ b/chain/chunks/src/lib.rs
@@ -1306,13 +1306,11 @@ impl ShardsManager {
         &mut self,
         header: &ShardChunkHeader,
     ) -> bool {
-        tracing::debug!(target: "waclaw", "insert_header_if_not_exists_and_process_cached_chunk_forwards");
         let header_known_before = self.encoded_chunks.get(&header.chunk_hash()).is_some();
         if self.encoded_chunks.get_or_insert_from_header(header).complete {
             return false;
         }
         if let Some(parts) = self.chunk_forwards_cache.pop(&header.chunk_hash()) {
-            tracing::debug!(target: "waclaw", parts_len=parts.len(), "insert_header_if_not_exists_and_process_cached_chunk_forwards");
             // Note that we don't need any further validation for the forwarded part.
             // The forwarded part was earlier validated via validate_partial_encoded_chunk_forward,
             // which checks the part against the merkle root in the forward message, and the merkle
@@ -1736,8 +1734,6 @@ impl ShardsManager {
                 continue;
             }
             next_chunk_producers.remove(&bp_account_id);
-
-            tracing::debug!(target: "waclaw", ?bp_account_id, shard_id=partial_encoded_chunk.header.shard_id(), parts=?forward.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "sending forward");
             // Technically, here we should check if the block producer actually cares about the shard.
             // We don't because with the current implementation, we force all validators to track all
             // shards by making their config tracking all shards.
@@ -1910,7 +1906,6 @@ impl ShardsManager {
                 );
 
             if Some(&to_whom) != self.me.as_ref() {
-                tracing::trace!(target: "waclaw", parts_len=partial_encoded_chunk.parts.len(), "sending partial encoded chunk");
                 self.peer_manager_adapter.send(PeerManagerMessageRequest::NetworkRequests(
                     NetworkRequests::PartialEncodedChunkMessage {
                         account_id: to_whom.clone(),
diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs
index e9f42e7b4fc..594b6d9417b 100644
--- a/chain/client/src/client.rs
+++ b/chain/client/src/client.rs
@@ -2136,7 +2136,6 @@ impl Client {
         state_parts_arbiter_handle: &ArbiterHandle,
     ) -> Result<(), Error> {
         let me = &self.validator_signer.as_ref().map(|x| x.validator_id().clone());
-        tracing::trace!(target: "waclaw", client=?me, "run_catchup");
         for (sync_hash, state_sync_info) in self.chain.store().iterate_state_sync_infos()? {
             assert_eq!(sync_hash, state_sync_info.epoch_tail_hash);
             let network_adapter = self.network_adapter.clone();
diff --git a/chain/client/src/test_utils.rs b/chain/client/src/test_utils.rs
index 936da6e8af1..f395270698c 100644
--- a/chain/client/src/test_utils.rs
+++ b/chain/client/src/test_utils.rs
@@ -2405,7 +2405,6 @@ pub fn run_catchup(
     client: &mut Client,
     highest_height_peers: &[HighestHeightPeerInfo],
 ) -> Result<(), Error> {
-    tracing::trace!(target: "waclaw", "run_catchup");
     let f = |_| {};
     let block_messages = Arc::new(RwLock::new(vec![]));
     let block_inside_messages = block_messages.clone();
diff --git a/chain/network/src/test_utils.rs b/chain/network/src/test_utils.rs
index b83c59c96cf..c2ba09351cb 100644
--- a/chain/network/src/test_utils.rs
+++ b/chain/network/src/test_utils.rs
@@ -241,7 +241,6 @@ impl CanSendAsync<PeerManagerMessageRequest, Result<PeerManagerMessageResponse,
         &self,
         message: PeerManagerMessageRequest,
     ) -> BoxFuture<'static, Result<PeerManagerMessageResponse, ()>> {
-        tracing::trace!(target: "waclaw", "send async");
         self.requests.write().unwrap().push_back(message);
         self.notify.notify_one();
         async { Ok(PeerManagerMessageResponse::NetworkResponses(NetworkResponses::NoResponse)) }
@@ -251,7 +250,6 @@ impl CanSendAsync<PeerManagerMessageRequest, Result<PeerManagerMessageResponse,
 
 impl CanSend<PeerManagerMessageRequest> for MockPeerManagerAdapter {
     fn send(&self, msg: PeerManagerMessageRequest) {
-        tracing::trace!(target: "waclaw", "send sync");
         self.requests.write().unwrap().push_back(msg);
         self.notify.notify_one();
     }
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index bede550658c..e04b8a956b4 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -17,10 +17,7 @@ use near_primitives::hash::CryptoHash;
 use near_primitives::serialize::to_base64;
 use near_primitives::shard_layout::{account_id_to_shard_id, account_id_to_shard_uid};
 use near_primitives::transaction::{
-    Action,
-    FunctionCallAction,
-    SignedTransaction,
-    // Action, DeployContractAction, FunctionCallAction, SignedTransaction,
+    Action, DeployContractAction, FunctionCallAction, SignedTransaction,
 };
 use near_primitives::types::{BlockHeight, NumShards, ProtocolVersion, ShardId};
 use near_primitives::utils::MaybeValidated;
@@ -867,8 +864,7 @@ fn setup_test_env_with_cross_contract_txs(
     genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
     let num = 4;
-    // let mut test_env = TestShardUpgradeEnv::new(epoch_length, 4, 4, 100, Some(100_000_000_000_000));
-    let test_env = TestShardUpgradeEnv::new_with_protocol_version(
+    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
         epoch_length,
         num,
         num,
@@ -888,33 +884,33 @@ fn setup_test_env_with_cross_contract_txs(
         test_env.initial_accounts[indices[0]].clone(),
         test_env.initial_accounts[indices[1]].clone(),
     ];
-    // test_env.set_init_tx(
-    //     contract_accounts
-    //         .iter()
-    //         .map(|account_id| {
-    //             let signer = InMemorySigner::from_seed(
-    //                 account_id.clone(),
-    //                 KeyType::ED25519,
-    //                 &account_id.to_string(),
-    //             );
-    //             SignedTransaction::from_actions(
-    //                 1,
-    //                 account_id.clone(),
-    //                 account_id.clone(),
-    //                 &signer,
-    //                 vec![Action::DeployContract(DeployContractAction {
-    //                     code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
-    //                 })],
-    //                 genesis_hash,
-    //             )
-    //         })
-    //         .collect(),
-    // );
+    test_env.set_init_tx(
+        contract_accounts
+            .iter()
+            .map(|account_id| {
+                let signer = InMemorySigner::from_seed(
+                    account_id.clone(),
+                    KeyType::ED25519,
+                    &account_id.to_string(),
+                );
+                SignedTransaction::from_actions(
+                    1,
+                    account_id.clone(),
+                    account_id.clone(),
+                    &signer,
+                    vec![Action::DeployContract(DeployContractAction {
+                        code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
+                    })],
+                    genesis_hash,
+                )
+            })
+            .collect(),
+    );
 
     let mut nonce = 100;
     let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
     let mut new_accounts = HashMap::new();
-    let _generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
+    let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
         &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
             let mut rng = thread_rng();
             let size = rng.gen_range(min_size..max_size + 1);
@@ -943,26 +939,26 @@ fn setup_test_env_with_cross_contract_txs(
             .collect()
         };
 
-    // // add a bunch of transactions before the two epoch boundaries
-    // for height in vec![
-    //     epoch_length - 2,
-    //     epoch_length - 1,
-    //     epoch_length,
-    //     2 * epoch_length - 2,
-    //     2 * epoch_length - 1,
-    //     2 * epoch_length,
-    // ] {
-    //     test_env.set_tx_at_height(height, generate_txs(5, 8));
-    // }
-
-    // // adds some transactions after sharding change finishes
-    // // but do not add too many because I want all transactions to
-    // // finish processing before epoch 5
-    // for height in 2 * epoch_length + 1..3 * epoch_length {
-    //     if rng.gen_bool(0.3) {
-    //         test_env.set_tx_at_height(height, generate_txs(5, 8));
-    //     }
-    // }
+    // add a bunch of transactions before the two epoch boundaries
+    for height in vec![
+        epoch_length - 2,
+        epoch_length - 1,
+        epoch_length,
+        2 * epoch_length - 2,
+        2 * epoch_length - 1,
+        2 * epoch_length,
+    ] {
+        test_env.set_tx_at_height(height, generate_txs(5, 8));
+    }
+
+    // adds some transactions after sharding change finishes
+    // but do not add too many because I want all transactions to
+    // finish processing before epoch 5
+    for height in 2 * epoch_length + 1..3 * epoch_length {
+        if rng.gen_bool(0.3) {
+            test_env.set_tx_at_height(height, generate_txs(5, 8));
+        }
+    }
 
     (test_env, new_accounts)
 }
@@ -1006,6 +1002,7 @@ fn test_shard_layout_upgrade_cross_contract_calls_v1() {
 
 // Test cross contract calls
 // This test case tests postponed receipts and delayed receipts
+#[cfg(feature = "protocol_feature_simple_nightshade_v2")]
 #[test]
 fn test_shard_layout_upgrade_cross_contract_calls_v2() {
     test_shard_layout_upgrade_cross_contract_calls_impl(ReshardingType::V2);
diff --git a/p.sh b/p.sh
new file mode 100755
index 00000000000..a2ace2b6609
--- /dev/null
+++ b/p.sh
@@ -0,0 +1,11 @@
+
+clear
+tmux clear-history
+
+RUST_BACKTRACE=all \
+RUST_LOG=debug,resharding=trace \
+cargo nextest run \
+    --package integration-tests \
+    --features nightly \
+    sharding_upgrade
+
diff --git a/r.sh b/r.sh
index b27ae18b39c..80133092509 100755
--- a/r.sh
+++ b/r.sh
@@ -14,8 +14,9 @@ for i in {1..1}
 
      # RUST_LOG=info \
      #RUST_LOG=info,catchup=trace,store=trace,client=debug,store=debug,test=debug,resharding=trace \
+     # RUST_LOG=debug,resharding=trace,waclaw=trace,catchup=trace \
      RUST_BACKTRACE=all \
-     RUST_LOG=debug,resharding=trace,waclaw=trace,catchup=trace,resharding=trace \
+     RUST_LOG=debug,resharding=trace \
      cargo nextest run -p integration-tests \
          --no-capture \
          --features nightly \

From 6b8b0ca8aa60ea7ec739c80d649e16a0c1ca4bf6 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Tue, 8 Aug 2023 11:14:41 +0000
Subject: [PATCH 07/15] reenable asserts - broken again

---
 integration-tests/src/tests/client/sharding_upgrade.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index e04b8a956b4..c2126cd10a1 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -249,8 +249,7 @@ impl TestShardUpgradeEnv {
                 .get_shard_layout_from_prev_block(block.hash())
                 .unwrap()
                 .num_shards();
-            tracing::trace!(target: "waclaw", ?num_shards, ?expected_num_shards, "boom");
-            // assert_eq!(num_shards, expected_num_shards);
+            assert_eq!(num_shards, expected_num_shards);
         }
 
         env.process_partial_encoded_chunks();

From bb682a5f0db7bcafe7aad1aaf49a643685178df6 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Tue, 8 Aug 2023 11:19:03 +0000
Subject: [PATCH 08/15] no tx

---
 .../src/tests/client/sharding_upgrade.rs      | 181 +++++++++---------
 1 file changed, 91 insertions(+), 90 deletions(-)

diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index c2126cd10a1..0037bd9cb1a 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -1,3 +1,5 @@
+#![allow(unused)]
+
 use borsh::BorshSerialize;
 use near_client::{Client, ProcessTxResponse};
 use near_primitives::epoch_manager::{AllEpochConfig, EpochConfig};
@@ -863,7 +865,7 @@ fn setup_test_env_with_cross_contract_txs(
     genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
     let num = 4;
-    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
+    let test_env = TestShardUpgradeEnv::new_with_protocol_version(
         epoch_length,
         num,
         num,
@@ -871,93 +873,94 @@ fn setup_test_env_with_cross_contract_txs(
         Some(100_000_000_000_000),
         genesis_protocol_version,
     );
-    let mut rng = thread_rng();
-
-    let genesis_hash = *test_env.env.clients[0].chain.genesis_block().hash();
-    // Use test0, test1 and two random accounts to deploy contracts because we want accounts on
-    // different shards.
-    let indices = (4..test_env.initial_accounts.len()).choose_multiple(&mut rng, 2);
-    let contract_accounts = vec![
-        test_env.initial_accounts[0].clone(),
-        test_env.initial_accounts[1].clone(),
-        test_env.initial_accounts[indices[0]].clone(),
-        test_env.initial_accounts[indices[1]].clone(),
-    ];
-    test_env.set_init_tx(
-        contract_accounts
-            .iter()
-            .map(|account_id| {
-                let signer = InMemorySigner::from_seed(
-                    account_id.clone(),
-                    KeyType::ED25519,
-                    &account_id.to_string(),
-                );
-                SignedTransaction::from_actions(
-                    1,
-                    account_id.clone(),
-                    account_id.clone(),
-                    &signer,
-                    vec![Action::DeployContract(DeployContractAction {
-                        code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
-                    })],
-                    genesis_hash,
-                )
-            })
-            .collect(),
-    );
-
-    let mut nonce = 100;
-    let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
-    let mut new_accounts = HashMap::new();
-    let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
-        &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
-            let mut rng = thread_rng();
-            let size = rng.gen_range(min_size..max_size + 1);
-            std::iter::repeat_with(|| loop {
-                let account_id = gen_account(&mut rng, b"abcdefghijkmn");
-                if all_accounts.insert(account_id.clone()) {
-                    nonce += 1;
-                    // randomly shuffle contract accounts
-                    // so that transactions are send to different shards
-                    let mut contract_accounts = contract_accounts.clone();
-                    contract_accounts.shuffle(&mut rng);
-                    let tx = gen_cross_contract_transaction(
-                        &contract_accounts[0],
-                        &contract_accounts[1],
-                        &contract_accounts[2],
-                        &contract_accounts[3],
-                        &account_id,
-                        nonce,
-                        &genesis_hash,
-                    );
-                    new_accounts.insert(tx.get_hash(), account_id);
-                    return tx;
-                }
-            })
-            .take(size)
-            .collect()
-        };
-
-    // add a bunch of transactions before the two epoch boundaries
-    for height in vec![
-        epoch_length - 2,
-        epoch_length - 1,
-        epoch_length,
-        2 * epoch_length - 2,
-        2 * epoch_length - 1,
-        2 * epoch_length,
-    ] {
-        test_env.set_tx_at_height(height, generate_txs(5, 8));
-    }
-
-    // adds some transactions after sharding change finishes
-    // but do not add too many because I want all transactions to
-    // finish processing before epoch 5
-    for height in 2 * epoch_length + 1..3 * epoch_length {
-        if rng.gen_bool(0.3) {
-            test_env.set_tx_at_height(height, generate_txs(5, 8));
-        }
-    }
+    // let mut rng = thread_rng();
+
+    // let genesis_hash = *test_env.env.clients[0].chain.genesis_block().hash();
+    // // Use test0, test1 and two random accounts to deploy contracts because we want accounts on
+    // // different shards.
+    // let indices = (4..test_env.initial_accounts.len()).choose_multiple(&mut rng, 2);
+    // let contract_accounts = vec![
+    //     test_env.initial_accounts[0].clone(),
+    //     test_env.initial_accounts[1].clone(),
+    //     test_env.initial_accounts[indices[0]].clone(),
+    //     test_env.initial_accounts[indices[1]].clone(),
+    // ];
+    // test_env.set_init_tx(
+    //     contract_accounts
+    //         .iter()
+    //         .map(|account_id| {
+    //             let signer = InMemorySigner::from_seed(
+    //                 account_id.clone(),
+    //                 KeyType::ED25519,
+    //                 &account_id.to_string(),
+    //             );
+    //             SignedTransaction::from_actions(
+    //                 1,
+    //                 account_id.clone(),
+    //                 account_id.clone(),
+    //                 &signer,
+    //                 vec![Action::DeployContract(DeployContractAction {
+    //                     code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
+    //                 })],
+    //                 genesis_hash,
+    //             )
+    //         })
+    //         .collect(),
+    // );
+
+    let new_accounts = HashMap::new();
+    // let mut nonce = 100;
+    // let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
+    // let mut new_accounts = HashMap::new();
+    // let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
+    //     &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
+    //         let mut rng = thread_rng();
+    //         let size = rng.gen_range(min_size..max_size + 1);
+    //         std::iter::repeat_with(|| loop {
+    //             let account_id = gen_account(&mut rng, b"abcdefghijkmn");
+    //             if all_accounts.insert(account_id.clone()) {
+    //                 nonce += 1;
+    //                 // randomly shuffle contract accounts
+    //                 // so that transactions are send to different shards
+    //                 let mut contract_accounts = contract_accounts.clone();
+    //                 contract_accounts.shuffle(&mut rng);
+    //                 let tx = gen_cross_contract_transaction(
+    //                     &contract_accounts[0],
+    //                     &contract_accounts[1],
+    //                     &contract_accounts[2],
+    //                     &contract_accounts[3],
+    //                     &account_id,
+    //                     nonce,
+    //                     &genesis_hash,
+    //                 );
+    //                 new_accounts.insert(tx.get_hash(), account_id);
+    //                 return tx;
+    //             }
+    //         })
+    //         .take(size)
+    //         .collect()
+    //     };
+
+    // // add a bunch of transactions before the two epoch boundaries
+    // for height in vec![
+    //     epoch_length - 2,
+    //     epoch_length - 1,
+    //     epoch_length,
+    //     2 * epoch_length - 2,
+    //     2 * epoch_length - 1,
+    //     2 * epoch_length,
+    // ] {
+    //     test_env.set_tx_at_height(height, generate_txs(5, 8));
+    // }
+
+    // // adds some transactions after sharding change finishes
+    // // but do not add too many because I want all transactions to
+    // // finish processing before epoch 5
+    // for height in 2 * epoch_length + 1..3 * epoch_length {
+    //     if rng.gen_bool(0.3) {
+    //         test_env.set_tx_at_height(height, generate_txs(5, 8));
+    //     }
+    // }
 
     (test_env, new_accounts)
 }
@@ -972,8 +975,6 @@ fn test_shard_layout_upgrade_cross_contract_calls_impl(resharding_type: Reshardi
     let genesis_protocol_version = get_genesis_protocol_version(&resharding_type);
     let target_protocol_version = get_target_protocol_version(&resharding_type);
 
-    tracing::trace!(target: "waclaw",genesis_protocol_version, target_protocol_version, "protocol version");
-
     let (mut test_env, new_accounts) =
         setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 

From d2b6e3f20c0713ebc4bfb6e57214efed3aeae8dd Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 9 Aug 2023 10:43:58 +0000
Subject: [PATCH 09/15] protocol version upgrade works predictably

---
 Cargo.lock                                    |  1 +
 chain/chain/src/chain.rs                      |  7 +--
 chain/epoch-manager/Cargo.toml                |  5 +-
 chain/epoch-manager/src/lib.rs                | 46 ++++++++++++++++---
 .../epoch-manager/src/validator_selection.rs  |  9 ++++
 core/primitives/src/epoch_manager.rs          |  3 ++
 core/primitives/src/rand.rs                   | 21 +++++++++
 core/store/src/trie/state_parts.rs            |  2 +-
 .../src/tests/client/sharding_upgrade.rs      | 40 ++++++++++++++--
 9 files changed, 118 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8d916f43382..fbad8345040 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3721,6 +3721,7 @@ version = "0.0.0"
 dependencies = [
  "borsh 0.10.2",
  "chrono",
+ "itertools",
  "near-cache",
  "near-chain-configs",
  "near-chain-primitives",
diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index aa43b205ae6..63bf88ac236 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -2715,11 +2715,8 @@ impl Chain {
         let will_shard_layout_change = match result {
             Ok(will_shard_layout_change) => will_shard_layout_change,
             Err(err) => {
-                // This is a problem if happens continuously. If a node fails to
-                // reshard it will fall behind the network once it switches over
-                // to the new shard layout. Luckily this method is called
-                // periodically so the node has a chance to recover if it only
-                // happens a few times.
+                // TODO(resharding) This is a problem, if this happens the node
+                // will not perform resharding and fall behind the network.
                 tracing::error!(target: "chain", ?err, "failed to check if shard layout will change");
                 false
             }
diff --git a/chain/epoch-manager/Cargo.toml b/chain/epoch-manager/Cargo.toml
index e23aee47676..c318bbe0eb9 100644
--- a/chain/epoch-manager/Cargo.toml
+++ b/chain/epoch-manager/Cargo.toml
@@ -18,6 +18,7 @@ rand_hc.workspace = true
 serde_json.workspace = true
 smart-default.workspace = true
 tracing.workspace = true
+itertools.workspace = true
 
 near-crypto.workspace = true
 near-primitives.workspace = true
@@ -28,7 +29,9 @@ near-cache.workspace = true
 
 [features]
 expensive_tests = []
-protocol_feature_fix_staking_threshold = ["near-primitives/protocol_feature_fix_staking_threshold"]
+protocol_feature_fix_staking_threshold = [
+  "near-primitives/protocol_feature_fix_staking_threshold",
+]
 nightly = [
   "nightly_protocol",
   "protocol_feature_fix_staking_threshold",
diff --git a/chain/epoch-manager/src/lib.rs b/chain/epoch-manager/src/lib.rs
index bcaa23a103a..cacc802e8c2 100644
--- a/chain/epoch-manager/src/lib.rs
+++ b/chain/epoch-manager/src/lib.rs
@@ -46,6 +46,8 @@ mod tests;
 pub mod types;
 mod validator_selection;
 
+use itertools::Itertools;
+
 const EPOCH_CACHE_SIZE: usize = if cfg!(feature = "no_cache") { 1 } else { 50 };
 const BLOCK_CACHE_SIZE: usize = if cfg!(feature = "no_cache") { 5 } else { 1000 }; // TODO(#5080): fix this
 const AGGREGATOR_SAVE_PERIOD: u64 = 1000;
@@ -215,6 +217,10 @@ impl EpochManager {
                 genesis_protocol_version,
                 genesis_protocol_version,
             )?;
+
+            tracing::info!(block_producers_settlement=?epoch_info.block_producers_settlement(), "epoch manager new");
+            tracing::info!(chunk_producers_settlement=?epoch_info.chunk_producers_settlement(), "epoch manager new");
+            tracing::info!(validators=?epoch_info.validators_iter().collect_vec(), "epoch manager new");
             // Dummy block info.
             // Artificial block we add to simplify implementation: dummy block is the
             // parent of genesis block that points to itself.
@@ -532,6 +538,12 @@ impl EpochManager {
         for (validator_id, version) in version_tracker {
             let stake = epoch_info.validator_stake(validator_id);
             *versions.entry(version).or_insert(0) += stake;
+            tracing::info!(
+                validator_id,
+                version,
+                stake,
+                "voting on the next next epoch protocol version"
+            );
         }
         let total_block_producer_stake: u128 = epoch_info
             .block_producers_settlement()
@@ -553,20 +565,34 @@ impl EpochManager {
         // Note: non-deterministic iteration is fine here, there can be only one
         // version with large enough stake.
         let next_version = if let Some((version, stake)) =
-            versions.into_iter().max_by_key(|&(_version, stake)| stake)
+            versions.clone().into_iter().max_by_key(|&(_version, stake)| stake)
         {
-            if stake
-                > (total_block_producer_stake
-                    * *config.protocol_upgrade_stake_threshold.numer() as u128)
-                    / *config.protocol_upgrade_stake_threshold.denom() as u128
-            {
+            let numer = *config.protocol_upgrade_stake_threshold.numer() as u128;
+            let denom = *config.protocol_upgrade_stake_threshold.denom() as u128;
+            let threshold = total_block_producer_stake * numer / denom;
+            tracing::info!(
+                stake,
+                threshold,
+                numer,
+                denom,
+                total_block_producer_stake,
+                "voting on the next next epoch protocol version"
+            );
+            if stake > threshold {
                 version
             } else {
                 protocol_version
             }
         } else {
+            tracing::info!("voting on the next next epoch protocol version winner - no max");
+
             protocol_version
         };
+        tracing::info!(
+            ?next_version,
+            ?versions,
+            "voting on the next next epoch protocol version winner"
+        );
 
         // Gather slashed validators and add them to kick out first.
         let slashed_validators = last_block_info.slashed();
@@ -632,6 +658,8 @@ impl EpochManager {
         let next_epoch_info = self.get_epoch_info(&next_epoch_id)?;
         self.save_epoch_validator_info(store_update, block_info.epoch_id(), &epoch_summary)?;
 
+        tracing::info!(?validator_stake, "finalize epoch");
+
         let EpochSummary {
             all_proposals,
             validator_kickout,
@@ -1744,6 +1772,9 @@ impl EpochManager {
 
         let mut aggregator = EpochInfoAggregator::new(epoch_id.clone(), *block_hash);
         let mut cur_hash = *block_hash;
+
+        tracing::info!(?block_hash, "aggregate_epoch_info_upto");
+
         Ok(Some(loop {
             #[cfg(test)]
             {
@@ -1765,6 +1796,8 @@ impl EpochManager {
                 // belongs to different epoch or we’re on different fork (though
                 // the latter should never happen).  In either case, the
                 // aggregator contains full epoch information.
+                tracing::info!(?block_hash, version_tracker=?aggregator.version_tracker, "aggregate_epoch_info_upto break epoch or genesis");
+
                 break (aggregator, true);
             }
 
@@ -1779,6 +1812,7 @@ impl EpochManager {
                 // We’ve reached sync point of the old aggregator.  If old
                 // aggregator was for a different epoch, we have full info in
                 // our aggregator; otherwise we don’t.
+                tracing::info!(?block_hash, version_tracker=?aggregator.version_tracker, "aggregate_epoch_info_upto break last block hash");
                 break (aggregator, epoch_id != prev_epoch);
             }
 
diff --git a/chain/epoch-manager/src/validator_selection.rs b/chain/epoch-manager/src/validator_selection.rs
index 41499687e22..fbb9945a74d 100644
--- a/chain/epoch-manager/src/validator_selection.rs
+++ b/chain/epoch-manager/src/validator_selection.rs
@@ -53,6 +53,7 @@ pub fn proposals_to_epoch_info(
         min_stake_ratio,
         last_version,
     );
+    tracing::info!(?block_producers, "select_block_producers");
     let (chunk_producer_proposals, chunk_producers, cp_stake_threshold) =
         if checked_feature!("stable", ChunkOnlyProducers, next_version) {
             let mut chunk_producer_proposals = order_proposals(proposals.into_values());
@@ -107,6 +108,13 @@ pub fn proposals_to_epoch_info(
         all_validators.push(bp);
     }
 
+    tracing::info!(
+        ?validator_to_index,
+        ?block_producers_settlement,
+        ?all_validators,
+        "select_block_producers"
+    );
+
     let chunk_producers_settlement = if checked_feature!("stable", ChunkOnlyProducers, next_version)
     {
         let minimum_validators_per_shard =
@@ -118,6 +126,7 @@ pub fn proposals_to_epoch_info(
                     num_shards,
                 },
             )?;
+        tracing::info!(?shard_assignment, "select_block_producers");
 
         let mut chunk_producers_settlement: Vec<Vec<ValidatorId>> =
             shard_assignment.iter().map(|vs| Vec::with_capacity(vs.len())).collect();
diff --git a/core/primitives/src/epoch_manager.rs b/core/primitives/src/epoch_manager.rs
index f12672f7cd6..a5f43cd8e17 100644
--- a/core/primitives/src/epoch_manager.rs
+++ b/core/primitives/src/epoch_manager.rs
@@ -878,14 +878,17 @@ pub mod epoch_info {
         pub fn sample_block_producer(&self, height: BlockHeight) -> ValidatorId {
             match &self {
                 Self::V1(v1) => {
+                    tracing::info!("sample_block_producer v1");
                     let bp_settlement = &v1.block_producers_settlement;
                     bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
                 }
                 Self::V2(v2) => {
+                    tracing::info!("sample_block_producer v2");
                     let bp_settlement = &v2.block_producers_settlement;
                     bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
                 }
                 Self::V3(v3) => {
+                    tracing::info!(block_producers_settlement=?v3.block_producers_settlement, "sample_block_producer v3");
                     let seed = Self::block_produce_seed(height, &v3.rng_seed);
                     v3.block_producers_settlement[v3.block_producers_sampler.sample(seed)]
                 }
diff --git a/core/primitives/src/rand.rs b/core/primitives/src/rand.rs
index a79c8fd1b56..735ff2ef887 100644
--- a/core/primitives/src/rand.rs
+++ b/core/primitives/src/rand.rs
@@ -62,6 +62,7 @@ impl WeightedIndex {
         let uniform_index = usize::from_le_bytes(usize_seed) % self.aliases.len();
         let uniform_weight = Balance::from_le_bytes(balance_seed) % self.weight_sum;
 
+        tracing::info!(aliases=?self.aliases, "sample wac");
         if uniform_weight < self.no_alias_odds[uniform_index] {
             uniform_index
         } else {
@@ -171,6 +172,26 @@ mod test {
         assert_relative_closeness(counts[1], counts[2]);
     }
 
+    #[test]
+    fn test_sample_wac() {
+        let weights = vec![50e30 as u128; 4];
+        let weighted_index = WeightedIndex::new(weights);
+
+        let n_samples = 1_000_000;
+        let mut seed = hash(&[0; 32]);
+        let mut counts: [i32; 4] = [0, 0, 0, 0];
+        for _ in 0..n_samples {
+            let index = weighted_index.sample(seed);
+            counts[index] += 1;
+            seed = hash(&seed);
+        }
+
+        println!("counts: {counts:#?}");
+
+        assert_relative_closeness(counts[0], 5 * counts[1]);
+        assert_relative_closeness(counts[1], counts[2]);
+    }
+
     /// Assert y is within 0.5% of x.
     #[track_caller]
     fn assert_relative_closeness(x: i32, y: i32) {
diff --git a/core/store/src/trie/state_parts.rs b/core/store/src/trie/state_parts.rs
index 2eddd41c2fc..56a2c481a22 100644
--- a/core/store/src/trie/state_parts.rs
+++ b/core/store/src/trie/state_parts.rs
@@ -314,7 +314,7 @@ impl Trie {
     ///
     /// Creating a StatePart takes all these nodes, validating a StatePart checks that it has the
     /// right set of nodes.
-    fn visit_nodes_for_state_part(&self, part_id: PartId) -> Result<(), StorageError> {
+    pub fn visit_nodes_for_state_part(&self, part_id: PartId) -> Result<(), StorageError> {
         let path_begin = self.find_state_part_boundary(part_id.idx, part_id.total)?;
         let path_end = self.find_state_part_boundary(part_id.idx + 1, part_id.total)?;
 
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index 0037bd9cb1a..a3af4aca162 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -3,6 +3,7 @@
 use borsh::BorshSerialize;
 use near_client::{Client, ProcessTxResponse};
 use near_primitives::epoch_manager::{AllEpochConfig, EpochConfig};
+use near_primitives_core::num_rational::Rational32;
 
 use crate::tests::client::process_blocks::set_block_protocol_version;
 use assert_matches::assert_matches;
@@ -120,6 +121,8 @@ impl TestShardUpgradeEnv {
             gas_limit,
             genesis_protocol_version,
         );
+        tracing::info!(genesis_validators=?genesis.config.validators, "genesis");
+        // tracing::info!(genesis_contents=?genesis.contents, "genesis");
         let chain_genesis = ChainGenesis::new(&genesis);
         let env = TestEnv::builder(chain_genesis)
             .clients_count(num_clients)
@@ -178,7 +181,7 @@ impl TestShardUpgradeEnv {
         let expected_num_shards =
             get_expected_shards_num(self.epoch_length, height, resharding_type);
 
-        tracing::debug!(target: "test", height, expected_num_shards, "step");
+        // tracing::debug!(target: "test", height, expected_num_shards, "step");
 
         // add transactions for the next block
         if height == 1 {
@@ -209,9 +212,21 @@ impl TestShardUpgradeEnv {
             let block_producer_client = env.client(&block_producer);
             let mut block = block_producer_client.produce_block(height).unwrap().unwrap();
             set_block_protocol_version(&mut block, block_producer.clone(), protocol_version);
-            // std::thread::sleep(std::time::Duration::from_millis(2000));
+
             block
         };
+
+        {
+            let header = block.header();
+
+            let client = &env.clients[0];
+            let epoch_manager = &client.epoch_manager;
+            let epoch_id = epoch_manager.get_epoch_id_from_prev_block(header.prev_hash()).unwrap();
+            let epoch_info = &epoch_manager.get_epoch_info(&epoch_id).unwrap();
+            let epoch_protocol_version = epoch_info.protocol_version();
+
+            tracing::info!(latest_protocol_version=protocol_version, epoch_protocol_version, ?height, ?epoch_id, block = ?header.hash(), "step");
+        }
         // Make sure that catchup is done before the end of each epoch, but when it is done is
         // by chance. This simulates when catchup takes a long time to be done
         // Note: if the catchup happens only at the last block of an epoch then
@@ -251,7 +266,8 @@ impl TestShardUpgradeEnv {
                 .get_shard_layout_from_prev_block(block.hash())
                 .unwrap()
                 .num_shards();
-            assert_eq!(num_shards, expected_num_shards);
+            tracing::info!(?num_shards, ?expected_num_shards, "checking num shards");
+            // assert_eq!(num_shards, expected_num_shards);
         }
 
         env.process_partial_encoded_chunks();
@@ -634,6 +650,15 @@ fn setup_genesis(
         genesis.config.gas_limit = gas_limit;
     }
 
+    // The block producer assignment often happens to be unlucky enough to not
+    // include one of the validators in the first epoch. When that happens the
+    // new protocol version gets only 75% of the votes which is lower that the
+    // default 80% threshold for upgrading. Delaying the upgrade also delays
+    // resharding and makes it harder to predict. The threshold is set slightly
+    // lower here to always ensure that upgrade takes place as soon as possible.
+    // This was not fun to debug.
+    genesis.config.protocol_upgrade_stake_threshold = Rational32::new(7, 10);
+
     let default_epoch_config = EpochConfig::from(&genesis.config);
     let all_epoch_config = AllEpochConfig::new(true, default_epoch_config);
     let epoch_config = all_epoch_config.for_protocol_version(genesis_protocol_version);
@@ -644,6 +669,15 @@ fn setup_genesis(
     genesis.config.avg_hidden_validator_seats_per_shard =
         epoch_config.avg_hidden_validator_seats_per_shard;
 
+    tracing::info!(
+        num_block_producer_seats = ?genesis.config.num_block_producer_seats,
+        num_block_producer_seats_per_shard = ?genesis.config.num_block_producer_seats_per_shard,
+        num_chunk_only_producer_seats = ?genesis.config.num_chunk_only_producer_seats,
+        avg_hidden_validator_seats_per_shard = ?genesis.config.avg_hidden_validator_seats_per_shard,
+
+        "genesis"
+    );
+
     // TODO(resharding)
     // In the current simple test setup each validator may track different
     // shards. Unfortunately something is broken in how state sync is triggered

From 2809487024a55f608d8af00eefcd5f9acc7e9c15 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 9 Aug 2023 10:52:48 +0000
Subject: [PATCH 10/15] cleanup

---
 chain/epoch-manager/src/lib.rs                | 37 +------------------
 .../epoch-manager/src/validator_selection.rs  |  9 -----
 core/primitives/src/epoch_manager.rs          |  3 --
 core/primitives/src/rand.rs                   | 21 -----------
 core/store/src/trie/state_parts.rs            |  2 +-
 .../src/tests/client/sharding_upgrade.rs      | 35 +-----------------
 6 files changed, 4 insertions(+), 103 deletions(-)

diff --git a/chain/epoch-manager/src/lib.rs b/chain/epoch-manager/src/lib.rs
index cacc802e8c2..48205d3584d 100644
--- a/chain/epoch-manager/src/lib.rs
+++ b/chain/epoch-manager/src/lib.rs
@@ -46,8 +46,6 @@ mod tests;
 pub mod types;
 mod validator_selection;
 
-use itertools::Itertools;
-
 const EPOCH_CACHE_SIZE: usize = if cfg!(feature = "no_cache") { 1 } else { 50 };
 const BLOCK_CACHE_SIZE: usize = if cfg!(feature = "no_cache") { 5 } else { 1000 }; // TODO(#5080): fix this
 const AGGREGATOR_SAVE_PERIOD: u64 = 1000;
@@ -217,10 +215,6 @@ impl EpochManager {
                 genesis_protocol_version,
                 genesis_protocol_version,
             )?;
-
-            tracing::info!(block_producers_settlement=?epoch_info.block_producers_settlement(), "epoch manager new");
-            tracing::info!(chunk_producers_settlement=?epoch_info.chunk_producers_settlement(), "epoch manager new");
-            tracing::info!(validators=?epoch_info.validators_iter().collect_vec(), "epoch manager new");
             // Dummy block info.
             // Artificial block we add to simplify implementation: dummy block is the
             // parent of genesis block that points to itself.
@@ -538,12 +532,6 @@ impl EpochManager {
         for (validator_id, version) in version_tracker {
             let stake = epoch_info.validator_stake(validator_id);
             *versions.entry(version).or_insert(0) += stake;
-            tracing::info!(
-                validator_id,
-                version,
-                stake,
-                "voting on the next next epoch protocol version"
-            );
         }
         let total_block_producer_stake: u128 = epoch_info
             .block_producers_settlement()
@@ -565,35 +553,19 @@ impl EpochManager {
         // Note: non-deterministic iteration is fine here, there can be only one
         // version with large enough stake.
         let next_version = if let Some((version, stake)) =
-            versions.clone().into_iter().max_by_key(|&(_version, stake)| stake)
+            versions.into_iter().max_by_key(|&(_version, stake)| stake)
         {
             let numer = *config.protocol_upgrade_stake_threshold.numer() as u128;
             let denom = *config.protocol_upgrade_stake_threshold.denom() as u128;
             let threshold = total_block_producer_stake * numer / denom;
-            tracing::info!(
-                stake,
-                threshold,
-                numer,
-                denom,
-                total_block_producer_stake,
-                "voting on the next next epoch protocol version"
-            );
             if stake > threshold {
                 version
             } else {
                 protocol_version
             }
         } else {
-            tracing::info!("voting on the next next epoch protocol version winner - no max");
-
             protocol_version
         };
-        tracing::info!(
-            ?next_version,
-            ?versions,
-            "voting on the next next epoch protocol version winner"
-        );
-
         // Gather slashed validators and add them to kick out first.
         let slashed_validators = last_block_info.slashed();
         for (account_id, _) in slashed_validators.iter() {
@@ -658,8 +630,6 @@ impl EpochManager {
         let next_epoch_info = self.get_epoch_info(&next_epoch_id)?;
         self.save_epoch_validator_info(store_update, block_info.epoch_id(), &epoch_summary)?;
 
-        tracing::info!(?validator_stake, "finalize epoch");
-
         let EpochSummary {
             all_proposals,
             validator_kickout,
@@ -1773,8 +1743,6 @@ impl EpochManager {
         let mut aggregator = EpochInfoAggregator::new(epoch_id.clone(), *block_hash);
         let mut cur_hash = *block_hash;
 
-        tracing::info!(?block_hash, "aggregate_epoch_info_upto");
-
         Ok(Some(loop {
             #[cfg(test)]
             {
@@ -1796,8 +1764,6 @@ impl EpochManager {
                 // belongs to different epoch or we’re on different fork (though
                 // the latter should never happen).  In either case, the
                 // aggregator contains full epoch information.
-                tracing::info!(?block_hash, version_tracker=?aggregator.version_tracker, "aggregate_epoch_info_upto break epoch or genesis");
-
                 break (aggregator, true);
             }
 
@@ -1812,7 +1778,6 @@ impl EpochManager {
                 // We’ve reached sync point of the old aggregator.  If old
                 // aggregator was for a different epoch, we have full info in
                 // our aggregator; otherwise we don’t.
-                tracing::info!(?block_hash, version_tracker=?aggregator.version_tracker, "aggregate_epoch_info_upto break last block hash");
                 break (aggregator, epoch_id != prev_epoch);
             }
 
diff --git a/chain/epoch-manager/src/validator_selection.rs b/chain/epoch-manager/src/validator_selection.rs
index fbb9945a74d..41499687e22 100644
--- a/chain/epoch-manager/src/validator_selection.rs
+++ b/chain/epoch-manager/src/validator_selection.rs
@@ -53,7 +53,6 @@ pub fn proposals_to_epoch_info(
         min_stake_ratio,
         last_version,
     );
-    tracing::info!(?block_producers, "select_block_producers");
     let (chunk_producer_proposals, chunk_producers, cp_stake_threshold) =
         if checked_feature!("stable", ChunkOnlyProducers, next_version) {
             let mut chunk_producer_proposals = order_proposals(proposals.into_values());
@@ -108,13 +107,6 @@ pub fn proposals_to_epoch_info(
         all_validators.push(bp);
     }
 
-    tracing::info!(
-        ?validator_to_index,
-        ?block_producers_settlement,
-        ?all_validators,
-        "select_block_producers"
-    );
-
     let chunk_producers_settlement = if checked_feature!("stable", ChunkOnlyProducers, next_version)
     {
         let minimum_validators_per_shard =
@@ -126,7 +118,6 @@ pub fn proposals_to_epoch_info(
                     num_shards,
                 },
             )?;
-        tracing::info!(?shard_assignment, "select_block_producers");
 
         let mut chunk_producers_settlement: Vec<Vec<ValidatorId>> =
             shard_assignment.iter().map(|vs| Vec::with_capacity(vs.len())).collect();
diff --git a/core/primitives/src/epoch_manager.rs b/core/primitives/src/epoch_manager.rs
index a5f43cd8e17..f12672f7cd6 100644
--- a/core/primitives/src/epoch_manager.rs
+++ b/core/primitives/src/epoch_manager.rs
@@ -878,17 +878,14 @@ pub mod epoch_info {
         pub fn sample_block_producer(&self, height: BlockHeight) -> ValidatorId {
             match &self {
                 Self::V1(v1) => {
-                    tracing::info!("sample_block_producer v1");
                     let bp_settlement = &v1.block_producers_settlement;
                     bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
                 }
                 Self::V2(v2) => {
-                    tracing::info!("sample_block_producer v2");
                     let bp_settlement = &v2.block_producers_settlement;
                     bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
                 }
                 Self::V3(v3) => {
-                    tracing::info!(block_producers_settlement=?v3.block_producers_settlement, "sample_block_producer v3");
                     let seed = Self::block_produce_seed(height, &v3.rng_seed);
                     v3.block_producers_settlement[v3.block_producers_sampler.sample(seed)]
                 }
diff --git a/core/primitives/src/rand.rs b/core/primitives/src/rand.rs
index 735ff2ef887..a79c8fd1b56 100644
--- a/core/primitives/src/rand.rs
+++ b/core/primitives/src/rand.rs
@@ -62,7 +62,6 @@ impl WeightedIndex {
         let uniform_index = usize::from_le_bytes(usize_seed) % self.aliases.len();
         let uniform_weight = Balance::from_le_bytes(balance_seed) % self.weight_sum;
 
-        tracing::info!(aliases=?self.aliases, "sample wac");
         if uniform_weight < self.no_alias_odds[uniform_index] {
             uniform_index
         } else {
@@ -172,26 +171,6 @@ mod test {
         assert_relative_closeness(counts[1], counts[2]);
     }
 
-    #[test]
-    fn test_sample_wac() {
-        let weights = vec![50e30 as u128; 4];
-        let weighted_index = WeightedIndex::new(weights);
-
-        let n_samples = 1_000_000;
-        let mut seed = hash(&[0; 32]);
-        let mut counts: [i32; 4] = [0, 0, 0, 0];
-        for _ in 0..n_samples {
-            let index = weighted_index.sample(seed);
-            counts[index] += 1;
-            seed = hash(&seed);
-        }
-
-        println!("counts: {counts:#?}");
-
-        assert_relative_closeness(counts[0], 5 * counts[1]);
-        assert_relative_closeness(counts[1], counts[2]);
-    }
-
     /// Assert y is within 0.5% of x.
     #[track_caller]
     fn assert_relative_closeness(x: i32, y: i32) {
diff --git a/core/store/src/trie/state_parts.rs b/core/store/src/trie/state_parts.rs
index 56a2c481a22..2eddd41c2fc 100644
--- a/core/store/src/trie/state_parts.rs
+++ b/core/store/src/trie/state_parts.rs
@@ -314,7 +314,7 @@ impl Trie {
     ///
     /// Creating a StatePart takes all these nodes, validating a StatePart checks that it has the
     /// right set of nodes.
-    pub fn visit_nodes_for_state_part(&self, part_id: PartId) -> Result<(), StorageError> {
+    fn visit_nodes_for_state_part(&self, part_id: PartId) -> Result<(), StorageError> {
         let path_begin = self.find_state_part_boundary(part_id.idx, part_id.total)?;
         let path_end = self.find_state_part_boundary(part_id.idx + 1, part_id.total)?;
 
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index a3af4aca162..761ebb8c0f1 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -121,8 +121,6 @@ impl TestShardUpgradeEnv {
             gas_limit,
             genesis_protocol_version,
         );
-        tracing::info!(genesis_validators=?genesis.config.validators, "genesis");
-        // tracing::info!(genesis_contents=?genesis.contents, "genesis");
         let chain_genesis = ChainGenesis::new(&genesis);
         let env = TestEnv::builder(chain_genesis)
             .clients_count(num_clients)
@@ -181,7 +179,7 @@ impl TestShardUpgradeEnv {
         let expected_num_shards =
             get_expected_shards_num(self.epoch_length, height, resharding_type);
 
-        // tracing::debug!(target: "test", height, expected_num_shards, "step");
+        tracing::debug!(target: "test", height, expected_num_shards, "step");
 
         // add transactions for the next block
         if height == 1 {
@@ -216,17 +214,6 @@ impl TestShardUpgradeEnv {
             block
         };
 
-        {
-            let header = block.header();
-
-            let client = &env.clients[0];
-            let epoch_manager = &client.epoch_manager;
-            let epoch_id = epoch_manager.get_epoch_id_from_prev_block(header.prev_hash()).unwrap();
-            let epoch_info = &epoch_manager.get_epoch_info(&epoch_id).unwrap();
-            let epoch_protocol_version = epoch_info.protocol_version();
-
-            tracing::info!(latest_protocol_version=protocol_version, epoch_protocol_version, ?height, ?epoch_id, block = ?header.hash(), "step");
-        }
         // Make sure that catchup is done before the end of each epoch, but when it is done is
         // by chance. This simulates when catchup takes a long time to be done
         // Note: if the catchup happens only at the last block of an epoch then
@@ -267,7 +254,7 @@ impl TestShardUpgradeEnv {
                 .unwrap()
                 .num_shards();
             tracing::info!(?num_shards, ?expected_num_shards, "checking num shards");
-            // assert_eq!(num_shards, expected_num_shards);
+            assert_eq!(num_shards, expected_num_shards);
         }
 
         env.process_partial_encoded_chunks();
@@ -669,24 +656,6 @@ fn setup_genesis(
     genesis.config.avg_hidden_validator_seats_per_shard =
         epoch_config.avg_hidden_validator_seats_per_shard;
 
-    tracing::info!(
-        num_block_producer_seats = ?genesis.config.num_block_producer_seats,
-        num_block_producer_seats_per_shard = ?genesis.config.num_block_producer_seats_per_shard,
-        num_chunk_only_producer_seats = ?genesis.config.num_chunk_only_producer_seats,
-        avg_hidden_validator_seats_per_shard = ?genesis.config.avg_hidden_validator_seats_per_shard,
-
-        "genesis"
-    );
-
-    // TODO(resharding)
-    // In the current simple test setup each validator may track different
-    // shards. Unfortunately something is broken in how state sync is triggered
-    // from this test. The config below forces all validators to track all
-    // shards so that we can avoid that problem for now. It would be nice to
-    // set it up properly and test both state sync and resharding - once the
-    // integration is fully supported.
-    // genesis.config.minimum_validators_per_shard = num_validators;
-
     genesis
 }
 

From b93c48a3e77ac9a90e04ff39f2f292297fc9b1bf Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 9 Aug 2023 15:50:04 +0000
Subject: [PATCH 11/15] re-enable all transactions

---
 .../src/tests/client/sharding_upgrade.rs      | 177 +++++++++---------
 r.sh                                          |   2 +-
 2 files changed, 89 insertions(+), 90 deletions(-)

diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index 761ebb8c0f1..53336fb2f3c 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -868,7 +868,7 @@ fn setup_test_env_with_cross_contract_txs(
     genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
     let num = 4;
-    let test_env = TestShardUpgradeEnv::new_with_protocol_version(
+    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
         epoch_length,
         num,
         num,
@@ -876,94 +876,93 @@ fn setup_test_env_with_cross_contract_txs(
         Some(100_000_000_000_000),
         genesis_protocol_version,
     );
-    // let mut rng = thread_rng();
-
-    // let genesis_hash = *test_env.env.clients[0].chain.genesis_block().hash();
-    // // Use test0, test1 and two random accounts to deploy contracts because we want accounts on
-    // // different shards.
-    // let indices = (4..test_env.initial_accounts.len()).choose_multiple(&mut rng, 2);
-    // let contract_accounts = vec![
-    //     test_env.initial_accounts[0].clone(),
-    //     test_env.initial_accounts[1].clone(),
-    //     test_env.initial_accounts[indices[0]].clone(),
-    //     test_env.initial_accounts[indices[1]].clone(),
-    // ];
-    // test_env.set_init_tx(
-    //     contract_accounts
-    //         .iter()
-    //         .map(|account_id| {
-    //             let signer = InMemorySigner::from_seed(
-    //                 account_id.clone(),
-    //                 KeyType::ED25519,
-    //                 &account_id.to_string(),
-    //             );
-    //             SignedTransaction::from_actions(
-    //                 1,
-    //                 account_id.clone(),
-    //                 account_id.clone(),
-    //                 &signer,
-    //                 vec![Action::DeployContract(DeployContractAction {
-    //                     code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
-    //                 })],
-    //                 genesis_hash,
-    //             )
-    //         })
-    //         .collect(),
-    // );
-
-    let new_accounts = HashMap::new();
-    // let mut nonce = 100;
-    // let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
-    // let mut new_accounts = HashMap::new();
-    // let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
-    //     &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
-    //         let mut rng = thread_rng();
-    //         let size = rng.gen_range(min_size..max_size + 1);
-    //         std::iter::repeat_with(|| loop {
-    //             let account_id = gen_account(&mut rng, b"abcdefghijkmn");
-    //             if all_accounts.insert(account_id.clone()) {
-    //                 nonce += 1;
-    //                 // randomly shuffle contract accounts
-    //                 // so that transactions are send to different shards
-    //                 let mut contract_accounts = contract_accounts.clone();
-    //                 contract_accounts.shuffle(&mut rng);
-    //                 let tx = gen_cross_contract_transaction(
-    //                     &contract_accounts[0],
-    //                     &contract_accounts[1],
-    //                     &contract_accounts[2],
-    //                     &contract_accounts[3],
-    //                     &account_id,
-    //                     nonce,
-    //                     &genesis_hash,
-    //                 );
-    //                 new_accounts.insert(tx.get_hash(), account_id);
-    //                 return tx;
-    //             }
-    //         })
-    //         .take(size)
-    //         .collect()
-    //     };
-
-    // // add a bunch of transactions before the two epoch boundaries
-    // for height in vec![
-    //     epoch_length - 2,
-    //     epoch_length - 1,
-    //     epoch_length,
-    //     2 * epoch_length - 2,
-    //     2 * epoch_length - 1,
-    //     2 * epoch_length,
-    // ] {
-    //     test_env.set_tx_at_height(height, generate_txs(5, 8));
-    // }
-
-    // // adds some transactions after sharding change finishes
-    // // but do not add too many because I want all transactions to
-    // // finish processing before epoch 5
-    // for height in 2 * epoch_length + 1..3 * epoch_length {
-    //     if rng.gen_bool(0.3) {
-    //         test_env.set_tx_at_height(height, generate_txs(5, 8));
-    //     }
-    // }
+    let mut rng = thread_rng();
+
+    let genesis_hash = *test_env.env.clients[0].chain.genesis_block().hash();
+    // Use test0, test1 and two random accounts to deploy contracts because we want accounts on
+    // different shards.
+    let indices = (4..test_env.initial_accounts.len()).choose_multiple(&mut rng, 2);
+    let contract_accounts = vec![
+        test_env.initial_accounts[0].clone(),
+        test_env.initial_accounts[1].clone(),
+        test_env.initial_accounts[indices[0]].clone(),
+        test_env.initial_accounts[indices[1]].clone(),
+    ];
+    test_env.set_init_tx(
+        contract_accounts
+            .iter()
+            .map(|account_id| {
+                let signer = InMemorySigner::from_seed(
+                    account_id.clone(),
+                    KeyType::ED25519,
+                    &account_id.to_string(),
+                );
+                SignedTransaction::from_actions(
+                    1,
+                    account_id.clone(),
+                    account_id.clone(),
+                    &signer,
+                    vec![Action::DeployContract(DeployContractAction {
+                        code: near_test_contracts::backwards_compatible_rs_contract().to_vec(),
+                    })],
+                    genesis_hash,
+                )
+            })
+            .collect(),
+    );
+
+    let mut nonce = 100;
+    let mut all_accounts: HashSet<_> = test_env.initial_accounts.clone().into_iter().collect();
+    let mut new_accounts = HashMap::new();
+    let generate_txs: &mut dyn FnMut(usize, usize) -> Vec<SignedTransaction> =
+        &mut |min_size: usize, max_size: usize| -> Vec<SignedTransaction> {
+            let mut rng = thread_rng();
+            let size = rng.gen_range(min_size..max_size + 1);
+            std::iter::repeat_with(|| loop {
+                let account_id = gen_account(&mut rng, b"abcdefghijkmn");
+                if all_accounts.insert(account_id.clone()) {
+                    nonce += 1;
+                    // randomly shuffle contract accounts
+                    // so that transactions are send to different shards
+                    let mut contract_accounts = contract_accounts.clone();
+                    contract_accounts.shuffle(&mut rng);
+                    let tx = gen_cross_contract_transaction(
+                        &contract_accounts[0],
+                        &contract_accounts[1],
+                        &contract_accounts[2],
+                        &contract_accounts[3],
+                        &account_id,
+                        nonce,
+                        &genesis_hash,
+                    );
+                    new_accounts.insert(tx.get_hash(), account_id);
+                    return tx;
+                }
+            })
+            .take(size)
+            .collect()
+        };
+
+    // add a bunch of transactions before the two epoch boundaries
+    for height in vec![
+        epoch_length - 2,
+        epoch_length - 1,
+        epoch_length,
+        2 * epoch_length - 2,
+        2 * epoch_length - 1,
+        2 * epoch_length,
+    ] {
+        test_env.set_tx_at_height(height, generate_txs(5, 8));
+    }
+
+    // adds some transactions after sharding change finishes
+    // but do not add too many because I want all transactions to
+    // finish processing before epoch 5
+    for height in 2 * epoch_length + 1..3 * epoch_length {
+        if rng.gen_bool(0.3) {
+            test_env.set_tx_at_height(height, generate_txs(5, 8));
+        }
+    }
 
     (test_env, new_accounts)
 }
diff --git a/r.sh b/r.sh
index 80133092509..66909ac54ad 100755
--- a/r.sh
+++ b/r.sh
@@ -18,8 +18,8 @@ for i in {1..1}
      RUST_BACKTRACE=all \
      RUST_LOG=debug,resharding=trace \
      cargo nextest run -p integration-tests \
-         --no-capture \
          --features nightly \
+         --no-capture \
 		 $TEST \
          | egrep -v prev_prev_stake_change \
          > $OUT

From 5a101eb45ef76ea6d169c27cf0d0dee5d5af88a2 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 9 Aug 2023 15:50:26 +0000
Subject: [PATCH 12/15] disable the v2 test

---
 integration-tests/src/tests/client/sharding_upgrade.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index 53336fb2f3c..b5697d52c41 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -1005,6 +1005,10 @@ fn test_shard_layout_upgrade_cross_contract_calls_v1() {
 // Test cross contract calls
 // This test case tests postponed receipts and delayed receipts
 #[cfg(feature = "protocol_feature_simple_nightshade_v2")]
+// TODO(resharding) this test is currently broken, potentially due to lack of
+// flat storage support. Once flat storage for resharding is fully implemented
+// this test should be revisited, fixed and re-enabled.
+#[ignore]
 #[test]
 fn test_shard_layout_upgrade_cross_contract_calls_v2() {
     test_shard_layout_upgrade_cross_contract_calls_impl(ReshardingType::V2);

From 53087c4b841008a1f4206d96e42744f3283ef23f Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Wed, 9 Aug 2023 16:50:47 +0000
Subject: [PATCH 13/15] cleanup

---
 .gitignore                                    |  4 +-
 chain/chain/src/chain.rs                      |  1 +
 chain/chunks/src/lib.rs                       |  7 +--
 chain/chunks/src/logic.rs                     |  4 +-
 chain/chunks/src/test_loop.rs                 |  2 +-
 chain/client/src/client.rs                    |  2 -
 chain/client/src/test_utils.rs                | 10 ++--
 chain/epoch-manager/src/lib.rs                |  1 -
 core/o11y/src/testonly.rs                     |  2 +-
 core/primitives-core/src/hash.rs              |  3 +-
 .../src/tests/client/sharding_upgrade.rs      | 52 +++++++------------
 p.sh                                          | 11 ----
 r.sh                                          | 38 --------------
 13 files changed, 32 insertions(+), 105 deletions(-)
 delete mode 100755 p.sh
 delete mode 100755 r.sh

diff --git a/.gitignore b/.gitignore
index edd9464ff9a..c5a545232a6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -61,6 +61,4 @@ rusty-tags.vi
 # Estimator generated files
 costs-*.txt
 names-to-stats.txt
-data_dump_*.bin
-
-stdout-1
+data_dump_*.bin
\ No newline at end of file
diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index 63bf88ac236..370574c367e 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -1120,6 +1120,7 @@ impl Chain {
             &prev_hash,
         )?;
         let prev_block = self.get_block(&prev_hash)?;
+
         if prev_block.chunks().len() != block.chunks().len() && !shards_to_state_sync.is_empty() {
             // Currently, the state sync algorithm assumes that the number of chunks do not change
             // between the epoch being synced to and the last epoch.
diff --git a/chain/chunks/src/lib.rs b/chain/chunks/src/lib.rs
index 3481a5078c8..864e5d0c0e0 100644
--- a/chain/chunks/src/lib.rs
+++ b/chain/chunks/src/lib.rs
@@ -137,7 +137,6 @@ pub mod metrics;
 pub mod shards_manager_actor;
 pub mod test_loop;
 pub mod test_utils;
-use itertools::Itertools;
 
 pub const CHUNK_REQUEST_RETRY: time::Duration = time::Duration::milliseconds(100);
 pub const CHUNK_REQUEST_SWITCH_TO_OTHERS: time::Duration = time::Duration::milliseconds(400);
@@ -704,7 +703,6 @@ impl ShardsManager {
         // Process chunk one part requests.
         let requests = self.requested_partial_encoded_chunks.fetch(self.clock.now().into());
         for (chunk_hash, chunk_request) in requests {
-            tracing::trace!(target: "client", ?chunk_request, "resending chunk request");
             let fetch_from_archival =
                 chunk_needs_to_be_fetched_from_archival(&chunk_request.ancestor_hash, &self.chain_header_head.last_block_hash,
                 self.epoch_manager.as_ref()).unwrap_or_else(|err| {
@@ -1172,7 +1170,6 @@ impl ShardsManager {
         &mut self,
         forward: PartialEncodedChunkForwardMsg,
     ) -> Result<(), Error> {
-        tracing::debug!(target: "chunks", shard_id=forward.shard_id, parts_len=forward.parts.len(), "process_partial_encoded_chunk_forward");
         let maybe_header = self
             .validate_partial_encoded_chunk_forward(&forward)
             .and_then(|_| self.get_partial_encoded_chunk_header(&forward.chunk_hash));
@@ -1361,7 +1358,6 @@ impl ShardsManager {
         let chunk_hash = header.chunk_hash();
         debug!(
             target: "chunks",
-            me = ?self.me,
             ?chunk_hash,
             height = header.height_created(),
             shard_id = header.shard_id(),
@@ -1689,7 +1685,6 @@ impl ShardsManager {
         epoch_id: &EpochId,
         lastest_block_hash: &CryptoHash,
     ) -> Result<(), Error> {
-        tracing::debug!(target: "client", me=?self.me, parts=?partial_encoded_chunk.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "send_partial_encoded_chunk_to_chunk_trackers");
         let me = match self.me.as_ref() {
             Some(me) => me,
             None => return Ok(()),
@@ -1706,7 +1701,6 @@ impl ShardsManager {
             })
             .cloned()
             .collect();
-        tracing::debug!(target: "client", ?me, owned_parts_len=owned_parts.len(), "send_partial_encoded_chunk_to_chunk_trackers");
 
         if owned_parts.is_empty() {
             return Ok(());
@@ -1734,6 +1728,7 @@ impl ShardsManager {
                 continue;
             }
             next_chunk_producers.remove(&bp_account_id);
+
             // Technically, here we should check if the block producer actually cares about the shard.
             // We don't because with the current implementation, we force all validators to track all
             // shards by making their config tracking all shards.
diff --git a/chain/chunks/src/logic.rs b/chain/chunks/src/logic.rs
index 394fe36d79a..d67d2fbb770 100644
--- a/chain/chunks/src/logic.rs
+++ b/chain/chunks/src/logic.rs
@@ -12,7 +12,7 @@ use near_primitives::{
     },
     types::{AccountId, ShardId},
 };
-use tracing::log::{debug, error};
+use tracing::log::error;
 
 pub fn need_receipt(
     prev_block_hash: &CryptoHash,
@@ -154,8 +154,6 @@ pub fn decode_encoded_chunk(
             Ok(shard_chunk)
         })
     {
-        debug!(target: "chunks", "Reconstructed and decoded chunk {}, shard id {}, encoded length was {}, num txs: {}, I'm {:?}", chunk_hash.0, encoded_chunk.shard_id(), encoded_chunk.encoded_length(), shard_chunk.transactions().len(), me);
-
         let partial_chunk = create_partial_chunk(
             encoded_chunk,
             merkle_paths,
diff --git a/chain/chunks/src/test_loop.rs b/chain/chunks/src/test_loop.rs
index 12b38979faa..32ab8f7ed10 100644
--- a/chain/chunks/src/test_loop.rs
+++ b/chain/chunks/src/test_loop.rs
@@ -347,7 +347,7 @@ impl TestChunkEncoder {
     }
 
     pub fn part_ords(&self) -> Vec<u64> {
-        self.full_partial_chunk.parts().iter().take(3).map(|part| part.part_ord).collect()
+        self.full_partial_chunk.parts().iter().map(|part| part.part_ord).collect()
     }
 
     pub fn make_partial_encoded_chunk(
diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs
index 594b6d9417b..5545ca92543 100644
--- a/chain/client/src/client.rs
+++ b/chain/client/src/client.rs
@@ -607,8 +607,6 @@ impl Client {
             panic!("The client protocol version is older than the protocol version of the network. Please update nearcore. Client protocol version:{}, network protocol version {}", PROTOCOL_VERSION, protocol_version);
         }
 
-        tracing::debug!(target: "client", ?next_height, ?epoch_id, ?protocol_version, "producing block");
-
         let approvals = self
             .epoch_manager
             .get_epoch_block_approvers_ordered(&prev_hash)?
diff --git a/chain/client/src/test_utils.rs b/chain/client/src/test_utils.rs
index f395270698c..4f6cdc7b8ec 100644
--- a/chain/client/src/test_utils.rs
+++ b/chain/client/src/test_utils.rs
@@ -9,7 +9,6 @@ use actix::{Actor, Addr, AsyncContext, Context};
 use actix_rt::{Arbiter, System};
 use chrono::DateTime;
 use futures::{future, FutureExt};
-use itertools::Itertools;
 use near_async::actix::AddrWithAutoSpanContextExt;
 use near_async::messaging::{CanSend, IntoSender, LateBoundSender, Sender};
 use near_async::time;
@@ -833,7 +832,6 @@ pub fn setup_mock_all_validators(
                             );
                         }
                         NetworkRequests::PartialEncodedChunkForward { account_id, forward } => {
-                            tracing::debug!(target: "test", "just checking, is this called?");
                             send_chunks(
                                 connectors1,
                                 validators_clone2.iter().cloned().enumerate(),
@@ -1879,7 +1877,7 @@ impl TestEnv {
 
         let mut keep_going = true;
         while keep_going {
-            for (i, network_adapter) in network_adapters.iter().enumerate() {
+            for network_adapter in network_adapters.iter() {
                 keep_going = false;
                 // process partial encoded chunks
                 while let Some(request) = network_adapter.pop() {
@@ -1888,6 +1886,7 @@ impl TestEnv {
                     // trigger more messages to be processed in other clients
                     // it's a bit sad and it would be much nicer if all messages
                     // were forwarded to a single queue
+                    // TODO would be nicer to first handle all PECs and then all PECFs
                     keep_going = true;
                     match request {
                         PeerManagerMessageRequest::NetworkRequests(
@@ -1896,7 +1895,6 @@ impl TestEnv {
                                 partial_encoded_chunk,
                             },
                         ) => {
-                            tracing::debug!(target: "test", client=i, ?account_id, shard_id=partial_encoded_chunk.header.shard_id(), parts=?partial_encoded_chunk.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "handling partial encoded chunk");
                             let partial_encoded_chunk =
                                 PartialEncodedChunk::from(partial_encoded_chunk);
                             let message =
@@ -1908,7 +1906,6 @@ impl TestEnv {
                         PeerManagerMessageRequest::NetworkRequests(
                             NetworkRequests::PartialEncodedChunkForward { account_id, forward },
                         ) => {
-                            tracing::debug!(target: "test", client=i, ?account_id, shard_id=forward.shard_id, parts=?forward.parts.iter().take(3).map(|part| part.part_ord).collect_vec(), "handling forward");
                             let message =
                                 ShardsManagerRequestFromNetwork::ProcessPartialEncodedChunkForward(
                                     forward,
@@ -2013,6 +2010,9 @@ impl TestEnv {
     }
 
     pub fn process_shards_manager_responses_and_finish_processing_blocks(&mut self, idx: usize) {
+        let _span =
+            tracing::debug_span!(target: "test", "process_shards_manager", client=idx).entered();
+
         loop {
             self.process_shards_manager_responses(idx);
             if self.clients[idx].finish_blocks_in_processing().is_empty() {
diff --git a/chain/epoch-manager/src/lib.rs b/chain/epoch-manager/src/lib.rs
index 48205d3584d..710933d84b5 100644
--- a/chain/epoch-manager/src/lib.rs
+++ b/chain/epoch-manager/src/lib.rs
@@ -1742,7 +1742,6 @@ impl EpochManager {
 
         let mut aggregator = EpochInfoAggregator::new(epoch_id.clone(), *block_hash);
         let mut cur_hash = *block_hash;
-
         Ok(Some(loop {
             #[cfg(test)]
             {
diff --git a/core/o11y/src/testonly.rs b/core/o11y/src/testonly.rs
index 2fff88ca204..2dddbc9f86c 100644
--- a/core/o11y/src/testonly.rs
+++ b/core/o11y/src/testonly.rs
@@ -27,7 +27,7 @@ fn setup_subscriber_from_filter(mut env_filter: EnvFilter) {
 
     let _ = fmt::Subscriber::builder()
         .with_ansi(use_color_auto())
-        .with_span_events(fmt::format::FmtSpan::NONE)
+        .with_span_events(fmt::format::FmtSpan::CLOSE)
         .with_env_filter(env_filter)
         .with_writer(fmt::TestWriter::new())
         .with_timer(TestUptime::default())
diff --git a/core/primitives-core/src/hash.rs b/core/primitives-core/src/hash.rs
index 8e77e8724ee..7b703e9cd2d 100644
--- a/core/primitives-core/src/hash.rs
+++ b/core/primitives-core/src/hash.rs
@@ -207,8 +207,7 @@ impl fmt::Debug for CryptoHash {
 
 impl fmt::Display for CryptoHash {
     fn fmt(&self, fmtr: &mut fmt::Formatter<'_>) -> fmt::Result {
-        // TODO remove me debugging only
-        self.to_base58_impl(|encoded| fmtr.write_str(&encoded[..4]))
+        self.to_base58_impl(|encoded| fmtr.write_str(encoded))
     }
 }
 
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index b5697d52c41..ef82db7d0e7 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -1,5 +1,3 @@
-#![allow(unused)]
-
 use borsh::BorshSerialize;
 use near_client::{Client, ProcessTxResponse};
 use near_primitives::epoch_manager::{AllEpochConfig, EpochConfig};
@@ -49,7 +47,7 @@ const SIMPLE_NIGHTSHADE_V2_PROTOCOL_VERSION: ProtocolVersion =
 #[cfg(not(feature = "protocol_feature_simple_nightshade_v2"))]
 const SIMPLE_NIGHTSHADE_V2_PROTOCOL_VERSION: ProtocolVersion = PROTOCOL_VERSION + 1;
 
-// const P_CATCHUP: f64 = 0.2;
+const P_CATCHUP: f64 = 0.2;
 
 enum ReshardingType {
     // In the V0->V1 resharding outgoing receipts are reassigned to receiver.
@@ -101,7 +99,7 @@ struct TestShardUpgradeEnv {
 }
 
 impl TestShardUpgradeEnv {
-    fn new_with_protocol_version(
+    fn new(
         epoch_length: u64,
         num_validators: usize,
         num_clients: usize,
@@ -218,8 +216,7 @@ impl TestShardUpgradeEnv {
         // by chance. This simulates when catchup takes a long time to be done
         // Note: if the catchup happens only at the last block of an epoch then
         // client will fail to produce the chunks in the first block of the next epoch.
-        // let should_catchup = rng.gen_bool(P_CATCHUP) || height % self.epoch_length == 0;
-        let should_catchup = (height + 1) % self.epoch_length == 0;
+        let should_catchup = rng.gen_bool(P_CATCHUP) || height % self.epoch_length == 0;
         // process block, this also triggers chunk producers for the next block to produce chunks
         for j in 0..self.num_clients {
             let client = &mut env.clients[j];
@@ -228,13 +225,13 @@ impl TestShardUpgradeEnv {
             // Here we don't just call self.clients[i].process_block_sync_with_produce_chunk_options
             // because we want to call run_catchup before finish processing this block. This simulates
             // that catchup and block processing run in parallel.
-            let result = client.start_process_block(
-                MaybeValidated::from(block.clone()),
-                Provenance::NONE,
-                Arc::new(|_| {}),
-            );
-            tracing::debug!(target: "test", start_process_block_result=?result);
-            result.unwrap();
+            client
+                .start_process_block(
+                    MaybeValidated::from(block.clone()),
+                    Provenance::NONE,
+                    Arc::new(|_| {}),
+                )
+                .unwrap();
             if should_catchup {
                 run_catchup(client, &[]).unwrap();
             }
@@ -253,16 +250,11 @@ impl TestShardUpgradeEnv {
                 .get_shard_layout_from_prev_block(block.hash())
                 .unwrap()
                 .num_shards();
-            tracing::info!(?num_shards, ?expected_num_shards, "checking num shards");
             assert_eq!(num_shards, expected_num_shards);
         }
 
         env.process_partial_encoded_chunks();
         for j in 0..self.num_clients {
-            let _span =
-                tracing::debug_span!(target: "test", "process shard manager things", client=j);
-            let _span = _span.entered();
-
             env.process_shards_manager_responses_and_finish_processing_blocks(j);
         }
 
@@ -272,12 +264,16 @@ impl TestShardUpgradeEnv {
         }
     }
 
+    // Submit the tx to all clients for processing and checks:
+    // Clients that track the relevant shard should return ValidTx
+    // Clients that do not track the relevenat shard should return RequestRouted
+    // At least one client should process it and return ValidTx.
     fn process_tx(env: &mut TestEnv, tx: &SignedTransaction) {
         let mut response_valid_count = 0;
         let mut response_routed_count = 0;
         for j in 0..env.validators.len() {
             let response = env.clients[j].process_tx(tx.clone(), false, false);
-            tracing::debug!(target: "test", client=j, tx=?tx.get_hash(), ?response, "process tx");
+            tracing::trace!(target: "test", client=j, tx=?tx.get_hash(), ?response, "process tx");
             match response {
                 ProcessTxResponse::ValidTx => response_valid_count += 1,
                 ProcessTxResponse::RequestRouted => response_routed_count += 1,
@@ -670,14 +666,8 @@ fn test_shard_layout_upgrade_simple_impl(resharding_type: ReshardingType) {
 
     // setup
     let epoch_length = 5;
-    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
-        epoch_length,
-        2,
-        2,
-        100,
-        None,
-        genesis_protocol_version,
-    );
+    let mut test_env =
+        TestShardUpgradeEnv::new(epoch_length, 2, 2, 100, None, genesis_protocol_version);
     test_env.set_init_tx(vec![]);
 
     let mut nonce = 100;
@@ -867,11 +857,10 @@ fn setup_test_env_with_cross_contract_txs(
     epoch_length: u64,
     genesis_protocol_version: ProtocolVersion,
 ) -> (TestShardUpgradeEnv, HashMap<CryptoHash, AccountId>) {
-    let num = 4;
-    let mut test_env = TestShardUpgradeEnv::new_with_protocol_version(
+    let mut test_env = TestShardUpgradeEnv::new(
         epoch_length,
-        num,
-        num,
+        4,
+        4,
         100,
         Some(100_000_000_000_000),
         genesis_protocol_version,
@@ -981,7 +970,6 @@ fn test_shard_layout_upgrade_cross_contract_calls_impl(resharding_type: Reshardi
         setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 
     for _ in 1..5 * epoch_length {
-        // std::thread::sleep(std::time::Duration::from_millis(2000));
         test_env.step_impl(0., target_protocol_version, &resharding_type);
         test_env.check_receipt_id_to_shard_id();
     }
diff --git a/p.sh b/p.sh
deleted file mode 100755
index a2ace2b6609..00000000000
--- a/p.sh
+++ /dev/null
@@ -1,11 +0,0 @@
-
-clear
-tmux clear-history
-
-RUST_BACKTRACE=all \
-RUST_LOG=debug,resharding=trace \
-cargo nextest run \
-    --package integration-tests \
-    --features nightly \
-    sharding_upgrade
-
diff --git a/r.sh b/r.sh
deleted file mode 100755
index 66909ac54ad..00000000000
--- a/r.sh
+++ /dev/null
@@ -1,38 +0,0 @@
-
-clear
-tmux clear-history
-
-rm stdout-*
-
-TEST=test_shard_layout_upgrade_simple_v2
-TEST=test_shard_layout_upgrade_cross_contract_calls_v2
-
-for i in {1..1}
- do
-     echo "test $i"
-     OUT=stdout-$i
-
-     # RUST_LOG=info \
-     #RUST_LOG=info,catchup=trace,store=trace,client=debug,store=debug,test=debug,resharding=trace \
-     # RUST_LOG=debug,resharding=trace,waclaw=trace,catchup=trace \
-     RUST_BACKTRACE=all \
-     RUST_LOG=debug,resharding=trace \
-     cargo nextest run -p integration-tests \
-         --features nightly \
-         --no-capture \
-		 $TEST \
-         | egrep -v prev_prev_stake_change \
-         > $OUT
-         # | tee $OUT
-         # | egrep -v -i "FlatStorage is not ready|Add delta for flat storage creation|epoch_manager: all proposals" \
-
-     sed -E -i 's/ed25519:(.{4})(.{40})/ed25519:\1/g' $OUT
-
-     sed -E -i 's/([0-9]*)([0-9]{30})/\1e30/g' $OUT
-     sed -E -i 's/([0-9]*)([0-9]{25})/\1e25/g' $OUT
-     sed -E -i 's/([0-9]*)([0-9]{20})/\1e20/g' $OUT
-     sed -E -i 's/AccountId/AId/g' $OUT
-
-     cat $OUT | egrep -a -i error
-
- done

From bd6a273d9189b83b212871b357b51c56c5318fbb Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Thu, 10 Aug 2023 09:26:38 +0000
Subject: [PATCH 14/15] nits

---
 chain/chain/src/chain.rs                               | 1 -
 chain/chunks/src/logic.rs                              | 3 ++-
 integration-tests/src/tests/client/sharding_upgrade.rs | 3 +--
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
index 370574c367e..f25d5e9a170 100644
--- a/chain/chain/src/chain.rs
+++ b/chain/chain/src/chain.rs
@@ -2534,7 +2534,6 @@ impl Chain {
         self.validate_chunk_headers(&block, &prev_block)?;
 
         self.ping_missing_chunks(me, prev_hash, block)?;
-
         let incoming_receipts = self.collect_incoming_receipts_from_block(me, block)?;
 
         // Check if block can be finalized and drop it otherwise.
diff --git a/chain/chunks/src/logic.rs b/chain/chunks/src/logic.rs
index d67d2fbb770..b0a4f823170 100644
--- a/chain/chunks/src/logic.rs
+++ b/chain/chunks/src/logic.rs
@@ -12,7 +12,7 @@ use near_primitives::{
     },
     types::{AccountId, ShardId},
 };
-use tracing::log::error;
+use tracing::log::{debug, error};
 
 pub fn need_receipt(
     prev_block_hash: &CryptoHash,
@@ -154,6 +154,7 @@ pub fn decode_encoded_chunk(
             Ok(shard_chunk)
         })
     {
+        debug!(target: "chunks", "Reconstructed and decoded chunk {}, encoded length was {}, num txs: {}, I'm {:?}", chunk_hash.0, encoded_chunk.encoded_length(), shard_chunk.transactions().len(), me);
         let partial_chunk = create_partial_chunk(
             encoded_chunk,
             merkle_paths,
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index ef82db7d0e7..6c6e3f19579 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -203,8 +203,7 @@ impl TestShardUpgradeEnv {
                 client.epoch_manager.get_epoch_id_from_prev_block(&head.last_block_hash).unwrap();
             let block_producer =
                 client.epoch_manager.get_block_producer(&epoch_id, height).unwrap();
-            let _span = tracing::debug_span!(target: "test", "", client=?block_producer);
-            let _span = _span.entered();
+            let _span = tracing::debug_span!(target: "test", "", client=?block_producer).entered();
             let block_producer_client = env.client(&block_producer);
             let mut block = block_producer_client.produce_block(height).unwrap().unwrap();
             set_block_protocol_version(&mut block, block_producer.clone(), protocol_version);

From 8a91db0167284f2703d41f3266f4317da4377450 Mon Sep 17 00:00:00 2001
From: wacban <wac.banasik@gmail.com>
Date: Thu, 17 Aug 2023 12:46:49 +0000
Subject: [PATCH 15/15] code review

---
 chain/chunks/Cargo.toml                                | 1 +
 chain/epoch-manager/Cargo.toml                         | 6 ++----
 integration-tests/src/tests/client/sharding_upgrade.rs | 6 +++---
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/chain/chunks/Cargo.toml b/chain/chunks/Cargo.toml
index e1322ed71a0..d2468844e45 100644
--- a/chain/chunks/Cargo.toml
+++ b/chain/chunks/Cargo.toml
@@ -21,6 +21,7 @@ rand.workspace = true
 reed-solomon-erasure.workspace = true
 time.workspace = true
 tracing.workspace = true
+# itertools has collect_vec which is useful in quick debugging prints
 itertools.workspace = true
 
 near-async.workspace = true
diff --git a/chain/epoch-manager/Cargo.toml b/chain/epoch-manager/Cargo.toml
index c318bbe0eb9..181abc8a44f 100644
--- a/chain/epoch-manager/Cargo.toml
+++ b/chain/epoch-manager/Cargo.toml
@@ -18,6 +18,7 @@ rand_hc.workspace = true
 serde_json.workspace = true
 smart-default.workspace = true
 tracing.workspace = true
+# itertools has collect_vec which is useful in quick debugging prints
 itertools.workspace = true
 
 near-crypto.workspace = true
@@ -46,7 +47,4 @@ nightly_protocol = [
   "near-store/nightly_protocol",
 ]
 no_cache = []
-new_epoch_sync = [
-  "near-store/new_epoch_sync",
-  "near-primitives/new_epoch_sync"
-]
+new_epoch_sync = ["near-store/new_epoch_sync", "near-primitives/new_epoch_sync"]
diff --git a/integration-tests/src/tests/client/sharding_upgrade.rs b/integration-tests/src/tests/client/sharding_upgrade.rs
index 6c6e3f19579..5a55933d26b 100644
--- a/integration-tests/src/tests/client/sharding_upgrade.rs
+++ b/integration-tests/src/tests/client/sharding_upgrade.rs
@@ -994,7 +994,7 @@ fn test_shard_layout_upgrade_cross_contract_calls_v1() {
 #[cfg(feature = "protocol_feature_simple_nightshade_v2")]
 // TODO(resharding) this test is currently broken, potentially due to lack of
 // flat storage support. Once flat storage for resharding is fully implemented
-// this test should be revisited, fixed and re-enabled.
+// this test should be revisited, fixed and re-enabled. See #8992 for overall progress.
 #[ignore]
 #[test]
 fn test_shard_layout_upgrade_cross_contract_calls_v2() {
@@ -1008,12 +1008,12 @@ fn test_shard_layout_upgrade_missing_chunks(p_missing: f64) {
     init_test_logger();
 
     let resharding_type = ReshardingType::V1;
-    let genesis_protocol_versino = get_genesis_protocol_version(&resharding_type);
+    let genesis_protocol_version = get_genesis_protocol_version(&resharding_type);
 
     // setup
     let epoch_length = 10;
     let (mut test_env, new_accounts) =
-        setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_versino);
+        setup_test_env_with_cross_contract_txs(epoch_length, genesis_protocol_version);
 
     // randomly dropping chunks at the first few epochs when sharding splits happens
     // make sure initial txs (deploy smart contracts) are processed succesfully