feat: Add ScheduleCommitDiffAndUndelegate to efficiently commit changes in delegated accounts

Cargo.toml

@@ -108,10 +108,8 @@ magicblock-config = { path = "./magicblock-config" }
 magicblock-config-helpers = { path = "./magicblock-config-helpers" }
 magicblock-config-macro = { path = "./magicblock-config-macro" }
 magicblock-core = { path = "./magicblock-core" }
-magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "aa1de56d90c", features = [
-  "no-entrypoint",
-] }
 magicblock-aperture = { path = "./magicblock-aperture" }
+magicblock-delegation-program = { path="../delegation-program", features = ["no-entrypoint"] }
 magicblock-geyser-plugin = { path = "./magicblock-geyser-plugin" }
 magicblock-ledger = { path = "./magicblock-ledger" }
 magicblock-metrics = { path = "./magicblock-metrics" }

magicblock-accounts/src/scheduled_commits_processor.rs

@@ -343,6 +343,7 @@ impl ScheduledCommitsProcessorImpl {
             included_pubkeys: intent_meta.included_pubkeys,
             excluded_pubkeys: intent_meta.excluded_pubkeys,
             requested_undelegation: intent_meta.requested_undelegation,
+            commit_diff: intent_meta.commit_diff,
         }
     }
 }
@@ -412,6 +413,7 @@ struct ScheduledBaseIntentMeta {
     excluded_pubkeys: Vec<Pubkey>,
     intent_sent_transaction: Transaction,
     requested_undelegation: bool,
+    commit_diff: bool,
 }
 
 impl ScheduledBaseIntentMeta {
@@ -429,6 +431,7 @@ impl ScheduledBaseIntentMeta {
             excluded_pubkeys,
             intent_sent_transaction: intent.action_sent_transaction.clone(),
             requested_undelegation: intent.is_undelegate(),
+            commit_diff: intent.is_commit_diff(),
         }
     }
 }

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

@@ -1,20 +1,33 @@
-use dlp::args::{CallHandlerArgs, CommitStateArgs};
+use dlp::{
+    args::{CallHandlerArgs, CommitDiffArgs, CommitStateArgs},
+    compute_diff,
+};
+use solana_account::ReadableAccount;
 use solana_pubkey::Pubkey;
-use solana_sdk::instruction::{AccountMeta, Instruction};
+use solana_rpc_client::rpc_client::RpcClient;
+use solana_sdk::{
+    commitment_config::CommitmentConfig,
+    instruction::{AccountMeta, Instruction},
+};
 
 #[cfg(test)]
 use crate::tasks::TaskStrategy;
-use crate::tasks::{
-    buffer_task::{BufferTask, BufferTaskType},
-    visitor::Visitor,
-    BaseActionTask, BaseTask, BaseTaskError, BaseTaskResult, CommitTask,
-    FinalizeTask, PreparationState, TaskType, UndelegateTask,
+use crate::{
+    config::ChainConfig,
+    tasks::{
+        buffer_task::{BufferTask, BufferTaskType},
+        visitor::Visitor,
+        BaseActionTask, BaseTask, BaseTaskError, BaseTaskResult, CommitTask,
+        FinalizeTask, PreparationState, TaskType, UndelegateTask,
+    },
+    ComputeBudgetConfig,
 };
 
 /// Task that will be executed on Base layer via arguments
 #[derive(Clone)]
 pub enum ArgsTaskType {
     Commit(CommitTask),
+    CommitDiff(CommitTask),
     Finalize(FinalizeTask),
     Undelegate(UndelegateTask), // Special action really
     BaseAction(BaseActionTask),
@@ -58,6 +71,55 @@ impl BaseTask for ArgsTask {
                     args,
                 )
             }
+            ArgsTaskType::CommitDiff(value) => {
+                let chain_config =
+                    ChainConfig::local(ComputeBudgetConfig::new(1_000_000));
+
+                let rpc_client = RpcClient::new_with_commitment(
+                    chain_config.rpc_uri.to_string(),
+                    CommitmentConfig {
+                        commitment: chain_config.commitment,
+                    },
+                );
+
+                let account = match rpc_client
+                    .get_account(&value.committed_account.pubkey)
+                {
+                    Ok(account) => account,
+                    Err(e) => {
+                        log::warn!("Fallback to commit_state and send full-bytes, as rpc failed to fetch the delegated-account from base chain: {}", e);
+                        let args = CommitStateArgs {
+                            nonce: value.commit_id,
+                            lamports: value.committed_account.account.lamports,
+                            data: value.committed_account.account.data.clone(),
+                            allow_undelegation: value.allow_undelegation,
+                        };
+                        return dlp::instruction_builder::commit_state(
+                            *validator,
+                            value.committed_account.pubkey,
+                            value.committed_account.account.owner,
+                            args,
+                        );
+                    }
+                };
+
+                let args = CommitDiffArgs {
+                    nonce: value.commit_id,
+                    lamports: value.committed_account.account.lamports,
+                    diff: compute_diff(
+                        account.data(),
+                        value.committed_account.account.data(),
+                    ),
+                    allow_undelegation: value.allow_undelegation,
+                };
+                log::warn!("DIFF computed: {:?}", args.diff);
+                dlp::instruction_builder::commit_diff(
+                    *validator,
+                    value.committed_account.pubkey,
+                    value.committed_account.account.owner,
+                    args,
+                )
+            }
             ArgsTaskType::Finalize(value) => {
                 dlp::instruction_builder::finalize(
                     *validator,
@@ -106,6 +168,9 @@ impl BaseTask for ArgsTask {
                     BufferTaskType::Commit(value),
                 )))
             }
+            ArgsTaskType::CommitDiff(_) => {
+                panic!("ArgsTaskType::CommitDiff not handled")
+            }
             ArgsTaskType::BaseAction(_)
             | ArgsTaskType::Finalize(_)
             | ArgsTaskType::Undelegate(_) => Err(self),
@@ -132,6 +197,7 @@ impl BaseTask for ArgsTask {
     fn compute_units(&self) -> u32 {
         match &self.task_type {
             ArgsTaskType::Commit(_) => 70_000,
+            ArgsTaskType::CommitDiff(_) => 65_000,
             ArgsTaskType::BaseAction(task) => task.action.compute_units,
             ArgsTaskType::Undelegate(_) => 70_000,
             ArgsTaskType::Finalize(_) => 70_000,
@@ -146,6 +212,9 @@ impl BaseTask for ArgsTask {
     fn task_type(&self) -> TaskType {
         match &self.task_type {
             ArgsTaskType::Commit(_) => TaskType::Commit,
+            // TODO (snawaz): What should we use here? Commit (in the sense of "category of task"), or add a
+            // new variant "CommitDiff" to indicate a specific instruction?
+            ArgsTaskType::CommitDiff(_) => TaskType::Commit,
             ArgsTaskType::BaseAction(_) => TaskType::Action,
             ArgsTaskType::Undelegate(_) => TaskType::Undelegate,
             ArgsTaskType::Finalize(_) => TaskType::Finalize,
@@ -158,7 +227,9 @@ impl BaseTask for ArgsTask {
     }
 
     fn reset_commit_id(&mut self, commit_id: u64) {
+        // TODO (snawaz): handle CommitDiff as well?
         let ArgsTaskType::Commit(commit_task) = &mut self.task_type else {
+            log::error!("reset_commit_id");
             return;
         };
 

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

@@ -53,6 +53,8 @@ pub enum TaskStrategy {
 pub trait BaseTask: Send + Sync + DynClone {
     /// Gets all pubkeys that involved in Task's instruction
     fn involved_accounts(&self, validator: &Pubkey) -> Vec<Pubkey> {
+        // TODO (snawaz): rewrite it.
+        // currently it is slow as it discards heavy computations and memory allocations.
         self.instruction(validator)
             .accounts
             .iter()

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

@@ -47,25 +47,29 @@ impl TasksBuilder for TaskBuilderImpl {
         base_intent: &ScheduledBaseIntent,
         persister: &Option<P>,
     ) -> TaskBuilderResult<Vec<Box<dyn BaseTask>>> {
-        let (accounts, allow_undelegation) = match &base_intent.base_intent {
-            MagicBaseIntent::BaseActions(actions) => {
-                let tasks = actions
-                    .iter()
-                    .map(|el| {
-                        let task = BaseActionTask { action: el.clone() };
-                        let task =
-                            ArgsTask::new(ArgsTaskType::BaseAction(task));
-                        Box::new(task) as Box<dyn BaseTask>
-                    })
-                    .collect();
-
-                return Ok(tasks);
-            }
-            MagicBaseIntent::Commit(t) => (t.get_committed_accounts(), false),
-            MagicBaseIntent::CommitAndUndelegate(t) => {
-                (t.commit_action.get_committed_accounts(), true)
-            }
-        };
+        let (accounts, allow_undelegation, commit_diff) =
+            match &base_intent.base_intent {
+                MagicBaseIntent::BaseActions(actions) => {
+                    let tasks = actions
+                        .iter()
+                        .map(|el| {
+                            let task = BaseActionTask { action: el.clone() };
+                            let task =
+                                ArgsTask::new(ArgsTaskType::BaseAction(task));
+                            Box::new(task) as Box<dyn BaseTask>
+                        })
+                        .collect();
+                    return Ok(tasks);
+                }
+                MagicBaseIntent::Commit(t) => {
+                    (t.get_committed_accounts(), false, t.is_commit_diff())
+                }
+                MagicBaseIntent::CommitAndUndelegate(t) => (
+                    t.commit_action.get_committed_accounts(),
+                    true,
+                    t.commit_action.is_commit_diff(),
+                ),
+            };
 
         let committed_pubkeys = accounts
             .iter()
@@ -89,11 +93,16 @@ impl TasksBuilder for TaskBuilderImpl {
             .iter()
             .map(|account| {
                 let commit_id = *commit_ids.get(&account.pubkey).expect("CommitIdFetcher provide commit ids for all listed pubkeys, or errors!");
-                let task = ArgsTaskType::Commit(CommitTask {
+                let task = CommitTask {
                     commit_id,
                     allow_undelegation,
                     committed_account: account.clone(),
-                });
+                };
+                let task = if commit_diff {
+                    ArgsTaskType::CommitDiff(task)
+                } else {
+                    ArgsTaskType::Commit(task)
+                };
 
                 Box::new(ArgsTask::new(task)) as Box<dyn BaseTask>
             })
@@ -134,6 +143,9 @@ impl TasksBuilder for TaskBuilderImpl {
                 CommitType::Standalone(accounts) => {
                     accounts.iter().map(finalize_task).collect()
                 }
+                CommitType::StandaloneDiff(accounts) => {
+                    accounts.iter().map(finalize_task).collect()
+                }
                 CommitType::WithBaseActions {
                     committed_accounts,
                     base_actions,

magicblock-committor-service/src/tasks/task_visitors/persistor_visitor.rs

@@ -26,9 +26,10 @@ 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_task = match &task.task_type {
+                    ArgsTaskType::Commit(commit_task) => commit_task,
+                    ArgsTaskType::CommitDiff(commit_task) => commit_task,
+                    _ => return,
                 };
 
                 let commit_strategy = if uses_lookup_tables {

magicblock-magic-program-api/src/instruction.rs

@@ -107,6 +107,8 @@ pub enum MagicBlockInstruction {
     /// # Account references
     /// - **0.** `[SIGNER]`         Validator authority
     EnableExecutableCheck,
+
+    ScheduleCommitDiffAndUndelegate,
 }
 
 impl MagicBlockInstruction {

magicblock-rpc-client/src/lib.rs

@@ -427,6 +427,7 @@ impl MagicblockRpcClient {
             .await?;
 
         if let Err(err) = processed_status {
+            error!("> ERROR: {:?}", err);
             return Err(MagicBlockRpcClientError::SentTransactionError(
                 err, sig,
             ));