feat: Execute CommitDiff as BufferTask

snawaz · snawaz · commit eff43ea8c073 · 2025-12-19T22:14:41.000+05:30
diff --git a/magicblock-committor-service/src/tasks/args_task.rs b/magicblock-committor-service/src/tasks/args_task.rs
@@ -89,21 +89,9 @@ impl BaseTask for ArgsTask {
         self: Box<Self>,
     ) -> Result<Box<dyn BaseTask>, Box<dyn BaseTask>> {
         match self.task_type {
-            ArgsTaskType::Commit(mut value) if value.is_commit_diff() => {
-                // TODO (snawaz): Currently, we do not support executing CommitDiff
-                // as BufferTask, which is why we're forcing CommitTask to use CommitState
-                // before converting this task into BufferTask. Once CommitDiff is supported
-                // by BufferTask, we do not have to force_commit_state and we can remove
-                // force_commit_state stuff, as it's essentially a downgrade.
-
-                value.force_commit_state();
-                Ok(Box::new(BufferTask::new_preparation_required(
-                    BufferTaskType::Commit(value),
-                )))
-            }
             ArgsTaskType::Commit(value) => {
                 Ok(Box::new(BufferTask::new_preparation_required(
-                    BufferTaskType::Commit(value),
+                    BufferTaskType::Commit(value.switch_to_buffer_strategy()),
                 )))
             }
             ArgsTaskType::BaseAction(_)
diff --git a/magicblock-committor-service/src/tasks/buffer_task.rs b/magicblock-committor-service/src/tasks/buffer_task.rs
@@ -1,4 +1,3 @@
-use dlp::args::CommitStateFromBufferArgs;
 use magicblock_committor_program::Chunks;
 use magicblock_metrics::metrics::LabelValue;
 use solana_instruction::Instruction;
@@ -47,16 +46,18 @@ impl BufferTask {
 
     fn preparation_required(task_type: &BufferTaskType) -> PreparationState {
         let BufferTaskType::Commit(ref commit_task) = task_type;
-        let committed_data = commit_task.committed_account.account.data.clone();
-        let chunks = Chunks::from_data_length(
-            committed_data.len(),
-            MAX_WRITE_CHUNK_SIZE,
-        );
+        let state_or_diff = if let Some(diff) = commit_task.compute_diff() {
+            diff.to_vec()
+        } else {
+            commit_task.committed_account.account.data.clone()
+        };
+        let chunks =
+            Chunks::from_data_length(state_or_diff.len(), MAX_WRITE_CHUNK_SIZE);
 
         PreparationState::Required(PreparationTask {
             commit_id: commit_task.commit_id,
             pubkey: commit_task.committed_account.pubkey,
-            committed_data,
+            committed_data: state_or_diff,
             chunks,
         })
     }
@@ -65,31 +66,15 @@ impl BufferTask {
 impl BaseTask for BufferTask {
     fn instruction(&self, validator: &Pubkey) -> Instruction {
         let BufferTaskType::Commit(ref value) = self.task_type;
-        let commit_id_slice = value.commit_id.to_le_bytes();
-        let (commit_buffer_pubkey, _) =
-            magicblock_committor_program::pdas::buffer_pda(
-                validator,
-                &value.committed_account.pubkey,
-                &commit_id_slice,
-            );
-
-        dlp::instruction_builder::commit_state_from_buffer(
-            *validator,
-            value.committed_account.pubkey,
-            value.committed_account.account.owner,
-            commit_buffer_pubkey,
-            CommitStateFromBufferArgs {
-                nonce: value.commit_id,
-                lamports: value.committed_account.account.lamports,
-                allow_undelegation: value.allow_undelegation,
-            },
-        )
+        value.create_commit_ix(validator)
     }
 
     /// No further optimizations
     fn optimize(
         self: Box<Self>,
     ) -> Result<Box<dyn BaseTask>, Box<dyn BaseTask>> {
+        // Since the buffer in BufferTask doesn't contribute to the size of
+        // transaction, there is nothing we can do here to optimize/reduce the size.
         Err(self)
     }
 
diff --git a/magicblock-committor-service/src/tasks/mod.rs b/magicblock-committor-service/src/tasks/mod.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use dlp::{
-    args::{CommitDiffArgs, CommitStateArgs},
+    args::{CommitDiffArgs, CommitStateArgs, CommitStateFromBufferArgs},
     compute_diff,
 };
 use dyn_clone::DynClone;
@@ -53,7 +53,6 @@ pub enum PreparationState {
     Cleanup(CleanupTask),
 }
 
-#[cfg(test)]
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum TaskStrategy {
     Args,
@@ -153,7 +152,7 @@ impl CommitTaskBuilder {
             allow_undelegation,
             committed_account,
             base_account,
-            force_commit_state: false,
+            strategy: TaskStrategy::Args,
         }
     }
 }
@@ -164,29 +163,46 @@ pub struct CommitTask {
     pub allow_undelegation: bool,
     pub committed_account: CommittedAccount,
     base_account: Option<Account>,
-    force_commit_state: bool,
+    strategy: TaskStrategy,
 }
 
 impl CommitTask {
-    pub fn is_commit_diff(&self) -> bool {
-        !self.force_commit_state
-            && self.committed_account.account.data.len()
-                > CommitTaskBuilder::COMMIT_STATE_SIZE_THRESHOLD
-            && self.base_account.is_some()
-    }
-
-    pub fn force_commit_state(&mut self) {
-        self.force_commit_state = true;
+    pub fn switch_to_buffer_strategy(mut self) -> Self {
+        self.strategy = TaskStrategy::Buffer;
+        self
     }
 
     pub fn create_commit_ix(&self, validator: &Pubkey) -> Instruction {
-        if let Some(fetched_account) = self.base_account.as_ref() {
-            self.create_commit_diff_ix(validator, fetched_account)
-        } else {
-            self.create_commit_state_ix(validator)
+        match self.strategy {
+            TaskStrategy::Args => {
+                if let Some(base_account) = self.base_account.as_ref() {
+                    self.create_commit_diff_ix(validator, base_account)
+                } else {
+                    self.create_commit_state_ix(validator)
+                }
+            }
+            TaskStrategy::Buffer => {
+                if let Some(base_account) = self.base_account.as_ref() {
+                    self.create_commit_diff_from_buffer_ix(
+                        validator,
+                        base_account,
+                    )
+                } else {
+                    self.create_commit_state_from_buffer_ix(validator)
+                }
+            }
         }
     }
 
+    pub fn compute_diff(&self) -> Option<dlp::rkyv::AlignedVec> {
+        self.base_account.as_ref().map(|base_account| {
+            compute_diff(
+                base_account.data(),
+                self.committed_account.account.data(),
+            )
+        })
+    }
+
     fn create_commit_state_ix(&self, validator: &Pubkey) -> Instruction {
         let args = CommitStateArgs {
             nonce: self.commit_id,
@@ -205,17 +221,13 @@ impl CommitTask {
     fn create_commit_diff_ix(
         &self,
         validator: &Pubkey,
-        fetched_account: &Account,
+        base_account: &Account,
     ) -> Instruction {
-        if self.force_commit_state {
-            return self.create_commit_state_ix(validator);
-        }
-
         let args = CommitDiffArgs {
             nonce: self.commit_id,
             lamports: self.committed_account.account.lamports,
             diff: compute_diff(
-                fetched_account.data(),
+                base_account.data(),
                 self.committed_account.account.data(),
             )
             .to_vec(),
@@ -229,6 +241,57 @@ impl CommitTask {
             args,
         )
     }
+
+    fn create_commit_state_from_buffer_ix(
+        &self,
+        validator: &Pubkey,
+    ) -> Instruction {
+        let commit_id_slice = self.commit_id.to_le_bytes();
+        let (commit_buffer_pubkey, _) =
+            magicblock_committor_program::pdas::buffer_pda(
+                validator,
+                &self.committed_account.pubkey,
+                &commit_id_slice,
+            );
+
+        dlp::instruction_builder::commit_state_from_buffer(
+            *validator,
+            self.committed_account.pubkey,
+            self.committed_account.account.owner,
+            commit_buffer_pubkey,
+            CommitStateFromBufferArgs {
+                nonce: self.commit_id,
+                lamports: self.committed_account.account.lamports,
+                allow_undelegation: self.allow_undelegation,
+            },
+        )
+    }
+
+    fn create_commit_diff_from_buffer_ix(
+        &self,
+        validator: &Pubkey,
+        _fetched_account: &Account,
+    ) -> Instruction {
+        let commit_id_slice = self.commit_id.to_le_bytes();
+        let (commit_buffer_pubkey, _) =
+            magicblock_committor_program::pdas::buffer_pda(
+                validator,
+                &self.committed_account.pubkey,
+                &commit_id_slice,
+            );
+
+        dlp::instruction_builder::commit_diff_from_buffer(
+            *validator,
+            self.committed_account.pubkey,
+            self.committed_account.account.owner,
+            commit_buffer_pubkey,
+            CommitStateFromBufferArgs {
+                nonce: self.commit_id,
+                lamports: self.committed_account.account.lamports,
+                allow_undelegation: self.allow_undelegation,
+            },
+        )
+    }
 }
 
 #[derive(Clone)]
diff --git a/test-integration/Cargo.toml b/test-integration/Cargo.toml
@@ -57,8 +57,8 @@ magicblock-config = { path = "../magicblock-config" }
 magicblock-core = { path = "../magicblock-core" }
 magic-domain-program = { git = "https://github.com/magicblock-labs/magic-domain-program.git", rev = "ea04d46", default-features = false }
 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 = "e8d03936", features = [
-  "no-entrypoint",
+magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "ea1f2f916268132248fe8d5de5f07d76765dd937", features = [
+ "no-entrypoint",
 ] }
 magicblock-program = { path = "../programs/magicblock" }
 magicblock-rpc-client = { path = "../magicblock-rpc-client" }
diff --git a/test-integration/schedulecommit/elfs/dlp.so b/test-integration/schedulecommit/elfs/dlp.so
diff --git a/test-integration/schedulecommit/test-scenarios/tests/02_commit_and_undelegate.rs b/test-integration/schedulecommit/test-scenarios/tests/02_commit_and_undelegate.rs
@@ -303,13 +303,13 @@ fn test_committing_and_undelegating_huge_order_book_account() {
             println!("Important: use {rng_seed} as seed to regenerate the random inputs in case of test failure");
             let mut random = StdRng::seed_from_u64(rng_seed);
             let mut update = BookUpdate::default();
-            update.bids.extend((0..random.gen_range(5..10)).map(|_| {
+            update.bids.extend((0..random.gen_range(5..100)).map(|_| {
                 OrderLevel {
                     price: random.gen_range(75000..90000),
                     size: random.gen_range(1..10),
                 }
             }));
-            update.asks.extend((0..random.gen_range(5..10)).map(|_| {
+            update.asks.extend((0..random.gen_range(5..100)).map(|_| {
                 OrderLevel {
                     price: random.gen_range(125000..150000),
                     size: random.gen_range(1..10),
