runtime: implement a receipt preparation pipeline

core/store/src/contract.rs

@@ -0,0 +1,52 @@
+use crate::TrieStorage;
+use near_primitives::hash::CryptoHash;
+use near_vm_runner::ContractCode;
+use std::collections::BTreeMap;
+use std::sync::{Arc, Mutex};
+
+/// Reads contract code from the trie by its hash.
+///
+/// Cloning is cheap.
+#[derive(Clone)]
+pub struct Storage {
+    storage: Arc<dyn TrieStorage>,
+
+    /// During an apply of a single chunk contracts may be deployed through the
+    /// `Action::DeployContract`.
+    ///
+    /// Unfortunately `TrieStorage` does not have a way to write to the underlying storage, and the
+    /// `TrieUpdate` will only write the contract once the whole transaction is committed at the
+    /// end of the chunk's apply.
+    ///
+    /// As a temporary work-around while we're still involving `Trie` for `ContractCode` storage,
+    /// we'll keep a list of such deployed contracts here. Once the contracts are no longer part of
+    /// The State this field should be removed, and the `Storage::store` function should be
+    /// adjusted to write out the contract into the relevant part of the database immediately
+    /// (without going through transactional storage operations and such).
+    uncommitted_deploys: Arc<Mutex<BTreeMap<CryptoHash, ContractCode>>>,
+}
+
+impl Storage {
+    pub fn new(storage: Arc<dyn TrieStorage>) -> Self {
+        Self { storage, uncommitted_deploys: Default::default() }
+    }
+
+    pub fn get(&self, code_hash: CryptoHash) -> Option<ContractCode> {
+        {
+            let guard = self.uncommitted_deploys.lock().expect("no panics");
+            if let Some(v) = guard.get(&code_hash) {
+                return Some(ContractCode::new(v.code().to_vec(), Some(code_hash)));
+            }
+        }
+
+        match self.storage.retrieve_raw_bytes(&code_hash) {
+            Ok(raw_code) => Some(ContractCode::new(raw_code.to_vec(), Some(code_hash))),
+            Err(_) => None,
+        }
+    }
+
+    pub fn store(&self, code: ContractCode) {
+        let mut guard = self.uncommitted_deploys.lock().expect("no panics");
+        guard.insert(*code.hash(), code);
+    }
+}

core/store/src/lib.rs

@@ -44,6 +44,7 @@ use strum;
 pub mod cold_storage;
 mod columns;
 pub mod config;
+pub mod contract;
 pub mod db;
 pub mod flat;
 pub mod genesis;
@@ -950,15 +951,6 @@ pub fn set_code(state_update: &mut TrieUpdate, account_id: AccountId, code: &Con
     state_update.set(TrieKey::ContractCode { account_id }, code.code().to_vec());
 }
 
-pub fn get_code(
-    trie: &dyn TrieAccess,
-    account_id: &AccountId,
-    code_hash: Option<CryptoHash>,
-) -> Result<Option<ContractCode>, StorageError> {
-    let key = TrieKey::ContractCode { account_id: account_id.clone() };
-    trie.get(&key).map(|opt| opt.map(|code| ContractCode::new(code, code_hash)))
-}
-
 /// Removes account, code and all access keys associated to it.
 pub fn remove_account(
     state_update: &mut TrieUpdate,

core/store/src/trie/update.rs

@@ -2,40 +2,17 @@ pub use self::iterator::TrieUpdateIterator;
 use super::accounting_cache::TrieAccountingCacheSwitch;
 use super::{OptimizedValueRef, Trie, TrieWithReadLock};
 use crate::trie::{KeyLookupMode, TrieChanges};
-use crate::{StorageError, TrieStorage};
+use crate::StorageError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{
     AccountId, RawStateChange, RawStateChanges, RawStateChangesWithTrieKey, StateChangeCause,
     StateRoot, TrieCacheMode,
 };
-use near_vm_runner::ContractCode;
 use std::collections::BTreeMap;
-use std::sync::Arc;
 
 mod iterator;
 
-/// Reads contract code from the trie by its hash.
-/// Currently, uses `TrieStorage`. Consider implementing separate logic for
-/// requesting and compiling contracts, as any contract code read and
-/// compilation is a major bottleneck during chunk execution.
-struct ContractStorage {
-    storage: Arc<dyn TrieStorage>,
-}
-
-impl ContractStorage {
-    fn new(storage: Arc<dyn TrieStorage>) -> Self {
-        Self { storage }
-    }
-
-    pub fn get(&self, code_hash: CryptoHash) -> Option<ContractCode> {
-        match self.storage.retrieve_raw_bytes(&code_hash) {
-            Ok(raw_code) => Some(ContractCode::new(raw_code.to_vec(), Some(code_hash))),
-            Err(_) => None,
-        }
-    }
-}
-
 /// Key-value update. Contains a TrieKey and a value.
 pub struct TrieKeyValueUpdate {
     pub trie_key: TrieKey,
@@ -49,7 +26,7 @@ pub type TrieUpdates = BTreeMap<Vec<u8>, TrieKeyValueUpdate>;
 /// TODO (#7327): rename to StateUpdate
 pub struct TrieUpdate {
     pub trie: Trie,
-    contract_storage: ContractStorage,
+    pub contract_storage: crate::contract::Storage,
     committed: RawStateChanges,
     prospective: TrieUpdates,
 }
@@ -80,7 +57,7 @@ impl TrieUpdate {
         let trie_storage = trie.storage.clone();
         Self {
             trie,
-            contract_storage: ContractStorage::new(trie_storage),
+            contract_storage: crate::contract::Storage::new(trie_storage),
             committed: Default::default(),
             prospective: Default::default(),
         }

integration-tests/src/tests/runtime/state_viewer.rs

@@ -1,11 +1,12 @@
 use std::{collections::HashMap, io, sync::Arc};
 
 use borsh::BorshDeserialize;
+use near_vm_runner::ContractCode;
 
 use crate::runtime_utils::{get_runtime_and_trie, get_test_trie_viewer, TEST_SHARD_UID};
 use near_primitives::{
     account::Account,
-    hash::{hash as sha256, CryptoHash},
+    hash::CryptoHash,
     serialize::to_base64,
     trie_key::trie_key_parsers,
     types::{AccountId, StateRoot},
@@ -374,12 +375,14 @@ fn test_view_state_with_large_contract() {
     let (_, tries, root) = get_runtime_and_trie();
     let mut state_update = tries.new_trie_update(TEST_SHARD_UID, root);
     let contract_code = [0; Account::MAX_ACCOUNT_DELETION_STORAGE_USAGE as usize].to_vec();
+    let code = ContractCode::new(contract_code, None);
     set_account(
         &mut state_update,
         alice_account(),
-        &Account::new(0, 0, 0, sha256(&contract_code), 50_001, PROTOCOL_VERSION),
+        &Account::new(0, 0, 0, *code.hash(), 50_001, PROTOCOL_VERSION),
     );
-    state_update.set(TrieKey::ContractCode { account_id: alice_account() }, contract_code);
+    // FIXME: this really should use the deploy action.
+    state_update.contract_storage.store(code);
     let trie_viewer = TrieViewer::new(Some(50_000), None);
     let result = trie_viewer.view_state(&state_update, &alice_account(), b"", false);
     assert!(result.is_ok());

runtime/runtime/src/actions.rs

@@ -2,7 +2,7 @@ use crate::config::{
     safe_add_compute, safe_add_gas, total_prepaid_exec_fees, total_prepaid_gas,
     total_prepaid_send_fees,
 };
-use crate::ext::{ExternalError, RuntimeContractExt, RuntimeExt};
+use crate::ext::{ExternalError, RuntimeExt};
 use crate::receipt_manager::ReceiptManager;
 use crate::{metrics, ActionResult, ApplyState};
 use near_crypto::PublicKey;
@@ -30,16 +30,15 @@ use near_primitives::version::{
 };
 use near_primitives_core::account::id::AccountType;
 use near_store::{
-    enqueue_promise_yield_timeout, get_access_key, get_code, get_promise_yield_indices,
-    remove_access_key, remove_account, set_access_key, set_code, set_promise_yield_indices,
-    StorageError, TrieUpdate,
+    enqueue_promise_yield_timeout, get_access_key, get_promise_yield_indices, remove_access_key,
+    remove_account, set_access_key, set_code, set_promise_yield_indices, StorageError, TrieUpdate,
 };
 use near_vm_runner::logic::errors::{
     CompilationError, FunctionCallError, InconsistentStateError, VMRunnerError,
 };
 use near_vm_runner::logic::{VMContext, VMOutcome};
-use near_vm_runner::ContractCode;
 use near_vm_runner::{precompile_contract, PreparedContract};
+use near_vm_runner::{ContractCode, ContractRuntimeCache};
 use near_wallet_contract::{wallet_contract, wallet_contract_magic_bytes};
 use std::sync::Arc;
 
@@ -161,44 +160,6 @@ pub(crate) fn execute_function_call(
     Ok(outcome)
 }
 
-pub(crate) fn prepare_function_call(
-    state_update: &TrieUpdate,
-    apply_state: &ApplyState,
-    account: &Account,
-    account_id: &AccountId,
-    function_call: &FunctionCallAction,
-    config: &RuntimeConfig,
-    epoch_info_provider: &(dyn EpochInfoProvider),
-    view_config: Option<ViewConfig>,
-) -> Box<dyn PreparedContract> {
-    let max_gas_burnt = match view_config {
-        Some(ViewConfig { max_gas_burnt }) => max_gas_burnt,
-        None => config.wasm_config.limit_config.max_gas_burnt,
-    };
-    let gas_counter = near_vm_runner::logic::GasCounter::new(
-        config.wasm_config.ext_costs.clone(),
-        max_gas_burnt,
-        config.wasm_config.regular_op_cost,
-        function_call.gas,
-        view_config.is_some(),
-    );
-    let code_ext = RuntimeContractExt {
-        trie_update: state_update,
-        account_id,
-        account,
-        chain_id: &epoch_info_provider.chain_id(),
-        current_protocol_version: apply_state.current_protocol_version,
-    };
-    let contract = near_vm_runner::prepare(
-        &code_ext,
-        Arc::clone(&config.wasm_config),
-        apply_state.cache.as_deref(),
-        gas_counter,
-        &function_call.method_name,
-    );
-    contract
-}
-
 pub(crate) fn action_function_call(
     state_update: &mut TrieUpdate,
     apply_state: &ApplyState,
@@ -645,11 +606,12 @@ pub(crate) fn action_deploy_contract(
     account: &mut Account,
     account_id: &AccountId,
     deploy_contract: &DeployContractAction,
-    apply_state: &ApplyState,
+    config: Arc<near_parameters::vm::Config>,
+    cache: Option<&dyn ContractRuntimeCache>,
 ) -> Result<(), StorageError> {
     let _span = tracing::debug_span!(target: "runtime", "action_deploy_contract").entered();
     let code = ContractCode::new(deploy_contract.code.clone(), None);
-    let prev_code = get_code(state_update, account_id, Some(account.code_hash()))?;
+    let prev_code = state_update.contract_storage.get(account.code_hash());
     let prev_code_length = prev_code.map(|code| code.code().len() as u64).unwrap_or_default();
     account.set_storage_usage(account.storage_usage().saturating_sub(prev_code_length));
     account.set_storage_usage(
@@ -661,16 +623,20 @@ pub(crate) fn action_deploy_contract(
         })?,
     );
     account.set_code_hash(*code.hash());
+    // Legacy: populate the mapping from `AccountId => sha256(code)` thus making contracts part of
+    // The State. For the time being we are also relying on the `TrieUpdate` to actually write the
+    // contracts into the storage as part of the commit routine, however no code should be relying
+    // that the contracts are written to The State.
     set_code(state_update, account_id.clone(), &code);
-    // Precompile the contract and store result (compiled code or error) in the database.
-    // Note, that contract compilation costs are already accounted in deploy cost using
-    // special logic in estimator (see get_runtime_config() function).
-    precompile_contract(
-        &code,
-        Arc::clone(&apply_state.config.wasm_config),
-        apply_state.cache.as_deref(),
-    )
-    .ok();
+    // Precompile the contract and store result (compiled code or error) in the contract runtime
+    // cache.
+    // Note, that contract compilation costs are already accounted in deploy cost using special
+    // logic in estimator (see get_runtime_config() function).
+    precompile_contract(&code, config, cache).ok();
+    // Inform the `store::contract::Storage` about the new deploy (so that the `get` method can
+    // return the contract before the contract is written out to the underlying storage as part of
+    // the `TrieUpdate` commit.)
+    state_update.contract_storage.store(code);
     Ok(())
 }
 
@@ -687,7 +653,7 @@ pub(crate) fn action_delete_account(
     if current_protocol_version >= ProtocolFeature::DeleteActionRestriction.protocol_version() {
         let account = account.as_ref().unwrap();
         let mut account_storage_usage = account.storage_usage();
-        let contract_code = get_code(state_update, account_id, Some(account.code_hash()))?;
+        let contract_code = state_update.contract_storage.get(account.code_hash());
         if let Some(code) = contract_code {
             // account storage usage should be larger than code size
             let code_len = code.code().len() as u64;
@@ -1190,10 +1156,8 @@ mod tests {
     use near_primitives::action::delegate::NonDelegateAction;
     use near_primitives::congestion_info::BlockCongestionInfo;
     use near_primitives::errors::InvalidAccessKeyError;
-    use near_primitives::hash::hash;
     use near_primitives::runtime::migration_data::MigrationFlags;
     use near_primitives::transaction::CreateAccountAction;
-    use near_primitives::trie_key::TrieKey;
     use near_primitives::types::{EpochId, StateChangeCause};
     use near_primitives_core::version::PROTOCOL_VERSION;
     use near_store::set_account;
@@ -1355,11 +1319,25 @@ mod tests {
         let mut state_update =
             tries.new_trie_update(ShardUId::single_shard(), CryptoHash::default());
         let account_id = "alice".parse::<AccountId>().unwrap();
-        let trie_key = TrieKey::ContractCode { account_id: account_id.clone() };
-        let empty_contract = [0; 10_000].to_vec();
-        let contract_hash = hash(&empty_contract);
-        state_update.set(trie_key, empty_contract);
-        test_delete_large_account(&account_id, &contract_hash, storage_usage, &mut state_update)
+        let deploy_action = DeployContractAction { code: [0; 10_000].to_vec() };
+        let mut account =
+            Account::new(100, 0, 0, CryptoHash::default(), storage_usage, PROTOCOL_VERSION);
+        let apply_state = create_apply_state(0);
+        let res = action_deploy_contract(
+            &mut state_update,
+            &mut account,
+            &account_id,
+            &deploy_action,
+            Arc::clone(&apply_state.config.wasm_config),
+            None,
+        );
+        assert!(res.is_ok());
+        test_delete_large_account(
+            &account_id,
+            &account.code_hash(),
+            storage_usage,
+            &mut state_update,
+        )
     }
 
     #[test]

runtime/runtime/src/ext.rs

@@ -6,7 +6,7 @@ use near_primitives::checked_feature;
 use near_primitives::errors::{EpochError, StorageError};
 use near_primitives::hash::CryptoHash;
 use near_primitives::trie_key::{trie_key_parsers, TrieKey};
-use near_primitives::types::{AccountId, Balance, EpochId, EpochInfoProvider, Gas, TrieCacheMode};
+use near_primitives::types::{AccountId, Balance, EpochId, EpochInfoProvider, Gas};
 use near_primitives::utils::create_receipt_id_from_action_hash;
 use near_primitives::version::ProtocolVersion;
 use near_store::{has_promise_yield_receipt, KeyLookupMode, TrieUpdate, TrieUpdateValuePtr};
@@ -362,34 +362,28 @@ impl<'a> External for RuntimeExt<'a> {
 }
 
 pub(crate) struct RuntimeContractExt<'a> {
-    pub(crate) trie_update: &'a TrieUpdate,
+    pub(crate) storage: near_store::contract::Storage,
     pub(crate) account_id: &'a AccountId,
-    pub(crate) account: &'a Account,
+    pub(crate) code_hash: CryptoHash,
     pub(crate) chain_id: &'a str,
     pub(crate) current_protocol_version: ProtocolVersion,
 }
 
 impl<'a> Contract for RuntimeContractExt<'a> {
     fn hash(&self) -> CryptoHash {
-        self.account.code_hash()
+        self.code_hash
     }
 
     fn get_code(&self) -> Option<Arc<ContractCode>> {
         let account_id = self.account_id;
         let code_hash = self.hash();
-        let version = self.current_protocol_version;
         let chain_id = self.chain_id;
         if checked_feature!("stable", EthImplicitAccounts, self.current_protocol_version)
             && account_id.get_account_type() == AccountType::EthImplicitAccount
             && &code_hash == wallet_contract_magic_bytes(&chain_id).hash()
         {
             return Some(wallet_contract(&chain_id));
         }
-        let mode = match checked_feature!("stable", ChunkNodesCache, version) {
-            true => Some(TrieCacheMode::CachingShard),
-            false => None,
-        };
-        let _guard = self.trie_update.with_trie_cache_mode(mode);
-        self.trie_update.get_code(self.account_id.clone(), code_hash).map(Arc::new)
+        self.storage.get(code_hash).map(Arc::new)
     }
 }