feat: Execute CommitDiff as BufferTask

snawaz · snawaz · commit 9e0f949aeb4c · 2025-11-21T23:05:03.000+04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -109,8 +109,8 @@ magicblock-config-helpers = { path = "./magicblock-config-helpers" }
 magicblock-config-macro = { path = "./magicblock-config-macro" }
 magicblock-core = { path = "./magicblock-core" }
 magicblock-aperture = { path = "./magicblock-aperture" }
-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-geyser-plugin = { path = "./magicblock-geyser-plugin" }
 magicblock-ledger = { path = "./magicblock-ledger" }
diff --git a/magicblock-committor-service/src/tasks/args_task.rs b/magicblock-committor-service/src/tasks/args_task.rs
@@ -88,21 +88,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 solana_pubkey::Pubkey;
 use solana_sdk::instruction::Instruction;
@@ -46,16 +45,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,
         })
     }
@@ -64,31 +65,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;
@@ -52,7 +52,6 @@ pub enum PreparationState {
     Cleanup(CleanupTask),
 }
 
-#[cfg(test)]
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum TaskStrategy {
     Args,
@@ -152,7 +151,7 @@ impl CommitTaskBuilder {
             allow_undelegation,
             committed_account,
             base_account,
-            force_commit_state: false,
+            strategy: TaskStrategy::Args,
         }
     }
 }
@@ -174,18 +173,42 @@ impl CommitTask {
             && 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,
@@ -204,17 +227,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(),
@@ -228,6 +247,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.lock b/test-integration/Cargo.lock
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/test-scenarios/tests/02_commit_and_undelegate.rs b/test-integration/schedulecommit/test-scenarios/tests/02_commit_and_undelegate.rs
@@ -253,13 +253,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),
