prefetch: claim.sweat::claim

near · Apr 17, 2024 · 8a9786c · 8a9786c
1 parent 4d506a7
commit 8a9786c
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 3 deletions.
diff --git a/core/store/src/config.rs b/core/store/src/config.rs
@@ -259,6 +259,10 @@ impl Default for StoreConfig {
                 receiver: "claim.sweat".to_owned(),
                 sender: "token.sweat".to_owned(),
                 method_name: "record_batch_for_hold".to_owned(),
+            }, PrefetchConfig {
+                receiver: "claim.sweat".to_owned(),
+                sender: "token.sweat".to_owned(),
+                method_name: "claim".to_owned(),
             }],
             kaiching_prefetch_config: vec![PrefetchConfig {
                 receiver: "earn.kaiching".to_owned(),

diff --git a/core/store/src/trie/prefetching_trie_storage.rs b/core/store/src/trie/prefetching_trie_storage.rs
@@ -87,6 +87,7 @@ pub struct PrefetchApi {
     pub kaiching_prefetch_config: Vec<PrefetchConfig>,
 
     pub shard_uid: ShardUId,
+    pub store: Store,
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -421,6 +422,7 @@ impl PrefetchApi {
             claim_sweat_prefetch_config,
             kaiching_prefetch_config,
             shard_uid,
+            store,
         };
         let (shutdown_tx, shutdown_rx) = crossbeam::channel::bounded(1);
         let handles = (0..NUM_IO_THREADS)

diff --git a/runtime/runtime/src/prefetch.rs b/runtime/runtime/src/prefetch.rs
@@ -158,7 +158,18 @@ impl TriePrefetcher {
                         && cfg.receiver == account_id.as_str()
                         && cfg.method_name == fn_call.method_name
                 }) {
-                    self.prefetch_claim_sweat(account_id.clone(), &fn_call.args)?;
+                    match &*fn_call.method_name {
+                        "record_batch_for_hold" => self
+                            .prefetch_claim_sweat_record_batch_for_hold(
+                                account_id.clone(),
+                                &fn_call.args,
+                            )?,
+                        "claim" => self.prefetch_claim_sweat_claim(
+                            account_id.clone(),
+                            receipt.predecessor_id.clone(),
+                        )?,
+                        _ => {}
+                    }
                 }
 
                 if self.prefetch_api.kaiching_prefetch_config.iter().any(|cfg| {
@@ -274,10 +285,14 @@ impl TriePrefetcher {
         Ok(())
     }
 
-    /// Prefetcher tuned for claim.sweat contract calls.
+    /// Prefetcher tuned for claim.sweat::record_batch_for_hold contract calls.
     ///
     /// Remove after #10965 reaches mainnet.
-    fn prefetch_claim_sweat(&self, account_id: AccountId, arg: &[u8]) -> Result<(), PrefetchError> {
+    fn prefetch_claim_sweat_record_batch_for_hold(
+        &self,
+        account_id: AccountId,
+        arg: &[u8],
+    ) -> Result<(), PrefetchError> {
         let Ok(json) = serde_json::de::from_slice::<serde_json::Value>(arg) else {
             return Ok(());
         };
@@ -304,6 +319,51 @@ impl TriePrefetcher {
         Ok(())
     }
 
+    /// Prefetcher tuned for claim.sweat::claim contract calls.
+    ///
+    /// Remove after #10965 reaches mainnet.
+    fn prefetch_claim_sweat_claim(
+        &self,
+        account_id: AccountId,
+        predecessor: AccountId,
+    ) -> Result<(), PrefetchError> {
+        rayon::spawn(|| {
+            let account_data_key = Vec::with_capacity(4 + 8 + predecessor.len());
+            account_data_key.extend(borsh::to_vec(&0u8)); // StorageKey::Accounts
+            account_data_key.extend(predecessor.as_bytes());
+            let trie_key = TrieKey::ContractData { account_id: account_id.clone(), key };
+            // Just read this directly for now since this is temporary anyway
+            let trie = Trie::new(self.prefetch_api.store.clone(), self.trie_root, None);
+            let Some(account_record) = trie.get(&account_data_key.to_vec()) else { return };
+            #[derive(BorshDeserialize, BorshSerialize)]
+            pub struct AccountRecord {
+                pub accruals: Vec<(u32, u32)>,
+                // TODO maybe need the rest of the fields? to be found out.
+            }
+            let Ok(account_record) = borsh::from_slice::<AccountRecord>(account_record) else {
+                return;
+            };
+            for (dt, idx) in account_record.accruals {
+                let accruals_key = Vec::with_capacity(4 + 8);
+                accruals_key.extend(borsh::to_vec(&1u8)); // StorageKey::Accruals
+                accruals_key.extend(borsh::to_vec(&dt));
+                self.prefetch_trie_key(TrieKey::ContractData {
+                    account_id: account_id.clone(),
+                    key: accruals_key,
+                });
+                let amount_key = Vec::with_capacity(4 + 8 + 8);
+                amount_key.extend(borsh::to_vec(&2u8)); // StorageKey::AccrualsEntry as u8
+                amount_key.extend(borsh::to_vec(&dt)); // StorageKey::AccrualsEntry.0
+                amount_key.extend(&idx.to_le_bytes()); // index into Vector
+                self.prefetch_trie_key(TrieKey::ContractData {
+                    account_id: account_id.clone(),
+                    key: amount_key,
+                });
+            }
+        });
+        Ok(())
+    }
+
     /// Prefetcher tuned for kaiching contract calls.
     ///
     /// Remove after #10965 reaches mainnet.