Use existing style of tasks

Author: snawaz

magicblock-committor-service/src/intent_execution_manager/intent_execution_engine.rs

@@ -371,8 +371,6 @@ mod tests {
 
     use async_trait::async_trait;
     use magicblock_program::magic_scheduled_base_intent::ScheduledBaseIntent;
-    use magicblock_rpc_client::MagicBlockRpcClientResult;
-    use solana_account::Account;
     use solana_pubkey::{pubkey, Pubkey};
     use solana_signature::Signature;
     use solana_signer::SignerError;

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

@@ -1,5 +1,9 @@
-use dlp::args::CallHandlerArgs;
+use dlp::{
+    args::{CallHandlerArgs, CommitDiffArgs, CommitStateArgs},
+    compute_diff,
+};
 use magicblock_metrics::metrics::LabelValue;
+use solana_account::ReadableAccount;
 use solana_instruction::{AccountMeta, Instruction};
 use solana_pubkey::Pubkey;
 
@@ -8,14 +12,15 @@ use crate::tasks::TaskStrategy;
 use crate::tasks::{
     buffer_task::{BufferTask, BufferTaskType},
     visitor::Visitor,
-    BaseActionTask, BaseTask, BaseTaskError, BaseTaskResult, CommitTask,
-    FinalizeTask, PreparationState, TaskType, UndelegateTask,
+    BaseActionTask, BaseTask, BaseTaskError, BaseTaskResult, CommitDiffTask,
+    CommitTask, FinalizeTask, PreparationState, TaskType, UndelegateTask,
 };
 
 /// Task that will be executed on Base layer via arguments
 #[derive(Clone)]
 pub enum ArgsTaskType {
     Commit(CommitTask),
+    CommitDiff(CommitDiffTask),
     Finalize(FinalizeTask),
     Undelegate(UndelegateTask), // Special action really
     BaseAction(BaseActionTask),
@@ -45,7 +50,39 @@ impl ArgsTask {
 impl BaseTask for ArgsTask {
     fn instruction(&self, validator: &Pubkey) -> Instruction {
         match &self.task_type {
-            ArgsTaskType::Commit(value) => value.create_commit_ix(validator),
+            ArgsTaskType::Commit(value) => {
+                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,
+                };
+                dlp::instruction_builder::commit_state(
+                    *validator,
+                    value.committed_account.pubkey,
+                    value.committed_account.account.owner,
+                    args,
+                )
+            }
+            ArgsTaskType::CommitDiff(value) => {
+                let args = CommitDiffArgs {
+                    nonce: value.commit_id,
+                    lamports: value.committed_account.account.lamports,
+                    diff: compute_diff(
+                        value.base_account.data(),
+                        value.committed_account.account.data(),
+                    )
+                    .to_vec(),
+                    allow_undelegation: value.allow_undelegation,
+                };
+
+                dlp::instruction_builder::commit_diff(
+                    *validator,
+                    value.committed_account.pubkey,
+                    value.committed_account.account.owner,
+                    args,
+                )
+            }
             ArgsTaskType::Finalize(value) => {
                 dlp::instruction_builder::finalize(
                     *validator,
@@ -89,21 +126,22 @@ 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();
+            ArgsTaskType::Commit(value) => {
                 Ok(Box::new(BufferTask::new_preparation_required(
                     BufferTaskType::Commit(value),
                 )))
             }
-            ArgsTaskType::Commit(value) => {
+            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(value),
+                    BufferTaskType::Commit(CommitTask {
+                        commit_id: value.commit_id,
+                        allow_undelegation: value.allow_undelegation,
+                        committed_account: value.committed_account,
+                    }),
                 )))
             }
             ArgsTaskType::BaseAction(_)
@@ -132,6 +170,7 @@ impl BaseTask for ArgsTask {
     fn compute_units(&self) -> u32 {
         match &self.task_type {
             ArgsTaskType::Commit(_) => 70_000,
+            ArgsTaskType::CommitDiff(_) => 70_000,
             ArgsTaskType::BaseAction(task) => task.action.compute_units,
             ArgsTaskType::Undelegate(_) => 70_000,
             ArgsTaskType::Finalize(_) => 70_000,
@@ -146,6 +185,7 @@ impl BaseTask for ArgsTask {
     fn task_type(&self) -> TaskType {
         match &self.task_type {
             ArgsTaskType::Commit(_) => TaskType::Commit,
+            ArgsTaskType::CommitDiff(_) => TaskType::Commit,
             ArgsTaskType::BaseAction(_) => TaskType::Action,
             ArgsTaskType::Undelegate(_) => TaskType::Undelegate,
             ArgsTaskType::Finalize(_) => TaskType::Finalize,
@@ -170,6 +210,7 @@ impl LabelValue for ArgsTask {
     fn value(&self) -> &str {
         match self.task_type {
             ArgsTaskType::Commit(_) => "args_commit",
+            ArgsTaskType::CommitDiff(_) => "args_commit_diff",
             ArgsTaskType::BaseAction(_) => "args_action",
             ArgsTaskType::Finalize(_) => "args_finalize",
             ArgsTaskType::Undelegate(_) => "args_undelegate",

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

@@ -1,9 +1,3 @@
-use std::sync::Arc;
-
-use dlp::{
-    args::{CommitDiffArgs, CommitStateArgs},
-    compute_diff,
-};
 use dyn_clone::DynClone;
 use magicblock_committor_program::{
     instruction_builder::{
@@ -20,15 +14,12 @@ use magicblock_metrics::metrics::LabelValue;
 use magicblock_program::magic_scheduled_base_intent::{
     BaseAction, CommittedAccount,
 };
-use solana_account::{Account, ReadableAccount};
+use solana_account::Account;
 use solana_instruction::Instruction;
 use solana_pubkey::Pubkey;
 use thiserror::Error;
 
-use crate::{
-    intent_executor::task_info_fetcher::TaskInfoFetcher,
-    tasks::visitor::Visitor,
-};
+use crate::tasks::visitor::Visitor;
 
 pub mod args_task;
 pub mod buffer_task;
@@ -64,8 +55,6 @@ pub enum TaskStrategy {
 pub trait BaseTask: Send + Sync + DynClone + LabelValue {
     /// 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()
@@ -115,78 +104,14 @@ pub struct CommitTask {
     pub commit_id: u64,
     pub allow_undelegation: bool,
     pub committed_account: CommittedAccount,
-    base_account: Option<Account>,
-    force_commit_state: bool,
 }
 
-impl CommitTask {
-    // Accounts larger than COMMIT_STATE_SIZE_THRESHOLD, use CommitDiff to
-    // reduce instruction size. Below this, commit is sent as CommitState.
-    // Chose 256 as thresold seems good enough as it could hold 8 u32 fields
-    // or 4 u64 fields.
-    const COMMIT_STATE_SIZE_THRESHOLD: usize = 256;
-
-    pub fn is_commit_diff(&self) -> bool {
-        !self.force_commit_state
-            && self.committed_account.account.data.len()
-                > Self::COMMIT_STATE_SIZE_THRESHOLD
-            && self.base_account.is_some()
-    }
-
-    pub fn force_commit_state(&mut self) {
-        self.force_commit_state = true;
-    }
-
-    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)
-        }
-    }
-
-    fn create_commit_state_ix(&self, validator: &Pubkey) -> Instruction {
-        let args = CommitStateArgs {
-            nonce: self.commit_id,
-            lamports: self.committed_account.account.lamports,
-            data: self.committed_account.account.data.clone(),
-            allow_undelegation: self.allow_undelegation,
-        };
-        dlp::instruction_builder::commit_state(
-            *validator,
-            self.committed_account.pubkey,
-            self.committed_account.account.owner,
-            args,
-        )
-    }
-
-    fn create_commit_diff_ix(
-        &self,
-        validator: &Pubkey,
-        fetched_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(),
-                self.committed_account.account.data(),
-            )
-            .to_vec(),
-            allow_undelegation: self.allow_undelegation,
-        };
-
-        dlp::instruction_builder::commit_diff(
-            *validator,
-            self.committed_account.pubkey,
-            self.committed_account.account.owner,
-            args,
-        )
-    }
+#[derive(Clone)]
+pub struct CommitDiffTask {
+    pub commit_id: u64,
+    pub allow_undelegation: bool,
+    pub committed_account: CommittedAccount,
+    pub base_account: Account,
 }
 
 #[derive(Clone)]
@@ -397,25 +322,22 @@ mod serialization_safety_test {
         let validator = Pubkey::new_unique();
 
         // Test Commit variant
-        let commit_task: ArgsTask = ArgsTaskType::Commit(
-            TaskBuilderImpl::create_commit_task(
-                123,
-                true,
-                CommittedAccount {
-                    pubkey: Pubkey::new_unique(),
-                    account: Account {
-                        lamports: 1000,
-                        data: vec![1, 2, 3],
-                        owner: Pubkey::new_unique(),
-                        executable: false,
-                        rent_epoch: 0,
-                    },
+        let commit_task: ArgsTask = TaskBuilderImpl::create_commit_task(
+            123,
+            true,
+            CommittedAccount {
+                pubkey: Pubkey::new_unique(),
+                account: Account {
+                    lamports: 1000,
+                    data: vec![1, 2, 3],
+                    owner: Pubkey::new_unique(),
+                    executable: false,
+                    rent_epoch: 0,
                 },
-                &Arc::new(NullTaskInfoFetcher),
-            )
-            .await,
+            },
+            &Arc::new(NullTaskInfoFetcher),
         )
-        .into();
+        .await;
         assert_serializable(&commit_task.instruction(&validator));
 
         // Test Finalize variant