feat: Execute CommitDiff as BufferTask

magicblock-committor-service/src/persist/commit_persister.rs

@@ -585,14 +585,14 @@ mod tests {
         persister.set_commit_id(1, &pubkey, 100).unwrap();
 
         persister
-            .set_commit_strategy(100, &pubkey, CommitStrategy::Args)
+            .set_commit_strategy(100, &pubkey, CommitStrategy::StateArgs)
             .unwrap();
 
         let updated = persister
             .get_commit_status_by_message(1, &pubkey)
             .unwrap()
             .unwrap();
-        assert_eq!(updated.commit_strategy, CommitStrategy::Args);
+        assert_eq!(updated.commit_strategy, CommitStrategy::StateArgs);
     }
 
     #[test]

magicblock-committor-service/src/persist/db.rs

@@ -772,7 +772,7 @@ mod tests {
             commit_type: CommitType::DataAccount,
             created_at: 1000,
             commit_status: CommitStatus::Pending,
-            commit_strategy: CommitStrategy::Args,
+            commit_strategy: CommitStrategy::StateArgs,
             last_retried_at: 1000,
             retries_count: 0,
         }
@@ -907,7 +907,7 @@ mod tests {
         db.insert_commit_status_rows(std::slice::from_ref(&row))
             .unwrap();
 
-        let new_strategy = CommitStrategy::FromBuffer;
+        let new_strategy = CommitStrategy::StateBuffer;
         db.set_commit_strategy(100, &row.pubkey, new_strategy)
             .unwrap();
 

magicblock-committor-service/src/persist/types/commit_strategy.rs

@@ -4,39 +4,46 @@ use crate::persist::error::CommitPersistError;
 pub enum CommitStrategy {
     /// Args without the use of a lookup table
     #[default]
-    Args,
+    StateArgs,
     /// Args with the use of a lookup table
-    ArgsWithLookupTable,
+    StateArgsWithLookupTable,
     /// Buffer and chunks which has the most overhead
-    FromBuffer,
+    StateBuffer,
     /// Buffer and chunks with the use of a lookup table
-    FromBufferWithLookupTable,
+    StateBufferWithLookupTable,
+
+    /// Args without the use of a lookup table
+    DiffArgs,
+    /// Args with the use of a lookup table
+    DiffArgsWithLookupTable,
+    /// Buffer and chunks which has the most overhead
+    DiffBuffer,
+    /// Buffer and chunks with the use of a lookup table
+    DiffBufferWithLookupTable,
 }
 
 impl CommitStrategy {
-    pub fn args(use_lookup: bool) -> Self {
-        if use_lookup {
-            Self::ArgsWithLookupTable
-        } else {
-            Self::Args
-        }
-    }
-
     pub fn as_str(&self) -> &str {
         use CommitStrategy::*;
         match self {
-            Args => "Args",
-            ArgsWithLookupTable => "ArgsWithLookupTable",
-            FromBuffer => "FromBuffer",
-            FromBufferWithLookupTable => "FromBufferWithLookupTable",
+            StateArgs => "StateArgs",
+            StateArgsWithLookupTable => "StateArgsWithLookupTable",
+            StateBuffer => "StateBuffer",
+            StateBufferWithLookupTable => "StateBufferWithLookupTable",
+            DiffArgs => "DiffArgs",
+            DiffArgsWithLookupTable => "DiffArgsWithLookupTable",
+            DiffBuffer => "DiffBuffer",
+            DiffBufferWithLookupTable => "DiffBufferWithLookupTable",
         }
     }
 
     pub fn uses_lookup(&self) -> bool {
         matches!(
             self,
-            CommitStrategy::ArgsWithLookupTable
-                | CommitStrategy::FromBufferWithLookupTable
+            CommitStrategy::StateArgsWithLookupTable
+                | CommitStrategy::StateBufferWithLookupTable
+                | CommitStrategy::DiffArgsWithLookupTable
+                | CommitStrategy::DiffBufferWithLookupTable
         )
     }
 }
@@ -45,10 +52,18 @@ impl TryFrom<&str> for CommitStrategy {
     type Error = CommitPersistError;
     fn try_from(value: &str) -> Result<Self, CommitPersistError> {
         match value {
-            "Args" => Ok(Self::Args),
-            "ArgsWithLookupTable" => Ok(Self::ArgsWithLookupTable),
-            "FromBuffer" => Ok(Self::FromBuffer),
-            "FromBufferWithLookupTable" => Ok(Self::FromBufferWithLookupTable),
+            "Args" | "StateArgs" => Ok(Self::StateArgs),
+            "ArgsWithLookupTable" | "StateArgsWithLookupTable" => {
+                Ok(Self::StateArgsWithLookupTable)
+            }
+            "FromBuffer" | "StateBuffer" => Ok(Self::StateBuffer),
+            "FromBufferWithLookupTable" | "StateBufferWithLookupTable" => {
+                Ok(Self::StateBufferWithLookupTable)
+            }
+            "DiffArgs" => Ok(Self::DiffArgs),
+            "DiffArgsWithLookupTable" => Ok(Self::DiffArgsWithLookupTable),
+            "DiffBuffer" => Ok(Self::DiffBuffer),
+            "DiffBufferWithLookupTable" => Ok(Self::DiffBufferWithLookupTable),
             _ => Err(CommitPersistError::InvalidCommitStrategy(
                 value.to_string(),
             )),

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

@@ -122,7 +122,7 @@ impl BaseTask for ArgsTask {
         }
     }
 
-    fn optimize(
+    fn try_optimize_tx_size(
         self: Box<Self>,
     ) -> Result<Box<dyn BaseTask>, Box<dyn BaseTask>> {
         match self.task_type {
@@ -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(_)

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,38 +89,64 @@ 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
-    fn optimize(
+    fn try_optimize_tx_size(
         self: Box<Self>,
     ) -> Result<Box<dyn BaseTask>, Box<dyn BaseTask>> {
         Err(self)
@@ -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,15 @@ 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(task) => {
+                task.commit_id = commit_id;
+            }
+            BufferTaskType::CommitDiff(task) => {
+                task.commit_id = commit_id;
+            }
+        };
 
-        commit_task.commit_id = commit_id;
         self.preparation_state = Self::preparation_required(&self.task_type)
     }
 }
@@ -159,6 +210,7 @@ impl LabelValue for BufferTask {
     fn value(&self) -> &str {
         match self.task_type {
             BufferTaskType::Commit(_) => "buffer_commit",
+            BufferTaskType::CommitDiff(_) => "buffer_commit_diff",
         }
     }
 }

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,
@@ -67,8 +66,9 @@ pub trait BaseTask: Send + Sync + DynClone + LabelValue {
     /// Gets instruction for task execution
     fn instruction(&self, validator: &Pubkey) -> Instruction;
 
-    /// Optimizes Task strategy if possible, otherwise returns itself
-    fn optimize(
+    /// Optimize for transaction size so that more instructions can be buddled together in a single
+    /// transaction. Return Ok(new_tx_optimized_task), else Err(self) if task cannot be optimized.
+    fn try_optimize_tx_size(
         self: Box<Self>,
     ) -> Result<Box<dyn BaseTask>, Box<dyn BaseTask>>;