make near_store::contract::Storage public

As a somewhat notable shift in the approach, we'll start focusing on
contract wasm code management by directly addressing it by the code
hash, rather than going through The State (Trie).

One notable snag here is that the transaction runtime doesn't actually
have a write access to the underlying storage, so we continue relying on
the `TrieUpdate` to write these contracts out. In the meanwhile we
maintain an in-memory map of deployed contracts to make the storage
appear consistent.

core/store/src/lib.rs

@@ -55,6 +55,7 @@ mod rocksdb_metrics;
 mod sync_utils;
 pub mod test_utils;
 pub mod trie;
+pub mod contract;
 
 pub use crate::config::{Mode, StoreConfig};
 pub use crate::opener::{
@@ -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/sanity_checks.rs

@@ -115,8 +115,8 @@ fn test_cost_sanity() {
     let res = node
         .user()
         .function_call(
-            alice_account(),
-            test_contract_account(),
+            dbg!(alice_account()),
+            dbg!(test_contract_account()),
             "sanity_check",
             args.into_bytes(),
             MAX_GAS,

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::{ContractStorage, ExternalError, RuntimeContractExt, RuntimeExt};
+use crate::ext::{ExternalError, RuntimeContractExt, RuntimeExt};
 use crate::receipt_manager::ReceiptManager;
 use crate::{metrics, ActionResult, ApplyState};
 use near_crypto::PublicKey;
@@ -30,15 +30,14 @@ 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::{ContractCode, ContractRuntimeCache};
 use near_vm_runner::{precompile_contract, PreparedContract};
 use near_wallet_contract::{wallet_contract, wallet_contract_magic_bytes};
 use std::sync::Arc;
@@ -183,7 +182,7 @@ pub(crate) fn prepare_function_call(
         view_config.is_some(),
     );
     let code_ext = RuntimeContractExt {
-        storage: ContractStorage::Transaction(state_update),
+        storage: state_update.contract_storage.clone(),
         account_id,
         account,
         chain_id: &epoch_info_provider.chain_id(),
@@ -645,11 +644,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 +661,25 @@ 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
+    // on this per se.
     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 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,
-        Arc::clone(&apply_state.config.wasm_config),
-        apply_state.cache.as_deref(),
+        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 +696,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 +1199,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 +1362,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,12 +6,10 @@ 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, TrieStorage, TrieUpdate, TrieUpdateValuePtr,
-};
+use near_store::{has_promise_yield_receipt, KeyLookupMode, TrieUpdate, TrieUpdateValuePtr};
 use near_vm_runner::logic::errors::{AnyError, VMLogicError};
 use near_vm_runner::logic::types::ReceiptIndex;
 use near_vm_runner::logic::{External, StorageGetMode, ValuePtr};
@@ -363,19 +361,8 @@ impl<'a> External for RuntimeExt<'a> {
     }
 }
 
-pub(crate) enum ContractStorage<'a> {
-    /// Access the transaction first.
-    Transaction(&'a TrieUpdate),
-    /// Access the storage layer directly, forgoing the transaction.
-    ///
-    /// When using this, care must be taken to not accidentally access outdated contract code (i.e.
-    /// old code for an account that has deployed a contract but for which the code has not been
-    /// committed yet.)
-    DB(&'a dyn TrieStorage),
-}
-
 pub(crate) struct RuntimeContractExt<'a> {
-    pub(crate) storage: ContractStorage<'a>,
+    pub(crate) storage: near_store::contract::Storage,
     pub(crate) account_id: &'a AccountId,
     pub(crate) account: &'a Account,
     pub(crate) chain_id: &'a str,
@@ -390,29 +377,13 @@ impl<'a> Contract for RuntimeContractExt<'a> {
     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,
-        };
-        match self.storage {
-            ContractStorage::Transaction(trie_update) => {
-                let _guard = trie_update.with_trie_cache_mode(mode);
-                trie_update.get_code(self.account_id.clone(), code_hash).map(Arc::new)
-            }
-            ContractStorage::DB(db) => match db.retrieve_raw_bytes(&code_hash) {
-                Ok(raw_code) => {
-                    Some(Arc::new(ContractCode::new(raw_code.to_vec(), Some(code_hash))))
-                }
-                Err(_) => None,
-            },
-        }
+        self.storage.get(code_hash).map(Arc::new)
     }
 }

runtime/runtime/src/lib.rs

@@ -431,7 +431,8 @@ impl Runtime {
                     account.as_mut().expect(EXPECT_ACCOUNT_EXISTS),
                     account_id,
                     deploy_contract,
-                    apply_state,
+                    Arc::clone(&apply_state.config.wasm_config),
+                    apply_state.cache.as_deref(),
                 )?;
             }
             Action::FunctionCall(function_call) => {
@@ -2296,8 +2297,7 @@ impl<'a> ApplyProcessingState<'a> {
             self.apply_state.cache.as_ref().map(|v| v.handle()),
             self.epoch_info_provider.chain_id(),
             self.apply_state.current_protocol_version,
-            todo!(),
-            // self.state_update.trie.storage, // TODO??
+            self.state_update.contract_storage.clone(),
         );
         ApplyProcessingReceiptState {
             pipeline_manager,
@@ -2334,14 +2334,15 @@ struct ApplyProcessingReceiptState<'a> {
     local_receipts: VecDeque<Receipt>,
     incoming_receipts: &'a [Receipt],
     delayed_receipts: DelayedReceiptQueueWrapper,
-
     pipeline_manager: pipelining::ReceiptPreparationPipeline,
 }
 
 impl<'a> ApplyProcessingReceiptState<'a> {
     /// Obtain the next receipt that should be executed.
     fn next_local_receipt(&mut self) -> Option<Receipt> {
         let receipt = self.local_receipts.pop_front()?;
+        #[allow(dropping_references)]
+        drop(&self.pipeline_manager);
         Some(receipt)
     }
 }