fix(resharding): storage costs

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));
         }
@@ -701,6 +706,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