Always record storage when stateless validation is enabled (#10859)

For stateless validation we need to record all trie reads performed when applying a chunk in order to prepare a `PartialState` that will be included in `ChunkStateWitness` Initially it was only necessary to record trie reads when producing a chunk. Validators could read all the values from the provided `PartialState` without recording anything. A recent change (#10703) introduced a soft limit on the size of `PartialState`. When applying a chunk we watch how much state was recorded, and once the amount of state exceeds the soft limit we stop applying the receipts in this chunk. This needs to be done on both the chunk producer and chunk validator - if the chunk validator doesn't record reads and enforce the limit, it will produce a different result of chunk application, which would lead to validation failure. This means that state should be recorded in all cases when a statelessly validated chunk is applied. Let's remove the configuration option that controls whether trie reads should be recorded (`record_storage`) and just record the reads on every chunk application (when `statelessnet_protocol` feature is enabled). Refs: [zulip conversation](https://near.zulipchat.com/#narrow/stream/407237-core.2Fstateless-validation/topic/State.20witness.20size.20limit/near/428313518)
near · Mar 26, 2024 · c2f9695 · wacban · Mar 28, 2024 · wacban
1 parent 155acad
commit c2f9695
Show file tree

Hide file tree

Showing 10 changed files with 81 additions and 55 deletions.
diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs
@@ -1738,9 +1738,18 @@ impl Chain {
         block_preprocess_info: BlockPreprocessInfo,
         apply_results: Vec<(ShardId, Result<ShardUpdateResult, Error>)>,
     ) -> Result<Option<Tip>, Error> {
+        // Save state transition data to the database only if it might later be needed
+        // for generating a state witness. Storage space optimization.
+        let should_save_state_transition_data =
+            self.should_produce_state_witness_for_this_or_next_epoch(me, block.header())?;
         let mut chain_update = self.chain_update();
-        let new_head =
-            chain_update.postprocess_block(me, &block, block_preprocess_info, apply_results)?;
+        let new_head = chain_update.postprocess_block(
+            me,
+            &block,
+            block_preprocess_info,
+            apply_results,
+            should_save_state_transition_data,
+        )?;
         chain_update.commit()?;
         Ok(new_head)
     }
@@ -2954,9 +2963,17 @@ impl Chain {
         results: Vec<Result<ShardUpdateResult, Error>>,
     ) -> Result<(), Error> {
         let block = self.chain_store.get_block(block_hash)?;
+        // Save state transition data to the database only if it might later be needed
+        // for generating a state witness. Storage space optimization.
+        let should_save_state_transition_data =
+            self.should_produce_state_witness_for_this_or_next_epoch(me, block.header())?;
         let mut chain_update = self.chain_update();
         let results = results.into_iter().collect::<Result<Vec<_>, Error>>()?;
-        chain_update.apply_chunk_postprocessing(&block, results)?;
+        chain_update.apply_chunk_postprocessing(
+            &block,
+            results,
+            should_save_state_transition_data,
+        )?;
         chain_update.commit()?;
 
         let epoch_id = block.header().epoch_id();
@@ -3323,12 +3340,8 @@ impl Chain {
             // only for a single shard. This so far has been enough.
             let state_patch = state_patch.take();
 
-            let storage_context = StorageContext {
-                storage_data_source: StorageDataSource::Db,
-                state_patch,
-                record_storage: self
-                    .should_produce_state_witness_for_this_or_next_epoch(me, block.header())?,
-            };
+            let storage_context =
+                StorageContext { storage_data_source: StorageDataSource::Db, state_patch };
             let stateful_job = self.get_update_shard_job(
                 me,
                 block,

diff --git a/chain/chain/src/chain_update.rs b/chain/chain/src/chain_update.rs
@@ -145,10 +145,11 @@ impl<'a> ChainUpdate<'a> {
         &mut self,
         block: &Block,
         apply_results: Vec<ShardUpdateResult>,
+        should_save_state_transition_data: bool,
     ) -> Result<(), Error> {
         let _span = tracing::debug_span!(target: "chain", "apply_chunk_postprocessing").entered();
         for result in apply_results {
-            self.process_apply_chunk_result(block, result)?;
+            self.process_apply_chunk_result(block, result, should_save_state_transition_data)?;
         }
         Ok(())
     }
@@ -299,6 +300,7 @@ impl<'a> ChainUpdate<'a> {
         &mut self,
         block: &Block,
         result: ShardUpdateResult,
+        should_save_state_transition_data: bool,
     ) -> Result<(), Error> {
         let block_hash = block.hash();
         let prev_hash = block.header().prev_hash();
@@ -351,12 +353,14 @@ impl<'a> ChainUpdate<'a> {
                     apply_result.outcomes,
                     outcome_paths,
                 );
-                self.chain_store_update.save_state_transition_data(
-                    *block_hash,
-                    shard_id,
-                    apply_result.proof,
-                    apply_result.applied_receipts_hash,
-                );
+                if should_save_state_transition_data {
+                    self.chain_store_update.save_state_transition_data(
+                        *block_hash,
+                        shard_id,
+                        apply_result.proof,
+                        apply_result.applied_receipts_hash,
+                    );
+                }
                 if let Some(resharding_results) = resharding_results {
                     self.process_resharding_results(block, &shard_uid, resharding_results)?;
                 }
@@ -383,12 +387,14 @@ impl<'a> ChainUpdate<'a> {
 
                 self.chain_store_update.save_chunk_extra(block_hash, &shard_uid, new_extra);
                 self.chain_store_update.save_trie_changes(apply_result.trie_changes);
-                self.chain_store_update.save_state_transition_data(
-                    *block_hash,
-                    shard_uid.shard_id(),
-                    apply_result.proof,
-                    apply_result.applied_receipts_hash,
-                );
+                if should_save_state_transition_data {
+                    self.chain_store_update.save_state_transition_data(
+                        *block_hash,
+                        shard_uid.shard_id(),
+                        apply_result.proof,
+                        apply_result.applied_receipts_hash,
+                    );
+                }
 
                 if let Some(resharding_config) = resharding_results {
                     self.process_resharding_results(block, &shard_uid, resharding_config)?;
@@ -413,6 +419,7 @@ impl<'a> ChainUpdate<'a> {
         block: &Block,
         block_preprocess_info: BlockPreprocessInfo,
         apply_chunks_results: Vec<(ShardId, Result<ShardUpdateResult, Error>)>,
+        should_save_state_transition_data: bool,
     ) -> Result<Option<Tip>, Error> {
         let shard_ids = self.epoch_manager.shard_ids(block.header().epoch_id())?;
         let prev_hash = block.header().prev_hash();
@@ -422,7 +429,7 @@ impl<'a> ChainUpdate<'a> {
             }
             x
         }).collect::<Result<Vec<_>, Error>>()?;
-        self.apply_chunk_postprocessing(block, results)?;
+        self.apply_chunk_postprocessing(block, results, should_save_state_transition_data)?;
 
         let BlockPreprocessInfo {
             is_caught_up,

diff --git a/chain/chain/src/runtime/mod.rs b/chain/chain/src/runtime/mod.rs
@@ -15,6 +15,7 @@ use near_epoch_manager::{EpochManagerAdapter, EpochManagerHandle};
 use near_parameters::{ActionCosts, ExtCosts, RuntimeConfigStore};
 use near_pool::types::TransactionGroupIterator;
 use near_primitives::account::{AccessKey, Account};
+use near_primitives::checked_feature;
 use near_primitives::errors::{InvalidTxError, RuntimeError, StorageError};
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::receipt::{DelayedReceiptIndices, Receipt};
@@ -30,7 +31,7 @@ use near_primitives::types::{
     AccountId, Balance, BlockHeight, EpochHeight, EpochId, EpochInfoProvider, Gas, MerkleHash,
     ShardId, StateChangeCause, StateChangesForResharding, StateRoot, StateRootNode,
 };
-use near_primitives::version::ProtocolVersion;
+use near_primitives::version::{ProtocolVersion, PROTOCOL_VERSION};
 use near_primitives::views::{
     AccessKeyInfoView, CallResult, ContractCodeView, QueryRequest, QueryResponse,
     QueryResponseKind, ViewApplyState, ViewStateResult,
@@ -704,7 +705,9 @@ impl RuntimeAdapter for NightshadeRuntime {
                 storage_config.use_flat_storage,
             ),
         };
-        if storage_config.record_storage {
+        if checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION)
+            || cfg!(feature = "shadow_chunk_validation")
+        {
             trie = trie.recording_reads();
         }
         let mut state_update = TrieUpdate::new(trie);
@@ -865,7 +868,9 @@ impl RuntimeAdapter for NightshadeRuntime {
                 storage_config.use_flat_storage,
             ),
         };
-        if storage_config.record_storage {
+        if checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION)
+            || cfg!(feature = "shadow_chunk_validation")
+        {
             trie = trie.recording_reads();
         }
 

diff --git a/chain/chain/src/runtime/tests.rs b/chain/chain/src/runtime/tests.rs
@@ -9,8 +9,10 @@ use near_epoch_manager::{EpochManager, RngSeed};
 use near_pool::{
     InsertTransactionResult, PoolIteratorWrapper, TransactionGroupIteratorWrapper, TransactionPool,
 };
+use near_primitives::checked_feature;
 use near_primitives::test_utils::create_test_signer;
 use near_primitives::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
+use near_primitives::version::PROTOCOL_VERSION;
 use near_store::flat::{FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata};
 use near_store::genesis::initialize_genesis_state;
 use num_rational::Ratio;
@@ -1601,6 +1603,11 @@ fn prepare_transactions(
 /// Check that transactions validation works the same when using recorded storage proof instead of db.
 #[test]
 fn test_prepare_transactions_storage_proof() {
+    if !checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION) {
+        println!("Test not applicable without StatelessValidation enabled");
+        return;
+    }
+
     let (env, chain, mut transaction_pool) = get_test_env_with_chain_and_pool();
     let transactions_count = transaction_pool.len();
 
@@ -1609,7 +1616,6 @@ fn test_prepare_transactions_storage_proof() {
         use_flat_storage: true,
         source: StorageDataSource::Db,
         state_patch: Default::default(),
-        record_storage: true,
     };
 
     let proposed_transactions = prepare_transactions(
@@ -1630,7 +1636,6 @@ fn test_prepare_transactions_storage_proof() {
             nodes: proposed_transactions.storage_proof.unwrap(),
         }),
         state_patch: Default::default(),
-        record_storage: false,
     };
 
     let validated_transactions = prepare_transactions(
@@ -1647,6 +1652,11 @@ fn test_prepare_transactions_storage_proof() {
 /// Check that transactions validation fails if provided empty storage proof.
 #[test]
 fn test_prepare_transactions_empty_storage_proof() {
+    if !checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION) {
+        println!("Test not applicable without StatelessValidation enabled");
+        return;
+    }
+
     let (env, chain, mut transaction_pool) = get_test_env_with_chain_and_pool();
     let transactions_count = transaction_pool.len();
 
@@ -1655,7 +1665,6 @@ fn test_prepare_transactions_empty_storage_proof() {
         use_flat_storage: true,
         source: StorageDataSource::Db,
         state_patch: Default::default(),
-        record_storage: true,
     };
 
     let proposed_transactions = prepare_transactions(
@@ -1676,7 +1685,6 @@ fn test_prepare_transactions_empty_storage_proof() {
             nodes: PartialState::default(), // We use empty storage proof here.
         }),
         state_patch: Default::default(),
-        record_storage: false,
     };
 
     let validation_result = prepare_transactions(

diff --git a/chain/chain/src/test_utils/kv_runtime.rs b/chain/chain/src/test_utils/kv_runtime.rs
@@ -24,7 +24,6 @@ use near_primitives::epoch_manager::ValidatorSelectionConfig;
 use near_primitives::errors::{EpochError, InvalidTxError};
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::receipt::{ActionReceipt, Receipt, ReceiptEnum};
-use near_primitives::shard_layout;
 use near_primitives::shard_layout::{ShardLayout, ShardUId};
 use near_primitives::sharding::{ChunkHash, ShardChunkHeader};
 use near_primitives::state_part::PartId;
@@ -45,6 +44,7 @@ use near_primitives::views::{
     AccessKeyInfoView, AccessKeyList, CallResult, ContractCodeView, EpochValidatorInfo,
     QueryRequest, QueryResponse, QueryResponseKind, ViewStateResult,
 };
+use near_primitives::{checked_feature, shard_layout};
 use near_store::test_utils::TestTriesBuilder;
 use near_store::{
     set_genesis_hash, set_genesis_state_roots, DBCol, ShardTries, StorageError, Store, StoreUpdate,
@@ -1083,7 +1083,7 @@ impl RuntimeAdapter for KeyValueRuntime {
 
     fn prepare_transactions(
         &self,
-        storage: RuntimeStorageConfig,
+        _storage: RuntimeStorageConfig,
         _chunk: PrepareTransactionsChunkContext,
         _prev_block: PrepareTransactionsBlockContext,
         transaction_groups: &mut dyn TransactionGroupIterator,
@@ -1094,11 +1094,14 @@ impl RuntimeAdapter for KeyValueRuntime {
         while let Some(iter) = transaction_groups.next() {
             res.push(iter.next().unwrap());
         }
-        Ok(PreparedTransactions {
-            transactions: res,
-            limited_by: None,
-            storage_proof: if storage.record_storage { Some(Default::default()) } else { None },
-        })
+        let storage_proof = if checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION)
+            || cfg!(feature = "shadow_chunk_validation")
+        {
+            Some(Default::default())
+        } else {
+            None
+        };
+        Ok(PreparedTransactions { transactions: res, limited_by: None, storage_proof })
     }
 
     fn apply_chunk(
@@ -1242,7 +1245,13 @@ impl RuntimeAdapter for KeyValueRuntime {
         let state_root = hash(&data);
         self.state.write().unwrap().insert(state_root, state);
         self.state_size.write().unwrap().insert(state_root, state_size);
-
+        let storage_proof = if checked_feature!("stable", StatelessValidationV0, PROTOCOL_VERSION)
+            || cfg!(feature = "shadow_chunk_validation")
+        {
+            Some(Default::default())
+        } else {
+            None
+        };
         Ok(ApplyChunkResult {
             trie_changes: WrappedTrieChanges::new(
                 self.get_tries(),
@@ -1258,7 +1267,7 @@ impl RuntimeAdapter for KeyValueRuntime {
             validator_proposals: vec![],
             total_gas_burnt: 0,
             total_balance_burnt: 0,
-            proof: if storage_config.record_storage { Some(Default::default()) } else { None },
+            proof: storage_proof,
             processed_delayed_receipts: vec![],
             applied_receipts_hash: hash(&borsh::to_vec(receipts).unwrap()),
         })

diff --git a/chain/chain/src/types.rs b/chain/chain/src/types.rs
@@ -263,7 +263,6 @@ pub struct RuntimeStorageConfig {
     pub use_flat_storage: bool,
     pub source: StorageDataSource,
     pub state_patch: SandboxStatePatch,
-    pub record_storage: bool,
 }
 
 impl RuntimeStorageConfig {
@@ -273,7 +272,6 @@ impl RuntimeStorageConfig {
             use_flat_storage,
             source: StorageDataSource::Db,
             state_patch: Default::default(),
-            record_storage: false,
         }
     }
 }

diff --git a/chain/chain/src/update_shard.rs b/chain/chain/src/update_shard.rs
@@ -115,7 +115,6 @@ pub struct StorageContext {
     /// Data source used for processing shard update.
     pub storage_data_source: StorageDataSource,
     pub state_patch: SandboxStatePatch,
-    pub record_storage: bool,
 }
 
 /// Processes shard update with given block and shard.
@@ -185,7 +184,6 @@ pub fn apply_new_chunk(
         use_flat_storage: true,
         source: storage_context.storage_data_source,
         state_patch: storage_context.state_patch,
-        record_storage: storage_context.record_storage,
     };
     match runtime.apply_chunk(
         storage_config,
@@ -247,7 +245,6 @@ pub fn apply_old_chunk(
         use_flat_storage: true,
         source: storage_context.storage_data_source,
         state_patch: storage_context.state_patch,
-        record_storage: storage_context.record_storage,
     };
     match runtime.apply_chunk(
         storage_config,

diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs
@@ -999,18 +999,11 @@ impl Client {
         let prepared_transactions = if let Some(mut iter) =
             sharded_tx_pool.get_pool_iterator(shard_uid)
         {
-            let me = self
-                .validator_signer
-                .as_ref()
-                .map(|validator_signer| validator_signer.validator_id().clone());
-            let record_storage = chain
-                .should_produce_state_witness_for_this_or_next_epoch(&me, &prev_block_header)?;
             let storage_config = RuntimeStorageConfig {
                 state_root: *chunk_extra.state_root(),
                 use_flat_storage: true,
                 source: StorageDataSource::Db,
                 state_patch: Default::default(),
-                record_storage,
             };
             runtime.prepare_transactions(
                 storage_config,

diff --git a/chain/client/src/stateless_validation/chunk_validator/mod.rs b/chain/client/src/stateless_validation/chunk_validator/mod.rs
@@ -262,7 +262,6 @@ pub(crate) fn pre_validate_chunk_state_witness(
                 nodes: state_witness.new_transactions_validation_state.clone(),
             }),
             state_patch: Default::default(),
-            record_storage: false,
         };
 
         match validate_prepared_transactions(
@@ -314,7 +313,6 @@ pub(crate) fn pre_validate_chunk_state_witness(
                     nodes: state_witness.main_state_transition.base_state.clone(),
                 }),
                 state_patch: Default::default(),
-                record_storage: false,
             },
         })
     };
@@ -529,7 +527,6 @@ pub(crate) fn validate_chunk_state_witness(
                     nodes: transition.base_state,
                 }),
                 state_patch: Default::default(),
-                record_storage: false,
             },
         };
         let OldChunkResult { apply_result, .. } = apply_old_chunk(