feat: use SetLoadedAccountsDataSizeLimit for tx containing DLP ix

Author: snawaz

Cargo.lock

@@ -98,7 +98,8 @@ 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 = "2e52f858d86ec7420d187368cb1fd1edfccb9e2b", features = [
+# magicblock-delegation-program = { path = "../delegation-program", features = [
   "no-entrypoint",
 ] }
 magicblock-ledger = { path = "./magicblock-ledger" }

Cargo.toml

@@ -1,8 +1,14 @@
 use dlp::{
     args::{CallHandlerArgs, CommitDiffArgs, CommitStateArgs},
     compute_diff,
+    instruction_builder::{
+        commit_diff_from_buffer_size_budget, commit_diff_size_budget,
+        commit_size_budget, finalize_size_budget, undelegate_size_budget,
+    },
+    AccountSizeClass,
 };
 use magicblock_metrics::metrics::LabelValue;
+use rand::seq::IndexedRandom;
 use solana_account::ReadableAccount;
 use solana_instruction::{AccountMeta, Instruction};
 use solana_pubkey::Pubkey;
@@ -177,6 +183,28 @@ 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(_) => 0,
+            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;
@@ -128,6 +131,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

@@ -85,6 +85,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

@@ -224,8 +224,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(),
@@ -258,8 +262,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

@@ -66,6 +66,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 +128,28 @@ impl TransactionUtils {
         tasks.iter().map(|task| task.as_ref().compute_units()).sum()
     }
 
+    pub fn tasks_accounts_size_budget(
+        tasks: &[impl AsRef<dyn BaseTask>],
+    ) -> u32 {
+        tasks
+            .iter()
+            .map(|task| task.as_ref().accounts_size_budget())
+            .sum()
+    }
+
     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

@@ -101,17 +101,17 @@ async fn test_ix_commit_order_book_change_100_bytes() {
 }
 
 #[tokio::test]
-async fn test_ix_commit_order_book_change_679_bytes() {
-    commit_book_order_account(679, CommitStrategy::Args, false).await;
+async fn test_ix_commit_order_book_change_671_bytes() {
+    commit_book_order_account(671, CommitStrategy::Args, false).await;
 }
 
 #[tokio::test]
-async fn test_ix_commit_order_book_change_680_bytes() {
-    // We cannot use 680 as changed_len because that both 679 and 680 produce encoded tx
-    // of size 1644 (which is the max limit), but while the size of raw bytes for 679 is within
-    // 1232 limit, the size for 680 execeds by 1 (1233). That is why we used
-    // 681 as changed_len where CommitStrategy goes from Args to FromBuffer.
-    commit_book_order_account(681, CommitStrategy::FromBuffer, false).await;
+async fn test_ix_commit_order_book_change_673_bytes() {
+    // We cannot use 672 as changed_len because that both 671 and 672 produce encoded tx
+    // of size 1644 (which is the max limit), but while the size of raw bytes for 671 is within
+    // 1232 limit, the size for 672 execeds by 1 (1233). That is why we used
+    // 673 as changed_len where CommitStrategy goes from Args to FromBuffer.
+    commit_book_order_account(673, CommitStrategy::FromBuffer, false).await;
 }
 
 #[tokio::test]