Fetch base_accounts and commit_ids in parallel

Author: snawaz

magicblock-committor-service/src/intent_executor/mod.rs

@@ -6,9 +6,6 @@ pub mod two_stage_executor;
 
 use std::{mem, ops::ControlFlow, sync::Arc, time::Duration};
 
-#[cfg(test)]
-mod null_task_info_fetcher;
-
 use async_trait::async_trait;
 use futures_util::future::{join, try_join_all};
 use log::{trace, warn};
@@ -25,8 +22,6 @@ use magicblock_rpc_client::{
     MagicBlockRpcClientError, MagicBlockSendTransactionConfig,
     MagicBlockSendTransactionOutcome, MagicblockRpcClient,
 };
-#[cfg(test)]
-pub use null_task_info_fetcher::*;
 use solana_keypair::Keypair;
 use solana_message::VersionedMessage;
 use solana_pubkey::Pubkey;

magicblock-committor-service/src/intent_executor/null_task_info_fetcher.rs

@@ -9,9 +9,7 @@ use dlp::{
 use log::warn;
 use lru::LruCache;
 use magicblock_metrics::metrics;
-use magicblock_rpc_client::{
-    MagicBlockRpcClientError, MagicBlockRpcClientResult, MagicblockRpcClient,
-};
+use magicblock_rpc_client::{MagicBlockRpcClientError, MagicblockRpcClient};
 use solana_account::Account;
 use solana_pubkey::Pubkey;
 use solana_signature::Signature;
@@ -40,10 +38,10 @@ pub trait TaskInfoFetcher: Send + Sync + 'static {
     /// Resets cache for some or all accounts
     fn reset(&self, reset_type: ResetType);
 
-    async fn get_base_account(
+    async fn get_base_accounts(
         &self,
-        _pubkey: &Pubkey,
-    ) -> MagicBlockRpcClientResult<Option<Account>>;
+        pubkeys: &[Pubkey],
+    ) -> TaskInfoFetcherResult<HashMap<Pubkey, Account>>;
 }
 
 pub enum ResetType<'a> {
@@ -272,11 +270,23 @@ impl TaskInfoFetcher for CacheTaskInfoFetcher {
         }
     }
 
-    async fn get_base_account(
+    async fn get_base_accounts(
         &self,
-        pubkey: &Pubkey,
-    ) -> MagicBlockRpcClientResult<Option<Account>> {
-        self.rpc_client.get_account(pubkey).await
+        pubkeys: &[Pubkey],
+    ) -> TaskInfoFetcherResult<HashMap<Pubkey, Account>> {
+        self.rpc_client
+            .get_multiple_accounts(pubkeys, None)
+            .await
+            .map_err(|err| {
+                TaskInfoFetcherError::MagicBlockRpcClientError(Box::new(err))
+            })
+            .map(|accounts| {
+                pubkeys
+                    .iter()
+                    .zip(accounts.into_iter())
+                    .filter_map(|(key, value)| value.map(|value| (*key, value)))
+                    .collect()
+            })
     }
 }
 

magicblock-committor-service/src/intent_executor/task_info_fetcher.rs

@@ -1,12 +1,12 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use futures_util::future::join_all;
 use log::error;
 use magicblock_program::magic_scheduled_base_intent::{
     CommitType, CommittedAccount, MagicBaseIntent, ScheduledBaseIntent,
     UndelegateType,
 };
+use solana_account::Account;
 use solana_pubkey::Pubkey;
 use solana_signature::Signature;
 
@@ -51,31 +51,18 @@ pub struct TaskBuilderImpl;
 pub const COMMIT_STATE_SIZE_THRESHOLD: usize = 256;
 
 impl TaskBuilderImpl {
-    pub async fn create_commit_task<C: TaskInfoFetcher>(
+    pub fn create_commit_task(
         commit_id: u64,
         allow_undelegation: bool,
         account: CommittedAccount,
-        task_info_fetcher: &Arc<C>,
+        base_account: Option<Account>,
     ) -> ArgsTask {
-        let base_account = if account.account.data.len()
-            > COMMIT_STATE_SIZE_THRESHOLD
-        {
-            match task_info_fetcher.get_base_account(&account.pubkey).await {
-                Ok(Some(account)) => Some(account),
-                Ok(None) => {
-                    log::warn!("AccountNotFound for commit_diff, pubkey: {}, commit_id: {}, Falling back to commit_state.",
-                        account.pubkey, commit_id);
-                    None
-                }
-                Err(e) => {
-                    log::warn!("Failed to fetch base account for commit diff, pubkey: {}, commit_id: {}, error: {}. Falling back to commit_state.",
-                        account.pubkey, commit_id, e);
-                    None
-                }
-            }
-        } else {
-            None
-        };
+        let base_account =
+            if account.account.data.len() > COMMIT_STATE_SIZE_THRESHOLD {
+                base_account
+            } else {
+                None
+            };
 
         if let Some(base_account) = base_account {
             ArgsTaskType::CommitDiff(CommitDiffTask {
@@ -123,14 +110,37 @@ impl TasksBuilder for TaskBuilderImpl {
             }
         };
 
-        let committed_pubkeys = accounts
-            .iter()
-            .map(|account| account.pubkey)
-            .collect::<Vec<_>>();
-        let commit_ids = commit_id_fetcher
-            .fetch_next_commit_ids(&committed_pubkeys)
-            .await
-            .map_err(TaskBuilderError::CommitTasksBuildError)?;
+        let (commit_ids, base_accounts) = {
+            let committed_pubkeys = accounts
+                .iter()
+                .map(|account| account.pubkey)
+                .collect::<Vec<_>>();
+
+            let diffable_pubkeys = accounts
+                .iter()
+                .filter(|account| {
+                    account.account.data.len() >= COMMIT_STATE_SIZE_THRESHOLD
+                })
+                .map(|account| account.pubkey)
+                .collect::<Vec<_>>();
+
+            tokio::join!(
+                commit_id_fetcher.fetch_next_commit_ids(&committed_pubkeys),
+                commit_id_fetcher
+                    .get_base_accounts(diffable_pubkeys.as_slice())
+            )
+        };
+
+        let commit_ids =
+            commit_ids.map_err(TaskBuilderError::CommitTasksBuildError)?;
+
+        let base_accounts = match base_accounts {
+            Ok(map) => map,
+            Err(err) => {
+                log::warn!("Failed to fetch base accounts for CommitDiff (id={}): {}; falling back to CommitState", base_intent.id, err);
+                Default::default()
+            }
+        };
 
         // Persist commit ids for commitees
         commit_ids
@@ -141,13 +151,17 @@ impl TasksBuilder for TaskBuilderImpl {
                 }
             });
 
-        let tasks = join_all(accounts
+        let tasks = accounts
             .iter()
-            .map(|account| async {
+            .map(|account| {
                 let commit_id = *commit_ids.get(&account.pubkey).expect("CommitIdFetcher provide commit ids for all listed pubkeys, or errors!");
-                let task = Self::create_commit_task(commit_id, allow_undelegation, account.clone(), commit_id_fetcher).await;
+                // TODO (snawaz): if accounts do not have duplicate, then we can use remove
+                // instead:
+                //  let base_account = base_accounts.remove(&account.pubkey);
+                let base_account = base_accounts.get(&account.pubkey).cloned();
+                let task = Self::create_commit_task(commit_id, allow_undelegation, account.clone(), base_account);
                 Box::new(task) as Box<dyn BaseTask>
-            })).await;
+            }).collect();
 
         Ok(tasks)
     }