refactor: decouple preparation/cleanup tasks from BaseTasks#1382
Open
taco-paco wants to merge 2 commits into
Open
refactor: decouple preparation/cleanup tasks from BaseTasks#1382taco-paco wants to merge 2 commits into
taco-paco wants to merge 2 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR replaces the Changes
Sequence Diagram(s)sequenceDiagram
participant DeliveryPreparator
participant PreparationTask
participant CleanupTask
DeliveryPreparator->>PreparationTask: from_commit / from_commit_finalize
DeliveryPreparator->>PreparationTask: initialize_buffer_account
DeliveryPreparator->>PreparationTask: write_buffer_with_retries
DeliveryPreparator->>PreparationTask: done()
DeliveryPreparator->>CleanupTask: cleanup on AccountAlreadyInitializedError
Suggested reviewers: ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@magicblock-committor-service/src/tasks/commit_stage_task.rs`:
- Around line 34-116: The constructors in PreparationTask::from_commit and
PreparationTask::from_commit_finalize are duplicating the same matching and
buffer/chunk setup logic, with only the task wrapper and delivery field name
differing. Refactor this shared behavior into a small trait over both CommitTask
and CommitFinalizeTask that exposes commit_id, pubkey, committed_account data,
and mutable delivery access, then have both constructors delegate to one
implementation. Apply the same pattern to CleanupTask::from_commit and
CleanupTask::from_commit_finalize so the commit/finalize paths stay consistent
and easier to maintain.
In
`@magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs`:
- Around line 177-196: The retry cleanup path in
delivery_preparator::prepare_task is rebuilding a full PreparationTask just to
obtain a CleanupTask, which repeats the BaseTaskImpl matching and can trigger
expensive buffer cloning/compute_diff work. Change this branch to derive
CleanupTask directly from the existing task variants used in prepare_task, using
only the commit_id and pubkey needed for cleanup, and remove the intermediate
PreparationTask::from_commit/from_commit_finalize construction before calling
cleanup_buffers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 53bc9cb4-b5d2-4dfa-a247-906785281c37
📒 Files selected for processing (9)
magicblock-committor-service/src/tasks/commit_finalize_task.rsmagicblock-committor-service/src/tasks/commit_stage_task.rsmagicblock-committor-service/src/tasks/commit_task.rsmagicblock-committor-service/src/tasks/mod.rsmagicblock-committor-service/src/transaction_preparator/delivery_preparator.rsmagicblock-committor-service/src/transaction_preparator/mod.rstest-integration/test-committor-service/tests/common.rstest-integration/test-committor-service/tests/test_delivery_preparator.rstest-integration/test-committor-service/tests/test_transaction_preparator.rs
Comment on lines
+34
to
+116
| pub fn from_commit(task: &'a mut CommitTask) -> Option<Self> { | ||
| 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<Self> { | ||
| 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, | ||
| }) | ||
| } | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Duplicate logic between CommitTask and CommitFinalizeTask handlers.
PreparationTask::from_commit/from_commit_finalize (Lines 34-116) are near-identical (only differing in the task type and field name delivery_details vs delivery), and the same duplication repeats for CleanupTask::from_commit/from_commit_finalize (Lines 233-258). Consider a small trait abstracting commit_id(), pubkey(), committed_data(), and delivery_mut()/delivery() over both task types so both constructors can share one implementation.
♻️ Sketch of a shared-trait refactor
+trait DeliveryHolder {
+ fn commit_id(&self) -> u64;
+ fn pubkey(&self) -> Pubkey;
+ fn committed_data(&self) -> &[u8];
+ fn delivery_mut(&mut self) -> &mut CommitDelivery;
+ fn delivery(&self) -> &CommitDelivery;
+}
+
+impl DeliveryHolder for CommitTask { /* map to delivery_details */ }
+impl DeliveryHolder for CommitFinalizeTask { /* map to delivery */ }
+
impl<'a> PreparationTask<'a> {
- pub fn from_commit(task: &'a mut CommitTask) -> Option<Self> { ... }
- pub fn from_commit_finalize(task: &'a mut CommitFinalizeTask) -> Option<Self> { ... }
+ pub fn from_task<T: DeliveryHolder>(task: &'a mut T) -> Option<Self> {
+ let commit_id = task.commit_id();
+ let pubkey = task.pubkey();
+ let data = task.committed_data().to_vec();
+ match task.delivery_mut() {
+ CommitDelivery::StateInArgs | CommitDelivery::DiffInArgs { .. } => None,
+ CommitDelivery::StateInBuffer { prepared } => Some(Self { commit_id, pubkey, chunks: Chunks::from_data_length(data.len(), MAX_WRITE_CHUNK_SIZE), buffer_data: data, prepared }),
+ CommitDelivery::DiffInBuffer { base_account, prepared } => { /* compute diff against base_account.data */ }
+ }
+ }
}Same pattern applies to CleanupTask::from_commit/from_commit_finalize.
Note this touches code that may be further relocated to magic-program per the PR's stated plan; weigh against near-term rework.
Also applies to: 233-258
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-committor-service/src/tasks/commit_stage_task.rs` around lines 34
- 116, The constructors in PreparationTask::from_commit and
PreparationTask::from_commit_finalize are duplicating the same matching and
buffer/chunk setup logic, with only the task wrapper and delivery field name
differing. Refactor this shared behavior into a small trait over both CommitTask
and CommitFinalizeTask that exposes commit_id, pubkey, committed_account data,
and mutable delivery access, then have both constructors delegate to one
implementation. Apply the same pattern to CleanupTask::from_commit and
CleanupTask::from_commit_finalize so the commit/finalize paths stay consistent
and easier to maintain.
Comment on lines
+177
to
+196
| // 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 |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Avoid rebuilding a full PreparationTask just to get a CleanupTask.
This block duplicates the BaseTaskImpl → task-extraction match from prepare_task (lines 104-113), and worse, PreparationTask::from_commit/from_commit_finalize eagerly clone buffer_data and (for DiffInBuffer) run compute_diff over the full account data — all just to discard everything except commit_id/pubkey via cleanup_task(). For large committed accounts (tests use up to 500KB) this is wasted CPU/memory on every AccountAlreadyInitializedError retry.
Since CleanupTask only needs commit_id and pubkey, construct it directly from the task without going through PreparationTask:
♻️ Proposed fix
- // 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) => {
- PreparationTask::from_commit_finalize(commit_finalize_task)
- }
- _ => None,
- };
- let Some(preparation_task) = preparation_task else {
- return Ok(());
- };
-
- // Cleanup
- let cleanup_task = preparation_task.cleanup_task();
+ // Preparation failed due to buffer existing - cleanup and retry
+ let cleanup_task = match task {
+ BaseTaskImpl::Commit(commit_task) => CleanupTask {
+ commit_id: commit_task.commit_id,
+ pubkey: commit_task.committed_account.pubkey,
+ },
+ BaseTaskImpl::CommitFinalize(commit_finalize_task) => CleanupTask {
+ commit_id: commit_finalize_task.commit_id,
+ pubkey: commit_finalize_task.committed_account.pubkey,
+ },
+ _ => return Ok(()),
+ };
self.cleanup_buffers(authority, &[cleanup_task]).await?;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // 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 | |
| // Preparation failed due to buffer existing - cleanup and retry | |
| let cleanup_task = match task { | |
| BaseTaskImpl::Commit(commit_task) => CleanupTask { | |
| commit_id: commit_task.commit_id, | |
| pubkey: commit_task.committed_account.pubkey, | |
| }, | |
| BaseTaskImpl::CommitFinalize(commit_finalize_task) => CleanupTask { | |
| commit_id: commit_finalize_task.commit_id, | |
| pubkey: commit_finalize_task.committed_account.pubkey, | |
| }, | |
| _ => return Ok(()), | |
| }; | |
| self.cleanup_buffers(authority, &[cleanup_task]).await?; | |
| self.rpc_client.invalidate_cached_blockhash().await; | |
| // Retry preparation |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@magicblock-committor-service/src/transaction_preparator/delivery_preparator.rs`
around lines 177 - 196, The retry cleanup path in
delivery_preparator::prepare_task is rebuilding a full PreparationTask just to
obtain a CleanupTask, which repeats the BaseTaskImpl matching and can trigger
expensive buffer cloning/compute_diff work. Change this branch to derive
CleanupTask directly from the existing task variants used in prepare_task, using
only the commit_id and pubkey needed for cleanup, and remove the intermediate
PreparationTask::from_commit/from_commit_finalize construction before calling
cleanup_buffers.
taco-paco
requested review from
GabrielePicco and
snawaz
and removed request for
snawaz
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Breaking Changes
Test Plan
Summary by CodeRabbit
New Features
Bug Fixes
Tests