feat: use SetLoadedAccountsDataSizeLimit for tx containing DLP ix (#800)

Author: snawaz

Cargo.lock

@@ -106,7 +106,7 @@ magicblock-committor-program = { path = "./magicblock-committor-program", featur
 magicblock-committor-service = { path = "./magicblock-committor-service" }
 magicblock-config = { path = "./magicblock-config" }
 magicblock-core = { path = "./magicblock-core" }
-magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "3edb41022", features = [
+magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "1874b4f5f5f55cb9ab54b64de2cc0d41107d1435", features = [
   "no-entrypoint",
 ] }
 magicblock-ledger = { path = "./magicblock-ledger" }

Cargo.toml

@@ -179,8 +179,10 @@ impl TransactionStrategyExecutionError {
         signature: Option<Signature>,
         tasks: &[Box<dyn BaseTask>],
     ) -> Result<Self, TransactionError> {
-        // There's always 2 budget instructions in front
-        const OFFSET: u8 = 2;
+        // There's always 3 budget instructions in front
+        // TODO (snawaz): this is offset-sensitive, if we add or remove any instruction from the
+        // front, that leads to incorrect error reporting. so if possible, make it offset-insensitive.
+        const OFFSET: u8 = 3;
         const NONCE_OUT_OF_ORDER: u32 =
             dlp::error::DlpError::NonceOutOfOrder as u32;
 

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

@@ -2,6 +2,11 @@ use compressed_delegation_client::types::{CommitArgs, FinalizeArgs};
 use dlp::{
     args::{CallHandlerArgs, CommitDiffArgs, CommitStateArgs},
     compute_diff,
+    instruction_builder::{
+        call_handler_size_budget, commit_diff_size_budget, commit_size_budget,
+        finalize_size_budget, undelegate_size_budget,
+    },
+    AccountSizeClass,
 };
 use magicblock_metrics::metrics::LabelValue;
 use solana_account::ReadableAccount;
@@ -251,6 +256,38 @@ impl BaseTask for ArgsTask {
         }
     }
 
+    fn accounts_size_budget(&self) -> u32 {
+        match &self.task_type {
+            ArgsTaskType::Commit(task) => {
+                commit_size_budget(AccountSizeClass::Dynamic(
+                    task.committed_account.account.data.len() as u32,
+                ))
+            }
+            ArgsTaskType::CommitDiff(task) => {
+                commit_diff_size_budget(AccountSizeClass::Dynamic(
+                    task.committed_account.account.data.len() as u32,
+                ))
+            }
+            ArgsTaskType::BaseAction(task) => {
+                // assume all other accounts are Small accounts.
+                let other_accounts_budget =
+                    task.action.account_metas_per_program.len() as u32
+                        * AccountSizeClass::Small.size_budget();
+
+                call_handler_size_budget(
+                    AccountSizeClass::Medium,
+                    other_accounts_budget,
+                )
+            }
+            ArgsTaskType::Undelegate(_) => {
+                undelegate_size_budget(AccountSizeClass::Huge)
+            }
+            ArgsTaskType::Finalize(_) => {
+                finalize_size_budget(AccountSizeClass::Huge)
+            }
+        }
+    }
+
     #[cfg(test)]
     fn strategy(&self) -> TaskStrategy {
         TaskStrategy::Args

magicblock-committor-service/src/tasks/args_task.rs

@@ -1,4 +1,7 @@
-use dlp::args::CommitStateFromBufferArgs;
+use dlp::{
+    args::CommitStateFromBufferArgs, instruction_builder::commit_size_budget,
+    AccountSizeClass,
+};
 use magicblock_committor_program::Chunks;
 use magicblock_metrics::metrics::LabelValue;
 use solana_instruction::Instruction;
@@ -131,6 +134,14 @@ impl BaseTask for BufferTask {
         }
     }
 
+    fn accounts_size_budget(&self) -> u32 {
+        match self.task_type {
+            BufferTaskType::Commit(_) => {
+                commit_size_budget(AccountSizeClass::Huge)
+            }
+        }
+    }
+
     #[cfg(test)]
     fn strategy(&self) -> TaskStrategy {
         TaskStrategy::Buffer

magicblock-committor-service/src/tasks/buffer_task.rs

@@ -89,6 +89,9 @@ pub trait BaseTask: Send + Sync + DynClone + LabelValue {
     /// Returns [`Task`] budget
     fn compute_units(&self) -> u32;
 
+    /// Returns the max accounts-data-size that can be used with SetLoadedAccountsDataSizeLimit
+    fn accounts_size_budget(&self) -> u32;
+
     /// Returns current [`TaskStrategy`]
     #[cfg(test)]
     fn strategy(&self) -> TaskStrategy;

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

@@ -255,8 +255,12 @@ impl TaskStrategist {
         let placeholder = Keypair::new();
         // Gather all involved keys in tx
         let budgets = TransactionUtils::tasks_compute_units(tasks);
-        let budget_instructions =
-            TransactionUtils::budget_instructions(budgets, u64::default());
+        let size_budgets = TransactionUtils::tasks_accounts_size_budget(tasks);
+        let budget_instructions = TransactionUtils::budget_instructions(
+            budgets,
+            u64::default(),
+            size_budgets,
+        );
         let unique_involved_pubkeys = TransactionUtils::unique_involved_pubkeys(
             tasks,
             &placeholder.pubkey(),
@@ -289,8 +293,12 @@ impl TaskStrategist {
         tasks: &[Box<dyn BaseTask>],
     ) -> Vec<Pubkey> {
         let budgets = TransactionUtils::tasks_compute_units(tasks);
-        let budget_instructions =
-            TransactionUtils::budget_instructions(budgets, u64::default());
+        let size_budgets = TransactionUtils::tasks_accounts_size_budget(tasks);
+        let budget_instructions = TransactionUtils::budget_instructions(
+            budgets,
+            u64::default(),
+            size_budgets,
+        );
 
         TransactionUtils::unique_involved_pubkeys(
             tasks,

magicblock-committor-service/src/tasks/task_strategist.rs

@@ -1,5 +1,6 @@
 use std::collections::HashSet;
 
+use dlp::DLP_PROGRAM_DATA_SIZE_CLASS;
 use solana_compute_budget_interface::ComputeBudgetInstruction;
 use solana_hash::Hash;
 use solana_instruction::Instruction;
@@ -66,6 +67,7 @@ impl TransactionUtils {
         let budget_instructions = Self::budget_instructions(
             Self::tasks_compute_units(tasks),
             compute_unit_price,
+            Self::tasks_accounts_size_budget(tasks),
         );
         let ixs = Self::tasks_instructions(&authority.pubkey(), tasks);
         Self::assemble_tx_raw(
@@ -127,16 +129,38 @@ impl TransactionUtils {
         tasks.iter().map(|task| task.as_ref().compute_units()).sum()
     }
 
+    pub fn tasks_accounts_size_budget(
+        tasks: &[impl AsRef<dyn BaseTask>],
+    ) -> u32 {
+        if tasks.is_empty() {
+            return 0;
+        }
+
+        let total_budget: u32 = tasks
+            .iter()
+            .map(|task| task.as_ref().accounts_size_budget())
+            .sum();
+
+        // DLP_PROGRAM_DATA_SIZE_CLASS has been added N times, once for each task.
+        // We need to add it once only, so minus (N-1) times.
+        total_budget
+            - (tasks.len() as u32 - 1)
+                * DLP_PROGRAM_DATA_SIZE_CLASS.size_budget()
+    }
+
     pub fn budget_instructions(
         compute_units: u32,
         compute_unit_price: u64,
-    ) -> [Instruction; 2] {
-        let compute_budget_ix =
-            ComputeBudgetInstruction::set_compute_unit_limit(compute_units);
-        let compute_unit_price_ix =
+        accounts_size_budget: u32,
+    ) -> [Instruction; 3] {
+        [
+            ComputeBudgetInstruction::set_compute_unit_limit(compute_units),
             ComputeBudgetInstruction::set_compute_unit_price(
                 compute_unit_price,
-            );
-        [compute_budget_ix, compute_unit_price_ix]
+            ),
+            ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(
+                accounts_size_budget,
+            ),
+        ]
     }
 }

magicblock-committor-service/src/tasks/utils.rs

@@ -68,7 +68,7 @@ magic-domain-program = { git = "https://github.com/magicblock-labs/magic-domain-
   "modular-sdk",
 ] }
 magicblock_magic_program_api = { package = "magicblock-magic-program-api", path = "../magicblock-magic-program-api" }
-magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "3edb41022", features = [
+magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "1874b4f5f5f55cb9ab54b64de2cc0d41107d1435", features = [
   "no-entrypoint",
 ] }
 magicblock-program = { path = "../programs/magicblock" }