contract_preparation: load account lazily (near#12574)

In workloads where we don't have any function calls these `Account` reads can add up to significant execution time cost. Fixes near#12573
kobayurii · Dec 8, 2024 · 3270551 · 3270551
1 parent 5196eb9
commit 3270551
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 11 deletions.
diff --git a/runtime/runtime/src/lib.rs b/runtime/runtime/src/lib.rs
@@ -2593,21 +2593,24 @@ fn schedule_contract_preparation<'b, R: MaybeRefReceipt>(
     let scheduled_receipt_offset = iterator.position(|peek| {
         let peek = peek.as_ref();
         let account_id = peek.receiver_id();
-        let key = TrieKey::Account { account_id: account_id.clone() };
-        let receiver = get_pure::<Account>(state_update, &key);
-        let Ok(Some(receiver)) = receiver else {
-            // Most likely reason this can happen is because the receipt is for an account that
-            // does not yet exist. This is a routine occurrence as accounts are created by sending
-            // some NEAR to a name that's about to be created.
-            return false;
-        };
+        let receiver = std::cell::LazyCell::new(|| {
+            let key = TrieKey::Account { account_id: account_id.clone() };
+            let receiver = get_pure::<Account>(state_update, &key);
+            let Ok(Some(receiver)) = receiver else {
+                // Most likely reason this can happen is because the receipt is for an account that
+                // does not yet exist. This is a routine occurrence as accounts are created by
+                // sending some NEAR to a name that's about to be created.
+                return None;
+            };
+            Some(receiver)
+        });
 
         // We need to inspect each receipt recursively in case these are data receipts, thus a
         // function.
         fn handle_receipt(
             mgr: &mut ReceiptPreparationPipeline,
             state_update: &TrieUpdate,
-            receiver: &Account,
+            receiver: &std::cell::LazyCell<Option<Account>, impl FnOnce() -> Option<Account>>,
             account_id: &AccountId,
             receipt: &Receipt,
         ) -> bool {

diff --git a/runtime/runtime/src/pipelining.rs b/runtime/runtime/src/pipelining.rs
@@ -111,7 +111,7 @@ impl ReceiptPreparationPipeline {
     pub(crate) fn submit(
         &mut self,
         receipt: &Receipt,
-        account: &Account,
+        account: &std::cell::LazyCell<Option<Account>, impl FnOnce() -> Option<Account>>,
         view_config: Option<ViewConfig>,
     ) -> bool {
         let account_id = receipt.receiver_id();
@@ -133,6 +133,8 @@ impl ReceiptPreparationPipeline {
                     return self.block_accounts.insert(account_id);
                 }
                 Action::FunctionCall(function_call) => {
+                    let Some(account) = &**account else { continue };
+                    let code_hash = account.code_hash();
                     let key = PrepareTaskKey { receipt_id: receipt.get_hash(), action_index };
                     let gas_counter = self.gas_counter(view_config.as_ref(), function_call.gas);
                     let entry = match self.map.entry(key) {
@@ -145,7 +147,6 @@ impl ReceiptPreparationPipeline {
                     let cache = self.contract_cache.as_ref().map(|c| c.handle());
                     let storage = self.storage.clone();
                     let protocol_version = self.protocol_version;
-                    let code_hash = account.code_hash();
                     let created = Instant::now();
                     let method_name = function_call.method_name.clone();
                     let status = Mutex::new(PrepareTaskStatus::Pending);