fix(resharding): storage costs (#12661)

Fixing old storage costs inaccuracy, which caused a failure during some
forknet experiment
https://near.zulipchat.com/#narrow/channel/407288-core.2Fresharding/topic/forknet/near/489699725

The fix is one-liner, `charge_gas_for_trie_node_access: false,`.
Originally it wasn't the case because we had a protocol upgrade from
trie to flat storage read costs, so we needed to compute costs
differently. But since flat storage costs were enabled, and we started
to use only the flat storage read costs. Moreover, later it became
controlled by runtime `Parameter::FlatStorageReads`, but the original
condition stayed.

And now, when we do flat storage resharding, flat storage indeed doesn't
exist for a while, which triggered trie costs for some blocks during
resharding again. However, on chunk validation, costs were correct, so
chain couldn't validate any chunks since resharding start.

I test this by calling a contract which reads a key and then writes
key-value pair back. For the old code, key read charges more cost than
it should, which causes InvalidOutcomesProof error.

---------

Co-authored-by: wacban <[email protected]>

chain/chain/src/runtime/mod.rs

@@ -635,7 +635,7 @@ impl RuntimeAdapter for NightshadeRuntime {
                 // flat storage by not charging gas for trie nodes.
                 // WARNING: should never be used in production! Consider this option only for debugging or replaying blocks.
                 let mut trie = self.tries.get_trie_for_shard(shard_uid, storage_config.state_root);
-                trie.dont_charge_gas_for_trie_node_access();
+                trie.set_charge_gas_for_trie_node_access(false);
                 trie
             }
             StorageDataSource::Recorded(storage) => Trie::from_recorded_storage(
@@ -862,7 +862,7 @@ impl RuntimeAdapter for NightshadeRuntime {
                     storage_config.state_root,
                     false,
                 )?;
-                trie.dont_charge_gas_for_trie_node_access();
+                trie.set_charge_gas_for_trie_node_access(false);
                 trie
             }
             StorageDataSource::Recorded(storage) => Trie::from_recorded_storage(

core/store/src/trie/mod.rs

@@ -694,11 +694,16 @@ impl Trie {
             )))),
             None => RefCell::new(TrieAccountingCache::new(None)),
         };
+        // Technically the charge_gas_for_trie_node_access should be set based
+        // on the flat storage protocol feature. When flat storage is enabled
+        // the trie node access should be free and the charge flag should be set
+        // to false.
+        let charge_gas_for_trie_node_access = false;
         Trie {
             storage,
             memtries,
             root,
-            charge_gas_for_trie_node_access: flat_storage_chunk_view.is_none(),
+            charge_gas_for_trie_node_access,
             flat_storage_chunk_view,
             accounting_cache,
             recorder: None,
@@ -711,8 +716,8 @@ impl Trie {
     }
 
     /// Helper to simulate gas costs as if flat storage was present.
-    pub fn dont_charge_gas_for_trie_node_access(&mut self) {
-        self.charge_gas_for_trie_node_access = false;
+    pub fn set_charge_gas_for_trie_node_access(&mut self, value: bool) {
+        self.charge_gas_for_trie_node_access = value;
     }
 
     /// Makes a new trie that has everything the same except that access

core/store/src/trie/trie_recording.rs

@@ -469,7 +469,9 @@ mod trie_recording_tests {
                 false,
             )
         } else {
-            tries.get_trie_for_shard(shard_uid, state_root)
+            let mut trie = tries.get_trie_for_shard(shard_uid, state_root);
+            trie.charge_gas_for_trie_node_access = true;
+            trie
         }
     }
 

core/store/src/trie/trie_tests.rs

@@ -183,7 +183,8 @@ mod nodes_counter_tests {
             (create_trie_key(&[0, 1, 1]), Some(vec![1])),
             (create_trie_key(&[1, 0, 0]), Some(vec![2])),
         ];
-        let trie = create_trie(&trie_items);
+        let mut trie = create_trie(&trie_items);
+        trie.charge_gas_for_trie_node_access = true;
         assert_eq!(get_touched_nodes_numbers(&trie, &trie_items), vec![5, 5, 4]);
     }
 
@@ -197,7 +198,8 @@ mod nodes_counter_tests {
             (create_trie_key(&[0, 0]), Some(vec![1])),
             (create_trie_key(&[1, 1]), Some(vec![1])),
         ];
-        let trie = create_trie(&trie_items);
+        let mut trie = create_trie(&trie_items);
+        trie.charge_gas_for_trie_node_access = true;
         assert_eq!(get_touched_nodes_numbers(&trie, &trie_items), vec![4, 4]);
     }
 }

integration-tests/src/test_loop/tests/resharding_v3.rs

@@ -22,7 +22,8 @@ use crate::test_loop::utils::receipts::{
 use crate::test_loop::utils::resharding::fork_before_resharding_block;
 use crate::test_loop::utils::resharding::{
     call_burn_gas_contract, call_promise_yield, check_state_cleanup_after_resharding,
-    execute_money_transfers, temporary_account_during_resharding, TrackedShardSchedule,
+    execute_money_transfers, execute_storage_operations, temporary_account_during_resharding,
+    TrackedShardSchedule,
 };
 use crate::test_loop::utils::sharding::print_and_assert_shard_accounts;
 use crate::test_loop::utils::transactions::{
@@ -482,9 +483,13 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
         }
         latest_block_height.set(tip.height);
 
-        // Check that all chunks are included.
         let client = clients[client_index];
         let block_header = client.chain.get_block_header(&tip.last_block_hash).unwrap();
+        let shard_layout = client.epoch_manager.get_shard_layout(&tip.epoch_id).unwrap();
+        println!("Block: {:?} {} {:?}", tip.last_block_hash, tip.height, block_header.chunk_mask());
+        println!("Shard IDs: {:?}", shard_layout.shard_ids().collect_vec());
+
+        // Check that all chunks are included.
         if params.all_chunks_expected && params.chunk_ranges_to_drop.is_empty() {
             assert!(block_header.chunk_mask().iter().all(|chunk_bit| *chunk_bit));
         }
@@ -733,6 +738,24 @@ fn test_resharding_v3_shard_shuffling_intense() {
     test_resharding_v3_base(params);
 }
 
+/// Executes storage operations at every block height.
+/// In particular, checks that storage gas costs are computed correctly during
+/// resharding. Caught a bug with invalid storage costs computed during flat
+/// storage resharding.
+#[test]
+fn test_resharding_v3_storage_operations() {
+    let sender_account: AccountId = "account1".parse().unwrap();
+    let account_in_parent: AccountId = "account4".parse().unwrap();
+    let params = TestReshardingParametersBuilder::default()
+        .deploy_test_contract(account_in_parent.clone())
+        .add_loop_action(execute_storage_operations(sender_account, account_in_parent))
+        .all_chunks_expected(true)
+        .delay_flat_state_resharding(2)
+        .epoch_length(13)
+        .build();
+    test_resharding_v3_base(params);
+}
+
 #[test]
 #[cfg_attr(not(feature = "test_features"), ignore)]
 fn test_resharding_v3_delayed_receipts_left_child() {

integration-tests/src/test_loop/utils/resharding.rs

@@ -10,6 +10,7 @@ use near_chain::ChainStoreAccess;
 use near_client::Client;
 use near_client::{Query, QueryError::GarbageCollectedBlock};
 use near_crypto::Signer;
+use near_primitives::action::{Action, FunctionCallAction};
 use near_primitives::hash::CryptoHash;
 use near_primitives::test_utils::create_user_test_signer;
 use near_primitives::transaction::SignedTransaction;
@@ -139,6 +140,97 @@ pub(crate) fn execute_money_transfers(account_ids: Vec<AccountId>) -> LoopAction
     LoopAction::new(action_fn, succeeded)
 }
 
+/// Returns a loop action that makes storage read and write at every block
+/// height.
+pub(crate) fn execute_storage_operations(
+    sender_id: AccountId,
+    receiver_id: AccountId,
+) -> LoopAction {
+    const TX_CHECK_DEADLINE: u64 = 5;
+    let latest_height = Cell::new(0);
+    let txs = Cell::new(vec![]);
+    let nonce = Cell::new(102);
+
+    let (ran_transfers, succeeded) = LoopAction::shared_success_flag();
+
+    let action_fn = Box::new(
+        move |node_datas: &[TestData],
+              test_loop_data: &mut TestLoopData,
+              client_account_id: AccountId| {
+            let client_actor =
+                retrieve_client_actor(node_datas, test_loop_data, &client_account_id);
+            let tip = client_actor.client.chain.head().unwrap();
+
+            // Run this action only once at every block height.
+            if latest_height.get() == tip.height {
+                return;
+            }
+            latest_height.set(tip.height);
+
+            let mut remaining_txs = vec![];
+            for (tx, tx_height) in txs.take() {
+                if tx_height + TX_CHECK_DEADLINE >= tip.height {
+                    remaining_txs.push((tx, tx_height));
+                    continue;
+                }
+
+                let tx_outcome = client_actor.client.chain.get_partial_transaction_result(&tx);
+                let status = tx_outcome.as_ref().map(|o| o.status.clone());
+                assert_matches!(status, Ok(FinalExecutionStatus::SuccessValue(_)));
+            }
+            txs.set(remaining_txs);
+
+            let clients = node_datas
+                .iter()
+                .map(|test_data| {
+                    &test_loop_data.get(&test_data.client_sender.actor_handle()).client
+                })
+                .collect_vec();
+
+            // Send transaction which reads a key and writes a key-value pair
+            // to the contract storage.
+            let anchor_hash = get_anchor_hash(&clients);
+            let gas = 20 * TGAS;
+            let salt = 2 * tip.height;
+            nonce.set(nonce.get() + 1);
+            let read_action = Action::FunctionCall(Box::new(FunctionCallAction {
+                args: near_primitives::test_utils::encode(&[salt]),
+                method_name: "read_value".to_string(),
+                gas,
+                deposit: 0,
+            }));
+            let write_action = Action::FunctionCall(Box::new(FunctionCallAction {
+                args: near_primitives::test_utils::encode(&[salt + 1, salt * 10]),
+                method_name: "write_key_value".to_string(),
+                gas,
+                deposit: 0,
+            }));
+            let tx = SignedTransaction::from_actions(
+                nonce.get(),
+                sender_id.clone(),
+                receiver_id.clone(),
+                &create_user_test_signer(&sender_id).into(),
+                vec![read_action, write_action],
+                anchor_hash,
+                0,
+            );
+
+            store_and_submit_tx(
+                &node_datas,
+                &client_account_id,
+                &txs,
+                &sender_id,
+                &receiver_id,
+                tip.height,
+                tx,
+            );
+            ran_transfers.set(true);
+        },
+    );
+
+    LoopAction::new(action_fn, succeeded)
+}
+
 /// Returns a loop action that invokes a costly method from a contract
 /// `CALLS_PER_BLOCK_HEIGHT` times per block height.
 ///

integration-tests/src/user/runtime_user.rs

@@ -102,7 +102,10 @@ impl RuntimeUser {
                     false,
                 )
             } else {
-                client.tries.get_trie_for_shard(ShardUId::single_shard(), client.state_root)
+                let shard_uid = ShardUId::single_shard();
+                let mut trie = client.tries.get_trie_for_shard(shard_uid, client.state_root);
+                trie.set_charge_gas_for_trie_node_access(true);
+                trie
             };
             let apply_result = client
                 .runtime