feat: stateless validation jobs in test mode

chain/chain/src/types.rs

@@ -248,7 +248,14 @@ impl ChainGenesis {
 }
 
 pub enum StorageDataSource {
+    /// Full state data is present in DB.
     Db,
+    /// Trie is present in DB and flat storage is not.
+    /// Used for testing stateless validation jobs, should be removed after
+    /// stateless validation release.
+    DbTrieOnly,
+    /// State data is supplied from state witness, there is no state data
+    /// stored on disk.
     Recorded(PartialStorage),
 }
 

chain/chain/src/update_shard.rs

@@ -2,7 +2,7 @@ use crate::crypto_hash_timer::CryptoHashTimer;
 use crate::types::{
     ApplySplitStateResult, ApplySplitStateResultOrStateChanges, ApplyTransactionResult,
     ApplyTransactionsBlockContext, ApplyTransactionsChunkContext, RuntimeAdapter,
-    RuntimeStorageConfig,
+    RuntimeStorageConfig, StorageDataSource,
 };
 use near_chain_primitives::Error;
 use near_epoch_manager::EpochManagerAdapter;
@@ -45,8 +45,20 @@ pub struct StateSplitResult {
     pub(crate) results: Vec<ApplySplitStateResult>,
 }
 
+/// Result of processing shard update, covering both stateful and stateless scenarios.
 #[derive(Debug)]
-pub enum ApplyChunkResult {
+pub enum ShardUpdateResult {
+    /// Stateful scenario - processed update for a single block.
+    Stateful(ShardBlockUpdateResult),
+    /// Stateless scenario - processed update based on state witness in a chunk.
+    /// Contains `ChunkExtra`s - results for processing updates corresponding
+    /// to state witness.
+    Stateless(Vec<(CryptoHash, ShardUId, ChunkExtra)>),
+}
+
+/// Result for a shard update for a single block.
+#[derive(Debug)]
+pub enum ShardBlockUpdateResult {
     NewChunk(NewChunkResult),
     OldChunk(OldChunkResult),
     StateSplit(StateSplitResult),
@@ -61,12 +73,14 @@ pub(crate) struct NewChunkData {
     pub split_state_roots: Option<SplitStateRoots>,
     pub block: ApplyTransactionsBlockContext,
     pub is_first_block_with_chunk_of_version: bool,
+    pub storage_context: StorageContext,
 }
 
 pub(crate) struct OldChunkData {
     pub prev_chunk_extra: ChunkExtra,
     pub split_state_roots: Option<SplitStateRoots>,
     pub block: ApplyTransactionsBlockContext,
+    pub storage_context: StorageContext,
 }
 
 pub(crate) struct StateSplitData {
@@ -94,8 +108,21 @@ pub(crate) enum ShardUpdateReason {
 /// Information about shard to update.
 pub(crate) struct ShardContext {
     pub shard_uid: ShardUId,
+    /// Whether node cares about shard in this epoch.
+    pub cares_about_shard_this_epoch: bool,
     /// Whether shard layout changes in the next epoch.
     pub will_shard_layout_change: bool,
+    /// Whether transactions should be applied.
+    pub should_apply_transactions: bool,
+    /// See comment in `get_update_shard_job`.
+    pub need_to_split_states: bool,
+}
+
+/// Information about storage used for applying txs and receipts.
+pub(crate) struct StorageContext {
+    /// Data source used for processing shard update.
+    pub storage_data_source: StorageDataSource,
+    pub state_patch: SandboxStatePatch,
 }
 
 /// Processes shard update with given block and shard.
@@ -106,14 +133,13 @@ pub(crate) fn process_shard_update(
     epoch_manager: &dyn EpochManagerAdapter,
     shard_update_reason: ShardUpdateReason,
     shard_context: ShardContext,
-    state_patch: SandboxStatePatch,
-) -> Result<ApplyChunkResult, Error> {
+) -> Result<ShardBlockUpdateResult, Error> {
     match shard_update_reason {
         ShardUpdateReason::NewChunk(data) => {
-            apply_new_chunk(parent_span, data, shard_context, state_patch, runtime, epoch_manager)
+            apply_new_chunk(parent_span, data, shard_context, runtime, epoch_manager)
         }
         ShardUpdateReason::OldChunk(data) => {
-            apply_old_chunk(parent_span, data, shard_context, state_patch, runtime, epoch_manager)
+            apply_old_chunk(parent_span, data, shard_context, runtime, epoch_manager)
         }
         ShardUpdateReason::StateSplit(data) => {
             apply_state_split(parent_span, data, shard_context.shard_uid, runtime, epoch_manager)
@@ -126,19 +152,19 @@ pub(crate) fn process_shard_update(
 fn apply_new_chunk(
     parent_span: &tracing::Span,
     data: NewChunkData,
-    shard_info: ShardContext,
-    state_patch: SandboxStatePatch,
+    shard_context: ShardContext,
     runtime: &dyn RuntimeAdapter,
     epoch_manager: &dyn EpochManagerAdapter,
-) -> Result<ApplyChunkResult, Error> {
+) -> Result<ShardBlockUpdateResult, Error> {
     let NewChunkData {
         block,
         chunk,
         receipts,
         split_state_roots,
         is_first_block_with_chunk_of_version,
+        storage_context,
     } = data;
-    let shard_id = shard_info.shard_uid.shard_id();
+    let shard_id = shard_context.shard_uid.shard_id();
     let _span = tracing::debug_span!(
         target: "chain",
         parent: parent_span,
@@ -152,8 +178,8 @@ fn apply_new_chunk(
     let storage_config = RuntimeStorageConfig {
         state_root: *chunk_inner.prev_state_root(),
         use_flat_storage: true,
-        source: crate::types::StorageDataSource::Db,
-        state_patch,
+        source: storage_context.storage_data_source,
+        state_patch: storage_context.state_patch,
         record_storage: false,
     };
     match runtime.apply_transactions(
@@ -170,7 +196,7 @@ fn apply_new_chunk(
         chunk.transactions(),
     ) {
         Ok(apply_result) => {
-            let apply_split_result_or_state_changes = if shard_info.will_shard_layout_change {
+            let apply_split_result_or_state_changes = if shard_context.will_shard_layout_change {
                 Some(apply_split_state_changes(
                     epoch_manager,
                     runtime,
@@ -181,9 +207,9 @@ fn apply_new_chunk(
             } else {
                 None
             };
-            Ok(ApplyChunkResult::NewChunk(NewChunkResult {
+            Ok(ShardBlockUpdateResult::NewChunk(NewChunkResult {
                 gas_limit,
-                shard_uid: shard_info.shard_uid,
+                shard_uid: shard_context.shard_uid,
                 apply_result,
                 apply_split_result_or_state_changes,
             }))
@@ -198,13 +224,12 @@ fn apply_new_chunk(
 fn apply_old_chunk(
     parent_span: &tracing::Span,
     data: OldChunkData,
-    shard_info: ShardContext,
-    state_patch: SandboxStatePatch,
+    shard_context: ShardContext,
     runtime: &dyn RuntimeAdapter,
     epoch_manager: &dyn EpochManagerAdapter,
-) -> Result<ApplyChunkResult, Error> {
-    let OldChunkData { prev_chunk_extra, split_state_roots, block } = data;
-    let shard_id = shard_info.shard_uid.shard_id();
+) -> Result<ShardBlockUpdateResult, Error> {
+    let OldChunkData { prev_chunk_extra, split_state_roots, block, storage_context } = data;
+    let shard_id = shard_context.shard_uid.shard_id();
     let _span = tracing::debug_span!(
         target: "chain",
         parent: parent_span,
@@ -215,8 +240,8 @@ fn apply_old_chunk(
     let storage_config = RuntimeStorageConfig {
         state_root: *prev_chunk_extra.state_root(),
         use_flat_storage: true,
-        source: crate::types::StorageDataSource::Db,
-        state_patch,
+        source: storage_context.storage_data_source,
+        state_patch: storage_context.state_patch,
         record_storage: false,
     };
     match runtime.apply_transactions(
@@ -233,7 +258,7 @@ fn apply_old_chunk(
         &[],
     ) {
         Ok(apply_result) => {
-            let apply_split_result_or_state_changes = if shard_info.will_shard_layout_change {
+            let apply_split_result_or_state_changes = if shard_context.will_shard_layout_change {
                 Some(apply_split_state_changes(
                     epoch_manager,
                     runtime,
@@ -244,8 +269,8 @@ fn apply_old_chunk(
             } else {
                 None
             };
-            Ok(ApplyChunkResult::OldChunk(OldChunkResult {
-                shard_uid: shard_info.shard_uid,
+            Ok(ShardBlockUpdateResult::OldChunk(OldChunkResult {
+                shard_uid: shard_context.shard_uid,
                 apply_result,
                 apply_split_result_or_state_changes,
             }))
@@ -261,7 +286,7 @@ fn apply_state_split(
     shard_uid: ShardUId,
     runtime: &dyn RuntimeAdapter,
     epoch_manager: &dyn EpochManagerAdapter,
-) -> Result<ApplyChunkResult, Error> {
+) -> Result<ShardBlockUpdateResult, Error> {
     let StateSplitData { split_state_roots, state_changes, block_height: height, block_hash } =
         data;
     let shard_id = shard_uid.shard_id();
@@ -281,7 +306,7 @@ fn apply_state_split(
         &next_epoch_shard_layout,
         state_changes,
     )?;
-    Ok(ApplyChunkResult::StateSplit(StateSplitResult { shard_uid, results }))
+    Ok(ShardBlockUpdateResult::StateSplit(StateSplitResult { shard_uid, results }))
 }
 
 /// Process ApplyTransactionResult to apply changes to split states

chain/client/src/client_actor.rs

@@ -22,6 +22,7 @@ use crate::{metrics, StatusResponse, SyncAdapter};
 use actix::{Actor, Addr, Arbiter, AsyncContext, Context, Handler};
 use actix_rt::ArbiterHandle;
 use chrono::{DateTime, Utc};
+use itertools::Itertools;
 use near_async::messaging::{CanSend, Sender};
 use near_chain::chain::{
     ApplyStatePartsRequest, ApplyStatePartsResponse, BlockCatchUpRequest, BlockCatchUpResponse,
@@ -1889,7 +1890,9 @@ impl Handler<WithSpanContext<BlockCatchUpResponse>> for ClientActor {
             self.client.catchup_state_syncs.get_mut(&msg.sync_hash)
         {
             assert!(blocks_catch_up_state.scheduled_blocks.remove(&msg.block_hash));
-            blocks_catch_up_state.processed_blocks.insert(msg.block_hash, msg.results);
+            blocks_catch_up_state
+                .processed_blocks
+                .insert(msg.block_hash, msg.results.into_iter().map(|res| res.1).collect_vec());
         } else {
             panic!("block catch up processing result from unknown sync hash");
         }

chain/client/src/test_utils/client.rs

@@ -7,6 +7,7 @@ use std::sync::{Arc, RwLock};
 
 use crate::Client;
 use actix_rt::{Arbiter, System};
+use itertools::Itertools;
 use near_chain::chain::{do_apply_chunks, BlockCatchUpRequest};
 use near_chain::resharding::StateSplitRequest;
 use near_chain::test_utils::{wait_for_all_blocks_in_processing, wait_for_block_in_processing};
@@ -242,7 +243,10 @@ pub fn run_catchup(
         )?;
         let mut catchup_done = true;
         for msg in block_messages.write().unwrap().drain(..) {
-            let results = do_apply_chunks(msg.block_hash, msg.block_height, msg.work);
+            let results = do_apply_chunks(msg.block_hash, msg.block_height, msg.work)
+                .into_iter()
+                .map(|res| res.1)
+                .collect_vec();
             if let Some((_, _, blocks_catch_up_state)) =
                 client.catchup_state_syncs.get_mut(&msg.sync_hash)
             {

core/primitives/src/challenge.rs

@@ -17,6 +17,12 @@ pub enum PartialState {
     TrieValues(Vec<TrieValue>),
 }
 
+impl Default for PartialState {
+    fn default() -> Self {
+        PartialState::TrieValues(vec![])
+    }
+}
+
 impl PartialState {
     pub fn len(&self) -> usize {
         let Self::TrieValues(values) = self;

core/store/src/trie/mod.rs

@@ -60,7 +60,7 @@ pub mod update;
 const POISONED_LOCK_ERR: &str = "The lock was poisoned.";
 
 /// For fraud proofs
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct PartialStorage {
     pub nodes: PartialState,
 }
@@ -649,6 +649,11 @@ impl Trie {
         }
     }
 
+    /// Temporary helper, must be removed after stateless validation release.
+    pub fn dont_charge_gas_for_trie_node_access(&mut self) {
+        self.charge_gas_for_trie_node_access = false;
+    }
+
     /// Makes a new trie that has everything the same except that access
     /// through that trie accumulates a state proof for all nodes accessed.
     pub fn recording_reads(&self) -> Self {

integration-tests/src/tests/client/process_blocks.rs

@@ -11,9 +11,12 @@ use near_async::messaging::IntoSender;
 use near_chain::chain::ApplyStatePartsRequest;
 use near_chain::test_utils::ValidatorSchedule;
 use near_chain::types::{LatestKnown, RuntimeAdapter};
+#[cfg(not(feature = "nightly"))]
 use near_chain::validate::validate_chunk_with_chunk_extra;
+#[cfg(not(feature = "nightly"))]
+use near_chain::ChainStore;
 use near_chain::{
-    Block, BlockProcessingArtifact, ChainGenesis, ChainStore, ChainStoreAccess, Error, Provenance,
+    Block, BlockProcessingArtifact, ChainGenesis, ChainStoreAccess, Error, Provenance,
 };
 use near_chain_configs::{Genesis, DEFAULT_GC_NUM_EPOCHS_TO_KEEP};
 use near_chunks::test_utils::MockClientAdapterForShardsManager;
@@ -2252,7 +2255,10 @@ fn test_block_height_processed_orphan() {
     assert!(env.clients[0].chain.mut_store().is_height_processed(block_height).unwrap());
 }
 
+// Disabled until stateless validation release, because the test relies on
+// logging which is impacted by the release process.
 #[test]
+#[cfg(not(feature = "nightly"))]
 fn test_validate_chunk_extra() {
     let mut capture = near_o11y::testonly::TracingCapture::enable();
 

nearcore/src/runtime/mod.rs

@@ -856,6 +856,18 @@ impl RuntimeAdapter for NightshadeRuntime {
                 storage_config.state_root,
                 storage_config.use_flat_storage,
             )?,
+            StorageDataSource::DbTrieOnly => {
+                // If there is no flat storage on disk, use trie but simulate costs with enabled
+                // flat storage by not charging gas for trie nodes.
+                let mut trie = self.get_trie_for_shard(
+                    shard_id,
+                    &block.prev_block_hash,
+                    storage_config.state_root,
+                    false,
+                )?;
+                trie.dont_charge_gas_for_trie_node_access();
+                trie
+            }
             StorageDataSource::Recorded(storage) => Trie::from_recorded_storage(
                 storage,
                 storage_config.state_root,