feat: Execute CommitDiff as BufferTask

snawaz · snawaz · commit 632a4f975131 · 2025-12-22T14:56:50.000+05:30
diff --git a/magicblock-committor-service/src/tasks/args_task.rs b/magicblock-committor-service/src/tasks/args_task.rs
@@ -132,16 +132,8 @@ impl BaseTask for ArgsTask {
                 )))
             }
             ArgsTaskType::CommitDiff(value) => {
-                // TODO (snawaz): Currently, we do not support executing CommitDiff
-                // as BufferTask, which is why we're forcing CommitDiffTask to become CommitTask
-                // before converting this task into BufferTask. Once CommitDiff is supported
-                // by BufferTask, we do not have to do this, as it's essentially a downgrade.
                 Ok(Box::new(BufferTask::new_preparation_required(
-                    BufferTaskType::Commit(CommitTask {
-                        commit_id: value.commit_id,
-                        allow_undelegation: value.allow_undelegation,
-                        committed_account: value.committed_account,
-                    }),
+                    BufferTaskType::CommitDiff(value),
                 )))
             }
             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,4 @@
-use dlp::args::CommitStateFromBufferArgs;
+use dlp::{args::CommitStateFromBufferArgs, compute_diff};
 use magicblock_committor_program::Chunks;
 use magicblock_metrics::metrics::LabelValue;
 use solana_instruction::Instruction;
@@ -11,15 +11,17 @@ use crate::tasks::TaskStrategy;
 use crate::{
     consts::MAX_WRITE_CHUNK_SIZE,
     tasks::{
-        visitor::Visitor, BaseTask, BaseTaskError, BaseTaskResult, CommitTask,
-        PreparationState, PreparationTask, TaskType,
+        visitor::Visitor, BaseTask, BaseTaskError, BaseTaskResult,
+        CommitDiffTask, CommitTask, PreparationState, PreparationTask,
+        TaskType,
     },
 };
 
 /// Tasks that could be executed using buffers
 #[derive(Clone)]
 pub enum BufferTaskType {
     Commit(CommitTask),
+    CommitDiff(CommitDiffTask),
     // Action in the future
 }
 
@@ -48,19 +50,37 @@ 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,
-        );
-
-        PreparationState::Required(PreparationTask {
-            commit_id: commit_task.commit_id,
-            pubkey: commit_task.committed_account.pubkey,
-            committed_data,
-            chunks,
-        })
+        match task_type {
+            BufferTaskType::Commit(task) => {
+                let data = task.committed_account.account.data.clone();
+                let chunks =
+                    Chunks::from_data_length(data.len(), MAX_WRITE_CHUNK_SIZE);
+
+                PreparationState::Required(PreparationTask {
+                    commit_id: task.commit_id,
+                    pubkey: task.committed_account.pubkey,
+                    committed_data: data,
+                    chunks,
+                })
+            }
+
+            BufferTaskType::CommitDiff(task) => {
+                let diff = compute_diff(
+                    &task.base_account.data,
+                    &task.committed_account.account.data,
+                )
+                .to_vec();
+                let chunks =
+                    Chunks::from_data_length(diff.len(), MAX_WRITE_CHUNK_SIZE);
+
+                PreparationState::Required(PreparationTask {
+                    commit_id: task.commit_id,
+                    pubkey: task.committed_account.pubkey,
+                    committed_data: diff,
+                    chunks,
+                })
+            }
+        }
     }
 }
 
@@ -69,34 +89,60 @@ impl From<ArgsTaskType> for BufferTaskType {
     fn from(value: ArgsTaskType) -> Self {
         match value {
             ArgsTaskType::Commit(task) => BufferTaskType::Commit(task),
-            ArgsTaskType::CommitDiff(_) => panic!("BufferTask doesn't support CommitDiff yet. Disable your tests (if any) temporarily till the next PR"),
-            _ => unimplemented!("Only commit task can be BufferTask currently. Fix your tests"),
+            ArgsTaskType::CommitDiff(task) => BufferTaskType::CommitDiff(task),
+            _ => unimplemented!(
+                "Only commit task can be BufferTask currently. Fix your tests"
+            ),
         }
     }
 }
 
 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,
-            },
-        )
+        match &self.task_type {
+            BufferTaskType::Commit(task) => {
+                let commit_id_slice = task.commit_id.to_le_bytes();
+                let (commit_buffer_pubkey, _) =
+                    magicblock_committor_program::pdas::buffer_pda(
+                        validator,
+                        &task.committed_account.pubkey,
+                        &commit_id_slice,
+                    );
+
+                dlp::instruction_builder::commit_state_from_buffer(
+                    *validator,
+                    task.committed_account.pubkey,
+                    task.committed_account.account.owner,
+                    commit_buffer_pubkey,
+                    CommitStateFromBufferArgs {
+                        nonce: task.commit_id,
+                        lamports: task.committed_account.account.lamports,
+                        allow_undelegation: task.allow_undelegation,
+                    },
+                )
+            }
+            BufferTaskType::CommitDiff(task) => {
+                let commit_id_slice = task.commit_id.to_le_bytes();
+                let (commit_buffer_pubkey, _) =
+                    magicblock_committor_program::pdas::buffer_pda(
+                        validator,
+                        &task.committed_account.pubkey,
+                        &commit_id_slice,
+                    );
+
+                dlp::instruction_builder::commit_diff_from_buffer(
+                    *validator,
+                    task.committed_account.pubkey,
+                    task.committed_account.account.owner,
+                    commit_buffer_pubkey,
+                    CommitStateFromBufferArgs {
+                        nonce: task.commit_id,
+                        lamports: task.committed_account.account.lamports,
+                        allow_undelegation: task.allow_undelegation,
+                    },
+                )
+            }
+        }
     }
 
     /// No further optimizations
@@ -125,6 +171,7 @@ impl BaseTask for BufferTask {
     fn compute_units(&self) -> u32 {
         match self.task_type {
             BufferTaskType::Commit(_) => 70_000,
+            BufferTaskType::CommitDiff(_) => 70_000,
         }
     }
 
@@ -136,6 +183,7 @@ impl BaseTask for BufferTask {
     fn task_type(&self) -> TaskType {
         match self.task_type {
             BufferTaskType::Commit(_) => TaskType::Commit,
+            BufferTaskType::CommitDiff(_) => TaskType::Commit,
         }
     }
 
@@ -145,12 +193,21 @@ impl BaseTask for BufferTask {
     }
 
     fn reset_commit_id(&mut self, commit_id: u64) {
-        let BufferTaskType::Commit(commit_task) = &mut self.task_type;
-        if commit_id == commit_task.commit_id {
-            return;
-        }
+        match &mut self.task_type {
+            BufferTaskType::Commit(commit_task) => {
+                if commit_id == commit_task.commit_id {
+                    return;
+                }
+                commit_task.commit_id = commit_id;
+            }
+            BufferTaskType::CommitDiff(task) => {
+                if commit_id == task.commit_id {
+                    return;
+                }
+                task.commit_id = commit_id;
+            }
+        };
 
-        commit_task.commit_id = commit_id;
         self.preparation_state = Self::preparation_required(&self.task_type)
     }
 }
@@ -159,6 +216,7 @@ impl LabelValue for BufferTask {
     fn value(&self) -> &str {
         match self.task_type {
             BufferTaskType::Commit(_) => "buffer_commit",
+            BufferTaskType::CommitDiff(_) => "buffer_commit_diff",
         }
     }
 }
diff --git a/magicblock-committor-service/src/tasks/mod.rs b/magicblock-committor-service/src/tasks/mod.rs
@@ -46,7 +46,6 @@ pub enum PreparationState {
     Cleanup(CleanupTask),
 }
 
-#[cfg(test)]
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum TaskStrategy {
     Args,
diff --git a/magicblock-committor-service/src/tasks/task_visitors/persistor_visitor.rs b/magicblock-committor-service/src/tasks/task_visitors/persistor_visitor.rs
@@ -26,27 +26,40 @@ where
     fn visit_args_task(&mut self, task: &ArgsTask) {
         match self.context {
             PersistorContext::PersistStrategy { uses_lookup_tables } => {
-                let ArgsTaskType::Commit(ref commit_task) = task.task_type
-                else {
-                    return;
-                };
-
                 let commit_strategy = if uses_lookup_tables {
                     CommitStrategy::ArgsWithLookupTable
                 } else {
                     CommitStrategy::Args
                 };
 
-                if let Err(err) = self.persistor.set_commit_strategy(
-                    commit_task.commit_id,
-                    &commit_task.committed_account.pubkey,
-                    commit_strategy,
-                ) {
-                    error!(
-                        "Failed to persist commit strategy {}: {}",
-                        commit_strategy.as_str(),
-                        err
-                    );
+                match &task.task_type {
+                    ArgsTaskType::Commit(task) => {
+                        if let Err(err) = self.persistor.set_commit_strategy(
+                            task.commit_id,
+                            &task.committed_account.pubkey,
+                            commit_strategy,
+                        ) {
+                            error!(
+                                "Failed to persist commit strategy {}: {}",
+                                commit_strategy.as_str(),
+                                err
+                            );
+                        }
+                    }
+                    ArgsTaskType::CommitDiff(task) => {
+                        if let Err(err) = self.persistor.set_commit_strategy(
+                            task.commit_id,
+                            &task.committed_account.pubkey,
+                            commit_strategy,
+                        ) {
+                            error!(
+                                "Failed to persist commit strategy {}: {}",
+                                commit_strategy.as_str(),
+                                err
+                            );
+                        }
+                    }
+                    _ => {}
                 }
             }
         }
@@ -55,23 +68,39 @@ where
     fn visit_buffer_task(&mut self, task: &BufferTask) {
         match self.context {
             PersistorContext::PersistStrategy { uses_lookup_tables } => {
-                let BufferTaskType::Commit(ref commit_task) = task.task_type;
                 let commit_strategy = if uses_lookup_tables {
                     CommitStrategy::FromBufferWithLookupTable
                 } else {
                     CommitStrategy::FromBuffer
                 };
 
-                if let Err(err) = self.persistor.set_commit_strategy(
-                    commit_task.commit_id,
-                    &commit_task.committed_account.pubkey,
-                    commit_strategy,
-                ) {
-                    error!(
-                        "Failed to persist commit strategy {}: {}",
-                        commit_strategy.as_str(),
-                        err
-                    );
+                match &task.task_type {
+                    BufferTaskType::Commit(task) => {
+                        if let Err(err) = self.persistor.set_commit_strategy(
+                            task.commit_id,
+                            &task.committed_account.pubkey,
+                            commit_strategy,
+                        ) {
+                            error!(
+                                "Failed to persist commit strategy {}: {}",
+                                commit_strategy.as_str(),
+                                err
+                            );
+                        }
+                    }
+                    BufferTaskType::CommitDiff(task) => {
+                        if let Err(err) = self.persistor.set_commit_strategy(
+                            task.commit_id,
+                            &task.committed_account.pubkey,
+                            commit_strategy,
+                        ) {
+                            error!(
+                                "Failed to persist commit strategy {}: {}",
+                                commit_strategy.as_str(),
+                                err
+                            );
+                        }
+                    }
                 }
             }
         }
diff --git a/magicblock-committor-service/src/tasks/task_visitors/utility_visitor.rs b/magicblock-committor-service/src/tasks/task_visitors/utility_visitor.rs
@@ -19,23 +19,39 @@ impl Visitor for TaskVisitorUtils {
     fn visit_args_task(&mut self, task: &ArgsTask) {
         let Self::GetCommitMeta(commit_meta) = self;
 
-        if let ArgsTaskType::Commit(ref commit_task) = task.task_type {
-            *commit_meta = Some(CommitMeta {
-                committed_pubkey: commit_task.committed_account.pubkey,
-                commit_id: commit_task.commit_id,
-            })
-        } else {
-            *commit_meta = None
+        match &task.task_type {
+            ArgsTaskType::Commit(task) => {
+                *commit_meta = Some(CommitMeta {
+                    committed_pubkey: task.committed_account.pubkey,
+                    commit_id: task.commit_id,
+                })
+            }
+            ArgsTaskType::CommitDiff(task) => {
+                *commit_meta = Some(CommitMeta {
+                    committed_pubkey: task.committed_account.pubkey,
+                    commit_id: task.commit_id,
+                })
+            }
+            _ => *commit_meta = None,
         }
     }
 
     fn visit_buffer_task(&mut self, task: &BufferTask) {
         let Self::GetCommitMeta(commit_meta) = self;
 
-        let BufferTaskType::Commit(ref commit_task) = task.task_type;
-        *commit_meta = Some(CommitMeta {
-            committed_pubkey: commit_task.committed_account.pubkey,
-            commit_id: commit_task.commit_id,
-        })
+        match &task.task_type {
+            BufferTaskType::Commit(task) => {
+                *commit_meta = Some(CommitMeta {
+                    committed_pubkey: task.committed_account.pubkey,
+                    commit_id: task.commit_id,
+                })
+            }
+            BufferTaskType::CommitDiff(task) => {
+                *commit_meta = Some(CommitMeta {
+                    committed_pubkey: task.committed_account.pubkey,
+                    commit_id: task.commit_id,
+                })
+            }
+        }
     }
 }
diff --git a/test-integration/Cargo.toml b/test-integration/Cargo.toml
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

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,6 @@ pub enum PreparationState {`
`46`	`46`	`Cleanup(CleanupTask),`
`47`	`47`	`}`
`48`	`48`
`49`		`-#[cfg(test)]`
`50`	`49`	`#[derive(Copy, Clone, PartialEq, Eq, Debug)]`
`51`	`50`	`pub enum TaskStrategy {`
`52`	`51`	`Args,`