diff --git a/magicblock-committor-service/src/tasks/commit_finalize_task.rs b/magicblock-committor-service/src/tasks/commit_finalize_task.rs index 6a174f582..b189f0c1f 100644 --- a/magicblock-committor-service/src/tasks/commit_finalize_task.rs +++ b/magicblock-committor-service/src/tasks/commit_finalize_task.rs @@ -4,19 +4,12 @@ use dlp_api::{ instruction_builder::{commit_diff_size_budget, commit_size_budget}, AccountSizeClass, }; -use magicblock_committor_program::Chunks; use magicblock_core::intent::CommittedAccount; use solana_account::{Account, ReadableAccount}; use solana_instruction::Instruction; use solana_pubkey::Pubkey; -use crate::{ - consts::MAX_WRITE_CHUNK_SIZE, - tasks::{ - commit_task::{CommitBufferStage, CommitDelivery}, - BaseTask, BaseTaskImpl, PreparationTask, - }, -}; +use crate::tasks::{commit_task::CommitDelivery, BaseTask, BaseTaskImpl}; /// A task that commits a delegated account's state to the base layer and finalizes it in the same /// instruction. @@ -100,32 +93,6 @@ impl CommitFinalizeTask { .0 } - pub fn stage(&self) -> Option<&CommitBufferStage> { - match &self.delivery { - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } - | CommitDelivery::StateInBuffer { stage } => Some(stage), - CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { - None - } - } - } - - pub fn stage_mut(&mut self) -> Option<&mut CommitBufferStage> { - match &mut self.delivery { - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } - | CommitDelivery::StateInBuffer { stage } => Some(stage), - CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { - None - } - } - } - pub fn is_buffer(&self) -> bool { matches!( self.delivery, @@ -134,57 +101,15 @@ impl CommitFinalizeTask { ) } - pub fn state_preparation_stage(&self) -> CommitBufferStage { - let data = self.committed_account.account.data.clone(); - self.preparation_stage(data) - } - - fn diff_preparation_stage(&self, base_data: &[u8]) -> CommitBufferStage { - let diff = - compute_diff(base_data, &self.committed_account.account.data) - .to_vec(); - self.preparation_stage(diff) - } - - fn preparation_stage(&self, buffer_data: Vec) -> CommitBufferStage { - let chunks = - Chunks::from_data_length(buffer_data.len(), MAX_WRITE_CHUNK_SIZE); - CommitBufferStage::Preparation(PreparationTask { - commit_id: self.commit_id, - pubkey: self.committed_account.pubkey, - buffer_data, - chunks, - }) - } - pub fn reset_commit_id(&mut self, commit_id: u64) { self.commit_id = commit_id; - let new_stage = match &self.delivery { - CommitDelivery::StateInBuffer { .. } => { - self.state_preparation_stage() - } - CommitDelivery::DiffInBuffer { - base_account, - stage: _, - } => { - let slice = base_account.data.as_slice(); - self.diff_preparation_stage(slice) - } - _ => return, - }; - match &mut self.delivery { - CommitDelivery::StateInBuffer { stage } => { - *stage = new_stage; - } - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } => { - *stage = new_stage; + CommitDelivery::StateInBuffer { prepared } + | CommitDelivery::DiffInBuffer { prepared, .. } => { + *prepared = false } _ => {} - } + }; } } @@ -214,15 +139,14 @@ impl BaseTask for CommitFinalizeTask { std::mem::replace(&mut self.delivery, CommitDelivery::StateInArgs); match delivery { CommitDelivery::StateInArgs => { - let stage = self.state_preparation_stage(); - self.delivery = CommitDelivery::StateInBuffer { stage }; + self.delivery = + CommitDelivery::StateInBuffer { prepared: false }; true } CommitDelivery::DiffInArgs { base_account } => { - let stage = self.diff_preparation_stage(base_account.data()); self.delivery = CommitDelivery::DiffInBuffer { base_account, - stage, + prepared: false, }; true } diff --git a/magicblock-committor-service/src/tasks/commit_stage_task.rs b/magicblock-committor-service/src/tasks/commit_stage_task.rs new file mode 100644 index 000000000..cc6d4156b --- /dev/null +++ b/magicblock-committor-service/src/tasks/commit_stage_task.rs @@ -0,0 +1,295 @@ +use dlp_api::diff::compute_diff; +use magicblock_committor_program::{ + instruction_builder::{ + close_buffer::{create_close_ix, CreateCloseIxArgs}, + init_buffer::{create_init_ix, CreateInitIxArgs}, + realloc_buffer::{ + create_realloc_buffer_ixs, CreateReallocBufferIxArgs, + }, + write_buffer::{create_write_ix, CreateWriteIxArgs}, + }, + pdas, ChangesetChunks, Chunks, +}; +use magicblock_program::Pubkey; +use solana_instruction::Instruction; + +use crate::{ + consts::MAX_WRITE_CHUNK_SIZE, + tasks::{ + commit_finalize_task::CommitFinalizeTask, + commit_task::{CommitDelivery, CommitTask}, + }, +}; + +#[derive(Debug)] +pub struct PreparationTask<'a> { + pub commit_id: u64, + pub pubkey: Pubkey, + pub chunks: Chunks, + buffer_data: Vec, + prepared: &'a mut bool, +} + +impl<'a> PreparationTask<'a> { + pub fn from_commit(task: &'a mut CommitTask) -> Option { + match &mut task.delivery_details { + CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { + None + } + CommitDelivery::StateInBuffer { prepared } => { + let buffer_data = task.committed_account.account.data.clone(); + let chunks = Chunks::from_data_length( + buffer_data.len(), + MAX_WRITE_CHUNK_SIZE, + ); + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + buffer_data, + chunks, + prepared, + }) + } + CommitDelivery::DiffInBuffer { + base_account, + prepared, + } => { + let diff = compute_diff( + base_account.data.as_ref(), + &task.committed_account.account.data, + ) + .to_vec(); + let chunks = + Chunks::from_data_length(diff.len(), MAX_WRITE_CHUNK_SIZE); + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + buffer_data: diff, + chunks, + prepared, + }) + } + } + } + + pub fn from_commit_finalize( + task: &'a mut CommitFinalizeTask, + ) -> Option { + match &mut task.delivery { + CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { + None + } + CommitDelivery::StateInBuffer { prepared } => { + let buffer_data = task.committed_account.account.data.clone(); + let chunks = Chunks::from_data_length( + buffer_data.len(), + MAX_WRITE_CHUNK_SIZE, + ); + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + buffer_data, + chunks, + prepared, + }) + } + CommitDelivery::DiffInBuffer { + base_account, + prepared, + } => { + let diff = compute_diff( + base_account.data.as_ref(), + &task.committed_account.account.data, + ) + .to_vec(); + let chunks = + Chunks::from_data_length(diff.len(), MAX_WRITE_CHUNK_SIZE); + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + buffer_data: diff, + chunks, + prepared, + }) + } + } + } + + /// Returns initialization [`Instruction`] + pub fn init_instruction(&self, authority: &Pubkey) -> Instruction { + // // SAFETY: as object_length internally uses only already allocated or static buffers, + // // and we don't use any fs writers, so the only error that may occur here is of kind + // // OutOfMemory or WriteZero. This is impossible due to: + // // Chunks::new panics if its size exceeds MAX_ACCOUNT_ALLOC_PER_INSTRUCTION_SIZE or 10_240 + // // https://github.com/near/borsh-rs/blob/f1b75a6b50740bfb6231b7d0b1bd93ea58ca5452/borsh/src/ser/helpers.rs#L59 + let chunks_account_size = + borsh::object_length(&self.chunks).unwrap() as u64; + let buffer_account_size = self.buffer_data.len() as u64; + + let (instruction, _, _) = create_init_ix(CreateInitIxArgs { + authority: *authority, + pubkey: self.pubkey, + chunks_account_size, + buffer_account_size, + commit_id: self.commit_id, + chunk_count: self.chunks.count(), + chunk_size: self.chunks.chunk_size(), + }); + + instruction + } + + /// Returns compute units required for realloc instruction + pub fn init_compute_units(&self) -> u32 { + 12_000 + } + + /// Returns realloc instruction required for Buffer preparation + #[allow(clippy::let_and_return)] + pub fn realloc_instructions(&self, authority: &Pubkey) -> Vec { + let buffer_account_size = self.buffer_data.len() as u64; + let realloc_instructions = + create_realloc_buffer_ixs(CreateReallocBufferIxArgs { + authority: *authority, + pubkey: self.pubkey, + buffer_account_size, + commit_id: self.commit_id, + }); + + realloc_instructions + } + + /// Returns compute units required for realloc instruction + pub fn realloc_compute_units(&self) -> u32 { + 6_000 + } + + /// Returns realloc instruction required for Buffer preparation + #[allow(clippy::let_and_return)] + pub fn write_instructions(&self, authority: &Pubkey) -> Vec { + let chunks_iter = + ChangesetChunks::new(&self.chunks, self.chunks.chunk_size()) + .iter(&self.buffer_data); + let write_instructions = chunks_iter + .map(|chunk| { + create_write_ix(CreateWriteIxArgs { + authority: *authority, + pubkey: self.pubkey, + offset: chunk.offset, + data_chunk: chunk.data_chunk, + commit_id: self.commit_id, + }) + }) + .collect::>(); + + write_instructions + } + + pub fn write_compute_units(&self, bytes_count: usize) -> u32 { + const PER_BYTE: u32 = 3; + + u32::try_from(bytes_count) + .ok() + .and_then(|bytes_count| bytes_count.checked_mul(PER_BYTE)) + .unwrap_or(u32::MAX) + } + + pub fn chunks_pda(&self, authority: &Pubkey) -> Pubkey { + pdas::chunks_pda( + authority, + &self.pubkey, + self.commit_id.to_le_bytes().as_slice(), + ) + .0 + } + + pub fn buffer_pda(&self, authority: &Pubkey) -> Pubkey { + pdas::buffer_pda( + authority, + &self.pubkey, + self.commit_id.to_le_bytes().as_slice(), + ) + .0 + } + + pub fn cleanup_task(&self) -> CleanupTask { + CleanupTask { + pubkey: self.pubkey, + commit_id: self.commit_id, + } + } + + pub fn done(self) { + *self.prepared = true; + } +} + +#[derive(Clone, Debug)] +pub struct CleanupTask { + pub pubkey: Pubkey, + pub commit_id: u64, +} + +impl CleanupTask { + pub fn from_commit(task: &CommitTask) -> Option { + match &task.delivery_details { + CommitDelivery::StateInBuffer { prepared: true } + | CommitDelivery::DiffInBuffer { prepared: true, .. } => { + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + }) + } + _ => None, + } + } + + pub fn from_commit_finalize(task: &CommitFinalizeTask) -> Option { + match &task.delivery { + CommitDelivery::StateInBuffer { prepared: true } + | CommitDelivery::DiffInBuffer { prepared: true, .. } => { + Some(Self { + commit_id: task.commit_id, + pubkey: task.committed_account.pubkey, + }) + } + _ => None, + } + } + + pub fn instruction(&self, authority: &Pubkey) -> Instruction { + create_close_ix(CreateCloseIxArgs { + authority: *authority, + pubkey: self.pubkey, + commit_id: self.commit_id, + }) + } + + /// Returns compute units required to execute [`crate::tasks::CleanupTask`] + pub fn compute_units(&self) -> u32 { + 30_000 + } + + /// Returns a number of [`CleanupTask`]s that is possible to fit in single + pub const fn max_tx_fit_count_with_budget() -> usize { + 8 + } + + pub fn chunks_pda(&self, authority: &Pubkey) -> Pubkey { + pdas::chunks_pda( + authority, + &self.pubkey, + self.commit_id.to_le_bytes().as_slice(), + ) + .0 + } + + pub fn buffer_pda(&self, authority: &Pubkey) -> Pubkey { + pdas::buffer_pda( + authority, + &self.pubkey, + self.commit_id.to_le_bytes().as_slice(), + ) + .0 + } +} diff --git a/magicblock-committor-service/src/tasks/commit_task.rs b/magicblock-committor-service/src/tasks/commit_task.rs index 4d4ee77d1..439999827 100644 --- a/magicblock-committor-service/src/tasks/commit_task.rs +++ b/magicblock-committor-service/src/tasks/commit_task.rs @@ -3,25 +3,12 @@ use dlp_api::{ diff::compute_diff, AccountSizeClass, }; -use magicblock_committor_program::Chunks; use magicblock_core::intent::CommittedAccount; use solana_account::{Account, ReadableAccount}; use solana_instruction::Instruction; use solana_pubkey::Pubkey; -use crate::{ - consts::MAX_WRITE_CHUNK_SIZE, - tasks::{BaseTask, BaseTaskImpl, CleanupTask, PreparationTask}, -}; - -/// Lifecycle stage of a buffer used for commit delivery. -/// Tracks whether the on-chain buffer still needs to be initialized -/// or is ready to be cleaned up after a successful commit. -#[derive(Clone, Debug)] -pub enum CommitBufferStage { - Preparation(PreparationTask), - Cleanup(CleanupTask), -} +use crate::tasks::{BaseTask, BaseTaskImpl}; /// Describes how commit data is delivered to the base layer. /// @@ -32,14 +19,14 @@ pub enum CommitBufferStage { pub enum CommitDelivery { StateInArgs, StateInBuffer { - stage: CommitBufferStage, + prepared: bool, }, DiffInArgs { base_account: Account, }, DiffInBuffer { base_account: Account, - stage: CommitBufferStage, + prepared: bool, }, } @@ -139,32 +126,6 @@ impl CommitTask { ) } - pub fn stage(&self) -> Option<&CommitBufferStage> { - match &self.delivery_details { - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } - | CommitDelivery::StateInBuffer { stage } => Some(stage), - CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { - None - } - } - } - - pub fn stage_mut(&mut self) -> Option<&mut CommitBufferStage> { - match &mut self.delivery_details { - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } - | CommitDelivery::StateInBuffer { stage } => Some(stage), - CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => { - None - } - } - } - pub fn is_buffer(&self) -> bool { matches!( self.delivery_details, @@ -173,57 +134,15 @@ impl CommitTask { ) } - pub fn state_preparation_stage(&self) -> CommitBufferStage { - let committed_data = self.committed_account.account.data.clone(); - self.preparation_stage(committed_data) - } - - fn diff_preparation_stage(&self, base_data: &[u8]) -> CommitBufferStage { - let diff = - compute_diff(base_data, &self.committed_account.account.data) - .to_vec(); - self.preparation_stage(diff) - } - - fn preparation_stage(&self, buffer_data: Vec) -> CommitBufferStage { - let chunks = - Chunks::from_data_length(buffer_data.len(), MAX_WRITE_CHUNK_SIZE); - CommitBufferStage::Preparation(PreparationTask { - commit_id: self.commit_id, - pubkey: self.committed_account.pubkey, - buffer_data, - chunks, - }) - } - pub fn reset_commit_id(&mut self, commit_id: u64) { self.commit_id = commit_id; - let new_stage = match &self.delivery_details { - CommitDelivery::StateInBuffer { .. } => { - self.state_preparation_stage() - } - CommitDelivery::DiffInBuffer { - base_account, - stage: _, - } => { - let slice = base_account.data.as_slice(); - self.diff_preparation_stage(slice) - } - _ => return, - }; - match &mut self.delivery_details { - CommitDelivery::StateInBuffer { stage } => { - *stage = new_stage; - } - CommitDelivery::DiffInBuffer { - base_account: _, - stage, - } => { - *stage = new_stage; + CommitDelivery::StateInBuffer { prepared } + | CommitDelivery::DiffInBuffer { prepared, .. } => { + *prepared = false } _ => {} - } + }; } } @@ -254,15 +173,14 @@ impl BaseTask for CommitTask { ); match details { CommitDelivery::StateInArgs => { - let stage = self.state_preparation_stage(); - self.delivery_details = CommitDelivery::StateInBuffer { stage }; + self.delivery_details = + CommitDelivery::StateInBuffer { prepared: false }; true } CommitDelivery::DiffInArgs { base_account } => { - let stage = self.diff_preparation_stage(base_account.data()); self.delivery_details = CommitDelivery::DiffInBuffer { base_account, - stage, + prepared: false, }; true } diff --git a/magicblock-committor-service/src/tasks/mod.rs b/magicblock-committor-service/src/tasks/mod.rs index 7640ff74e..31b8186eb 100644 --- a/magicblock-committor-service/src/tasks/mod.rs +++ b/magicblock-committor-service/src/tasks/mod.rs @@ -1,15 +1,4 @@ use dlp_api::{args::CallHandlerArgs, AccountSizeClass}; -use magicblock_committor_program::{ - instruction_builder::{ - close_buffer::{create_close_ix, CreateCloseIxArgs}, - init_buffer::{create_init_ix, CreateInitIxArgs}, - realloc_buffer::{ - create_realloc_buffer_ixs, CreateReallocBufferIxArgs, - }, - write_buffer::{create_write_ix, CreateWriteIxArgs}, - }, - pdas, ChangesetChunks, Chunks, -}; use magicblock_core::intent::BaseActionCallback; use magicblock_metrics::metrics::LabelValue; use magicblock_program::magic_scheduled_base_intent::BaseAction; @@ -17,6 +6,7 @@ use solana_instruction::{AccountMeta, Instruction}; use solana_pubkey::Pubkey; pub mod commit_finalize_task; +pub mod commit_stage_task; pub mod commit_task; pub mod task_builder; pub mod task_strategist; @@ -385,166 +375,6 @@ impl From for BaseActionTask { } } -#[derive(Clone, Debug)] -pub struct PreparationTask { - pub commit_id: u64, - pub pubkey: Pubkey, - pub chunks: Chunks, - - // TODO(edwin): replace with reference once done - pub buffer_data: Vec, -} - -impl PreparationTask { - /// Returns initialization [`Instruction`] - pub fn init_instruction(&self, authority: &Pubkey) -> Instruction { - // // SAFETY: as object_length internally uses only already allocated or static buffers, - // // and we don't use any fs writers, so the only error that may occur here is of kind - // // OutOfMemory or WriteZero. This is impossible due to: - // // Chunks::new panics if its size exceeds MAX_ACCOUNT_ALLOC_PER_INSTRUCTION_SIZE or 10_240 - // // https://github.com/near/borsh-rs/blob/f1b75a6b50740bfb6231b7d0b1bd93ea58ca5452/borsh/src/ser/helpers.rs#L59 - let chunks_account_size = - borsh::object_length(&self.chunks).unwrap() as u64; - let buffer_account_size = self.buffer_data.len() as u64; - - let (instruction, _, _) = create_init_ix(CreateInitIxArgs { - authority: *authority, - pubkey: self.pubkey, - chunks_account_size, - buffer_account_size, - commit_id: self.commit_id, - chunk_count: self.chunks.count(), - chunk_size: self.chunks.chunk_size(), - }); - - instruction - } - - /// Returns compute units required for realloc instruction - pub fn init_compute_units(&self) -> u32 { - 12_000 - } - - /// Returns realloc instruction required for Buffer preparation - #[allow(clippy::let_and_return)] - pub fn realloc_instructions(&self, authority: &Pubkey) -> Vec { - let buffer_account_size = self.buffer_data.len() as u64; - let realloc_instructions = - create_realloc_buffer_ixs(CreateReallocBufferIxArgs { - authority: *authority, - pubkey: self.pubkey, - buffer_account_size, - commit_id: self.commit_id, - }); - - realloc_instructions - } - - /// Returns compute units required for realloc instruction - pub fn realloc_compute_units(&self) -> u32 { - 6_000 - } - - /// Returns realloc instruction required for Buffer preparation - #[allow(clippy::let_and_return)] - pub fn write_instructions(&self, authority: &Pubkey) -> Vec { - let chunks_iter = - ChangesetChunks::new(&self.chunks, self.chunks.chunk_size()) - .iter(&self.buffer_data); - let write_instructions = chunks_iter - .map(|chunk| { - create_write_ix(CreateWriteIxArgs { - authority: *authority, - pubkey: self.pubkey, - offset: chunk.offset, - data_chunk: chunk.data_chunk, - commit_id: self.commit_id, - }) - }) - .collect::>(); - - write_instructions - } - - pub fn write_compute_units(&self, bytes_count: usize) -> u32 { - const PER_BYTE: u32 = 3; - - u32::try_from(bytes_count) - .ok() - .and_then(|bytes_count| bytes_count.checked_mul(PER_BYTE)) - .unwrap_or(u32::MAX) - } - - pub fn chunks_pda(&self, authority: &Pubkey) -> Pubkey { - pdas::chunks_pda( - authority, - &self.pubkey, - self.commit_id.to_le_bytes().as_slice(), - ) - .0 - } - - pub fn buffer_pda(&self, authority: &Pubkey) -> Pubkey { - pdas::buffer_pda( - authority, - &self.pubkey, - self.commit_id.to_le_bytes().as_slice(), - ) - .0 - } - - pub fn cleanup_task(&self) -> CleanupTask { - CleanupTask { - pubkey: self.pubkey, - commit_id: self.commit_id, - } - } -} - -#[derive(Clone, Debug)] -pub struct CleanupTask { - pub pubkey: Pubkey, - pub commit_id: u64, -} - -impl CleanupTask { - pub fn instruction(&self, authority: &Pubkey) -> Instruction { - create_close_ix(CreateCloseIxArgs { - authority: *authority, - pubkey: self.pubkey, - commit_id: self.commit_id, - }) - } - - /// Returns compute units required to execute [`CleanupTask`] - pub fn compute_units(&self) -> u32 { - 30_000 - } - - /// Returns a number of [`CleanupTask`]s that is possible to fit in single - pub const fn max_tx_fit_count_with_budget() -> usize { - 8 - } - - pub fn chunks_pda(&self, authority: &Pubkey) -> Pubkey { - pdas::chunks_pda( - authority, - &self.pubkey, - self.commit_id.to_le_bytes().as_slice(), - ) - .0 - } - - pub fn buffer_pda(&self, authority: &Pubkey) -> Pubkey { - pdas::buffer_pda( - authority, - &self.pubkey, - self.commit_id.to_le_bytes().as_slice(), - ) - .0 - } -} - #[cfg(test)] mod serialization_safety_test { @@ -556,7 +386,8 @@ mod serialization_safety_test { use crate::{ tasks::{ - commit_task::{CommitBufferStage, CommitDelivery, CommitTask}, + commit_stage_task::PreparationTask, + commit_task::{CommitDelivery, CommitTask}, *, }, test_utils, @@ -669,9 +500,8 @@ mod serialization_safety_test { ) -> CommitTask { let task = make_commit_task(commit_id, allow_undelegation, data, lamports); - let stage = task.state_preparation_stage(); CommitTask { - delivery_details: CommitDelivery::StateInBuffer { stage }, + delivery_details: CommitDelivery::StateInBuffer { prepared: false }, ..task } } @@ -690,11 +520,11 @@ mod serialization_safety_test { fn test_preparation_instructions_serialization() { let authority = Pubkey::new_unique(); - let commit_task = + let mut commit_task = make_buffer_commit_task(789, true, vec![0; 1024], 3000); - let Some(CommitBufferStage::Preparation(preparation_task)) = - commit_task.stage() + let Some(preparation_task) = + PreparationTask::from_commit(&mut commit_task) else { panic!("invalid preparation state on creation!"); }; @@ -723,6 +553,7 @@ fn test_close_buffer_limit() { use tracing::info; use crate::{ + tasks::commit_stage_task::CleanupTask, test_utils, transactions::{ serialized_transaction_size, MAX_TRANSACTION_WIRE_SIZE, diff --git a/magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs b/magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs index f5d7a8c9e..792ae9ee2 100644 --- a/magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs +++ b/magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs @@ -30,8 +30,9 @@ use tracing::{error, info}; use crate::{ persist::{CommitStatus, IntentPersister}, tasks::{ - commit_task::CommitBufferStage, task_strategist::TransactionStrategy, - BaseTaskImpl, CleanupTask, PreparationTask, + commit_stage_task::{CleanupTask, PreparationTask}, + task_strategist::TransactionStrategy, + BaseTaskImpl, }, utils::persist_status_update, ComputeBudgetConfig, @@ -100,17 +101,16 @@ impl DeliveryPreparator { task: &mut BaseTaskImpl, persister: &Option

, ) -> DeliveryPreparatorResult<(), InternalError> { - let stage = match task { - BaseTaskImpl::Commit(commit_task) => commit_task.stage_mut(), + let preparation_task = match task { + BaseTaskImpl::Commit(commit_task) => { + PreparationTask::from_commit(commit_task) + } BaseTaskImpl::CommitFinalize(commit_finalize_task) => { - commit_finalize_task.stage_mut() + PreparationTask::from_commit_finalize(commit_finalize_task) } _ => None, }; - let Some(stage) = stage else { - return Ok(()); - }; - let CommitBufferStage::Preparation(preparation_task) = stage else { + let Some(preparation_task) = preparation_task else { return Ok(()); }; @@ -124,7 +124,7 @@ impl DeliveryPreparator { ); // Initialize buffer account. Init + reallocs - self.initialize_buffer_account(authority, preparation_task) + self.initialize_buffer_account(authority, &preparation_task) .await?; // Persist initialization success @@ -137,7 +137,7 @@ impl DeliveryPreparator { ); // Writing chunks with some retries - self.write_buffer_with_retries(authority, preparation_task) + self.write_buffer_with_retries(authority, &preparation_task) .await?; // Persist that buffer account initiated successfully let update_status = CommitStatus::BufferAndChunkFullyInitialized; @@ -148,8 +148,7 @@ impl DeliveryPreparator { update_status, ); - let cleanup_task = preparation_task.cleanup_task(); - *stage = CommitBufferStage::Cleanup(cleanup_task); + preparation_task.done(); Ok(()) } @@ -175,29 +174,26 @@ impl DeliveryPreparator { res => return res, } - // Prepare cleanup task - set stage to Cleanup before calling cleanup - let stage = match task { - BaseTaskImpl::Commit(commit_task) => commit_task.stage_mut(), + // Preparation failed due to buffer existing - cleanup and retry + let preparation_task = match task { + BaseTaskImpl::Commit(commit_task) => { + PreparationTask::from_commit(commit_task) + } BaseTaskImpl::CommitFinalize(commit_finalize_task) => { - commit_finalize_task.stage_mut() + PreparationTask::from_commit_finalize(commit_finalize_task) } _ => None, }; - let Some(stage) = stage else { + let Some(preparation_task) = preparation_task else { return Ok(()); }; - let CommitBufferStage::Preparation(preparation_task) = stage else { - return Ok(()); - }; - let preparation_task = preparation_task.clone(); - let cleanup_task = preparation_task.cleanup_task(); + // Cleanup + let cleanup_task = preparation_task.cleanup_task(); self.cleanup_buffers(authority, &[cleanup_task]).await?; self.rpc_client.invalidate_cached_blockhash().await; - // Restore preparation stage for retry - *stage = CommitBufferStage::Preparation(preparation_task); - + // Retry preparation self.prepare_task(authority, task, persister).await } @@ -206,7 +202,7 @@ impl DeliveryPreparator { async fn initialize_buffer_account( &self, authority: &Keypair, - preparation_task: &PreparationTask, + preparation_task: &PreparationTask<'_>, ) -> DeliveryPreparatorResult<(), BufferExecutionError> { let authority_pubkey = authority.pubkey(); let init_instruction = @@ -259,7 +255,7 @@ impl DeliveryPreparator { async fn write_buffer_with_retries( &self, authority: &Keypair, - preparation_task: &PreparationTask, + preparation_task: &PreparationTask<'_>, ) -> DeliveryPreparatorResult<(), InternalError> { let authority_pubkey = authority.pubkey(); let write_instructions = diff --git a/magicblock-committor-service/src/transaction_preparator/mod.rs b/magicblock-committor-service/src/transaction_preparator/mod.rs index cc60e955c..54e9ce844 100644 --- a/magicblock-committor-service/src/transaction_preparator/mod.rs +++ b/magicblock-committor-service/src/transaction_preparator/mod.rs @@ -9,7 +9,7 @@ use solana_pubkey::Pubkey; use crate::{ persist::IntentPersister, tasks::{ - commit_task::CommitBufferStage, task_strategist::TransactionStrategy, + commit_stage_task::CleanupTask, task_strategist::TransactionStrategy, utils::TransactionUtils, BaseTaskImpl, }, transaction_preparator::{ @@ -127,15 +127,11 @@ impl TransactionPreparator for TransactionPreparatorImpl { let cleanup_tasks: Vec<_> = tasks .iter() .filter_map(|task| match task { - BaseTaskImpl::Commit(commit_task) => commit_task.stage(), - BaseTaskImpl::CommitFinalize(commit_finalize_task) => { - commit_finalize_task.stage() + BaseTaskImpl::Commit(commit_task) => { + CleanupTask::from_commit(commit_task) } - _ => None, - }) - .filter_map(|stage| match stage { - CommitBufferStage::Cleanup(cleanup_task) => { - Some(cleanup_task.clone()) + BaseTaskImpl::CommitFinalize(commit_finalize_task) => { + CleanupTask::from_commit_finalize(commit_finalize_task) } _ => None, }) diff --git a/test-integration/test-committor-service/tests/common.rs b/test-integration/test-committor-service/tests/common.rs index 66f896211..db64b1d8b 100644 --- a/test-integration/test-committor-service/tests/common.rs +++ b/test-integration/test-committor-service/tests/common.rs @@ -245,9 +245,8 @@ pub fn create_commit_task(data: &[u8]) -> CommitTask { #[allow(dead_code)] pub fn create_buffer_commit_task(data: &[u8]) -> CommitTask { let task = create_commit_task(data); - let stage = task.state_preparation_stage(); CommitTask { - delivery_details: CommitDelivery::StateInBuffer { stage }, + delivery_details: CommitDelivery::StateInBuffer { prepared: false }, ..task } } diff --git a/test-integration/test-committor-service/tests/test_delivery_preparator.rs b/test-integration/test-committor-service/tests/test_delivery_preparator.rs index 403cd6f7e..354905212 100644 --- a/test-integration/test-committor-service/tests/test_delivery_preparator.rs +++ b/test-integration/test-committor-service/tests/test_delivery_preparator.rs @@ -3,7 +3,8 @@ use magicblock_committor_program::Chunks; use magicblock_committor_service::{ persist::IntentPersisterImpl, tasks::{ - commit_task::{CommitBufferStage, CommitDelivery}, + commit_stage_task::CleanupTask, + commit_task::CommitDelivery, task_strategist::{TaskStrategist, TransactionStrategy}, BaseTaskImpl, }, @@ -45,8 +46,7 @@ async fn test_prepare_10kb_buffer() { else { panic!("unexpected task type"); }; - let Some(CommitBufferStage::Cleanup(cleanup_task)) = commit_task.stage() - else { + let Some(cleanup_task) = CleanupTask::from_commit(commit_task) else { panic!("unexpected CommitStage"); }; @@ -115,11 +115,9 @@ async fn test_prepare_multiple_buffers() { .optimized_tasks .iter() .filter_map(|el| match el { - BaseTaskImpl::Commit(commit_task) => commit_task.stage(), - _ => None, - }) - .filter_map(|stage| match stage { - CommitBufferStage::Cleanup(cleanup_task) => Some(cleanup_task), + BaseTaskImpl::Commit(commit_task) => { + CleanupTask::from_commit(commit_task) + } _ => None, }) .collect(); @@ -233,7 +231,7 @@ async fn test_already_initialized_error_handled() { let BaseTaskImpl::Commit(ref ct) = strategy.optimized_tasks[0] else { panic!("unexpected task type"); }; - let Some(CommitBufferStage::Cleanup(cleanup_task)) = ct.stage() else { + let Some(cleanup_task) = CleanupTask::from_commit(ct) else { panic!("unexpected CommitStage"); }; // Check buffer account exists @@ -252,9 +250,8 @@ async fn test_already_initialized_error_handled() { commit_task.committed_account.account.data.len() - 2, ); commit_task.committed_account.account.data = data.clone(); - commit_task.delivery_details = CommitDelivery::StateInBuffer { - stage: commit_task.state_preparation_stage(), - }; + commit_task.delivery_details = + CommitDelivery::StateInBuffer { prepared: false }; let mut strategy = TransactionStrategy { optimized_tasks: vec![commit_task.into()], lookup_tables_keys: vec![], @@ -275,7 +272,7 @@ async fn test_already_initialized_error_handled() { let BaseTaskImpl::Commit(ref ct) = strategy.optimized_tasks[0] else { panic!("unexpected task type"); }; - let Some(CommitBufferStage::Cleanup(cleanup_task)) = ct.stage() else { + let Some(cleanup_task) = CleanupTask::from_commit(ct) else { panic!("unexpected CommitStage"); }; @@ -334,11 +331,7 @@ async fn test_prepare_cleanup_and_reprepare_mixed_tasks() { .optimized_tasks .iter() .filter_map(|t| match t { - BaseTaskImpl::Commit(ct) => ct.stage(), - _ => None, - }) - .filter_map(|stage| match stage { - CommitBufferStage::Cleanup(c) => Some(c), + BaseTaskImpl::Commit(ct) => CleanupTask::from_commit(ct), _ => None, }) .collect(); @@ -405,12 +398,10 @@ async fn test_prepare_cleanup_and_reprepare_mixed_tasks() { } // Rebuild buffer stages with mutated data - commit_a.delivery_details = CommitDelivery::StateInBuffer { - stage: commit_a.state_preparation_stage(), - }; - commit_b.delivery_details = CommitDelivery::StateInBuffer { - stage: commit_b.state_preparation_stage(), - }; + commit_a.delivery_details = + CommitDelivery::StateInBuffer { prepared: false }; + commit_b.delivery_details = + CommitDelivery::StateInBuffer { prepared: false }; // --- Step 4: re-prepare with the same logical tasks (same commit IDs, mutated data) --- let mut strategy2 = TransactionStrategy { @@ -441,11 +432,7 @@ async fn test_prepare_cleanup_and_reprepare_mixed_tasks() { .optimized_tasks .iter() .filter_map(|t| match t { - BaseTaskImpl::Commit(ct) => ct.stage(), - _ => None, - }) - .filter_map(|stage| match stage { - CommitBufferStage::Cleanup(c) => Some(c), + BaseTaskImpl::Commit(ct) => CleanupTask::from_commit(ct), _ => None, }) .collect(); diff --git a/test-integration/test-committor-service/tests/test_transaction_preparator.rs b/test-integration/test-committor-service/tests/test_transaction_preparator.rs index 207f3fc48..341385363 100644 --- a/test-integration/test-committor-service/tests/test_transaction_preparator.rs +++ b/test-integration/test-committor-service/tests/test_transaction_preparator.rs @@ -3,7 +3,7 @@ use magicblock_committor_program::Chunks; use magicblock_committor_service::{ persist::IntentPersisterImpl, tasks::{ - commit_task::CommitBufferStage, + commit_stage_task::CleanupTask, task_strategist::{TaskStrategist, TransactionStrategy}, utils::TransactionUtils, BaseActionTask, BaseActionTaskV1, BaseTaskImpl, FinalizeTask, @@ -154,9 +154,7 @@ async fn test_prepare_commit_tx_with_multiple_accounts() { BaseTaskImpl::Commit(ct) => ct, _ => continue, }; - let Some(CommitBufferStage::Cleanup(cleanup_task)) = - commit_task.stage() - else { + let Some(cleanup_task) = CleanupTask::from_commit(commit_task) else { continue; }; let chunks_pda = cleanup_task.chunks_pda(&fixture.authority.pubkey()); @@ -250,9 +248,7 @@ async fn test_prepare_commit_tx_with_base_actions() { BaseTaskImpl::Commit(ct) => ct, _ => continue, }; - let Some(CommitBufferStage::Cleanup(cleanup_task)) = - commit_task.stage() - else { + let Some(cleanup_task) = CleanupTask::from_commit(commit_task) else { continue; }; let chunks_pda = cleanup_task.chunks_pda(&fixture.authority.pubkey());