Skip to content

Commit

Permalink
feat(starknet_batcher): implement sync_block
Browse files Browse the repository at this point in the history
  • Loading branch information
ArniStarkware committed Dec 15, 2024
1 parent e6a1c63 commit ef3461a
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 4 deletions.
39 changes: 35 additions & 4 deletions crates/starknet_batcher/src/batcher.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use blockifier::abi::constants;
Expand All @@ -8,8 +8,10 @@ use chrono::Utc;
use mockall::automock;
use papyrus_storage::state::{StateStorageReader, StateStorageWriter};
use starknet_api::block::{BlockHashAndNumber, BlockNumber};
use starknet_api::core::{ContractAddress, Nonce};
use starknet_api::executable_transaction::Transaction;
use starknet_api::state::ThinStateDiff;
use starknet_api::transaction::TransactionHash;
use starknet_batcher_types::batcher_types::{
BatcherResult,
DecisionReachedInput,
Expand Down Expand Up @@ -353,9 +355,28 @@ impl Batcher {
Ok(GetProposalContentResponse { content: GetProposalContent::Finished(commitment) })
}

// TODO(Arni): Impl add sync block
pub async fn add_sync_block(&mut self, _sync_block: SyncBlock) -> BatcherResult<()> {
todo!("Implement add sync block");
pub async fn add_sync_block(&mut self, sync_block: SyncBlock) -> BatcherResult<()> {
if let Some(height) = self.active_height {
info!("Aborting all work on height {} due to state sync.", height);
self.clear_previous_proposals().await;
self.active_height = None;
}

let SyncBlock { state_diff, transaction_hashes } = sync_block;
let address_to_nonce = state_diff.nonces.iter().map(|(k, v)| (*k, *v)).collect();
let tx_hashes = transaction_hashes.into_iter().collect();

// TODO(Arni): Assert the input sync_block corresponds to this `height`.
let height = match self.get_height_from_storage()?.next() {
Some(height) => Ok(height),
None => {
// We use an explicit match pattern instead of 'ok_or' for improved logging.
info!("Height in storage is at the upper limit of block number.");
Err(BatcherError::InternalError)
}
}?;
info!("Syncing block at height {} and notifying mempool of the block.", height);
self.commit_proposal_and_block(height, state_diff, address_to_nonce, tx_hashes).await
}

#[instrument(skip(self), err)]
Expand All @@ -374,6 +395,16 @@ impl Batcher {
"Committing proposal {} at height {} and notifying mempool of the block.",
proposal_id, height
);
self.commit_proposal_and_block(height, state_diff, address_to_nonce, tx_hashes).await
}

async fn commit_proposal_and_block(
&mut self,
height: BlockNumber,
state_diff: ThinStateDiff,
address_to_nonce: HashMap<ContractAddress, Nonce>,
tx_hashes: HashSet<TransactionHash>,
) -> BatcherResult<()> {
trace!("Transactions: {:#?}, State diff: {:#?}.", tx_hashes, state_diff);
self.storage_writer.commit_proposal(height, state_diff).map_err(|err| {
error!("Failed to commit proposal to storage: {}", err);
Expand Down
40 changes: 40 additions & 0 deletions crates/starknet_batcher/src/batcher_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ use starknet_batcher_types::batcher_types::{
use starknet_batcher_types::errors::BatcherError;
use starknet_mempool_types::communication::MockMempoolClient;
use starknet_mempool_types::mempool_types::CommitBlockArgs;
use starknet_state_sync_types::state_sync_types::SyncBlock;
use tokio::sync::Mutex;

use crate::batcher::{Batcher, MockBatcherStorageReaderTrait, MockBatcherStorageWriterTrait};
Expand Down Expand Up @@ -483,6 +484,41 @@ async fn get_content_from_unknown_proposal() {
assert_eq!(result, Err(BatcherError::ProposalNotFound { proposal_id: PROPOSAL_ID }));
}

#[rstest]
#[tokio::test]
async fn add_sync_block() {
let mut mock_dependencies = MockDependencies::default();

mock_dependencies
.storage_writer
.expect_commit_proposal()
.with(eq(INITIAL_HEIGHT), eq(ThinStateDiff::default()))
.returning(|_, _| Ok(()));

mock_dependencies
.storage_writer
.expect_commit_proposal()
.with(eq(INITIAL_HEIGHT.unchecked_next()), eq(get_state_diff()))
.returning(|_, _| Ok(()));

mock_dependencies
.mempool_client
.expect_commit_block()
.with(eq(CommitBlockArgs {
address_to_nonce: test_contract_nonces(),
tx_hashes: test_tx_hashes(),
}))
.returning(|_| Ok(()));

let mut batcher = create_batcher(mock_dependencies);

let sync_block = SyncBlock {
state_diff: get_state_diff(),
transaction_hashes: test_tx_hashes().into_iter().collect(),
};
batcher.add_sync_block(sync_block).await.unwrap();
}

#[rstest]
#[tokio::test]
async fn decision_reached() {
Expand Down Expand Up @@ -620,3 +656,7 @@ fn test_tx_hashes() -> HashSet<TransactionHash> {
fn test_contract_nonces() -> HashMap<ContractAddress, Nonce> {
HashMap::from_iter((0..3u8).map(|i| (contract_address!(i + 33), nonce!(i + 9))))
}

pub fn get_state_diff() -> ThinStateDiff {
ThinStateDiff { nonces: test_contract_nonces().into_iter().collect(), ..Default::default() }
}

0 comments on commit ef3461a

Please sign in to comment.