starknet_os_runner: fetch classes concurrent

Author: AvivYossef-starkware

Cargo.lock

@@ -10,11 +10,13 @@ description = "Runs transactions through the Starknet OS and returns Cairo PIE a
 cairo_native = ["blockifier/cairo_native"]
 
 [dependencies]
+async-trait.workspace = true
 blockifier.workspace = true
 blockifier_reexecution.workspace = true
 cairo-lang-starknet-classes.workspace = true
 cairo-lang-utils.workspace = true
 cairo-vm.workspace = true
+futures.workspace = true
 indexmap.workspace = true
 shared_execution_objects.workspace = true
 starknet-rust.workspace = true

crates/starknet_os_runner/Cargo.toml

@@ -1,5 +1,7 @@
 use std::collections::{BTreeMap, HashSet};
+use std::sync::Arc;
 
+use async_trait::async_trait;
 use blockifier::execution::contract_class::{CompiledClassV1, RunnableCompiledClass};
 use blockifier::state::state_api::StateReader;
 use blockifier::state::state_reader_and_contract_manager::{
@@ -9,6 +11,7 @@ use blockifier::state::state_reader_and_contract_manager::{
 use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use cairo_lang_utils::bigint::BigUintAsHex;
 use cairo_vm::types::relocatable::MaybeRelocatable;
+use futures::future::try_join_all;
 use starknet_api::core::{ClassHash, CompiledClassHash};
 use starknet_types_core::felt::Felt;
 
@@ -51,19 +54,31 @@ pub struct ClassesInput {
     pub compiled_classes: BTreeMap<CompiledClassHash, CasmContractClass>,
 }
 
-pub trait ClassesProvider {
+#[async_trait]
+pub trait ClassesProvider: Sized + Clone + Send + Sync + 'static {
     /// Fetches all classes required for the OS run based on the executed class hashes.
-    fn get_classes(
+    /// This default implementation parallelizes fetching by spawning blocking tasks.
+    async fn get_classes(
         &self,
         executed_class_hashes: &HashSet<ClassHash>,
     ) -> Result<ClassesInput, ClassesProviderError> {
-        let mut compiled_classes = BTreeMap::new();
+        // Creating tasks to fetch classes in parallel.
+        let tasks = executed_class_hashes.iter().map(|&class_hash| {
+            let provider = self.clone();
+            tokio::task::spawn_blocking(move || provider.fetch_class(class_hash))
+        });
+
+        // Fetching classes in parallel.
+        // If any task fails, the entire operation fails.
+        let results = try_join_all(tasks)
+            .await
+            .map_err(|e| ClassesProviderError::GetClassesError(format!("Task join error: {e}")))?;
+
+        // Collecting results into a BTreeMap.
+        let compiled_classes = results
+            .into_iter()
+            .collect::<Result<BTreeMap<CompiledClassHash, CasmContractClass>, ClassesProviderError>>()?;
 
-        // TODO(Aviv): Parallelize the fetching of classes.
-        for &class_hash in executed_class_hashes {
-            let (compiled_class_hash, casm) = self.fetch_class(class_hash)?;
-            compiled_classes.insert(compiled_class_hash, casm);
-        }
         Ok(ClassesInput { compiled_classes })
     }
 
@@ -74,7 +89,10 @@ pub trait ClassesProvider {
     ) -> Result<(CompiledClassHash, CasmContractClass), ClassesProviderError>;
 }
 
-impl<S: FetchCompiledClasses> ClassesProvider for StateReaderAndContractManager<S> {
+#[async_trait]
+impl<S: FetchCompiledClasses + Send + Sync + 'static> ClassesProvider
+    for Arc<StateReaderAndContractManager<S>>
+{
     /// Fetch class from the state reader and contract manager.
     /// Returns error if the class is deprecated.
     fn fetch_class(

crates/starknet_os_runner/src/classes_provider.rs

@@ -79,8 +79,8 @@ impl From<VirtualOsBlockInput> for StarknetOsInput {
 
 pub struct Runner<C, S, V>
 where
-    C: ClassesProvider,
-    S: StorageProofProvider,
+    C: ClassesProvider + Sync,
+    S: StorageProofProvider + Sync,
     V: VirtualBlockExecutor,
 {
     pub classes_provider: C,
@@ -90,8 +90,8 @@ where
 
 impl<C, S, V> Runner<C, S, V>
 where
-    C: ClassesProvider,
-    S: StorageProofProvider,
+    C: ClassesProvider + Sync,
+    S: StorageProofProvider + Sync,
     V: VirtualBlockExecutor,
 {
     pub fn new(classes_provider: C, storage_proofs_provider: S, virtual_block_executor: V) -> Self {
@@ -100,7 +100,7 @@ where
 
     /// Creates the OS hints required to run the given transactions virtually
     /// on top of the given block number.
-    pub fn create_os_hints(
+    pub async fn create_os_hints(
         &self,
         block_number: BlockNumber,
         contract_class_manager: ContractClassManager,
@@ -120,12 +120,13 @@ where
             strk_fee_token_address: chain_info.fee_token_addresses.strk_fee_token_address,
         };
 
-        // Fetch classes.
-        let classes = self.classes_provider.get_classes(&execution_data.executed_class_hashes)?;
-
-        // Fetch storage proofs.
-        let storage_proofs =
-            self.storage_proofs_provider.get_storage_proofs(block_number, &execution_data)?;
+        // Fetch classes and storage proofs in parallel.
+        let (classes, storage_proofs) = tokio::join!(
+            self.classes_provider.get_classes(&execution_data.executed_class_hashes),
+            self.storage_proofs_provider.get_storage_proofs(block_number, &execution_data)
+        );
+        let classes = classes?;
+        let storage_proofs = storage_proofs?;
 
         // Convert execution outputs to CentralTransactionExecutionInfo.
         let tx_execution_infos =
@@ -177,13 +178,13 @@ where
     /// 4. Runs the OS in stateless mode (all state pre-loaded in input)
     ///
     /// Returns the OS output containing the Cairo PIE and execution metrics.
-    pub fn run_os(
+    pub async fn run_os(
         &self,
         block_number: BlockNumber,
         contract_class_manager: ContractClassManager,
         txs: Vec<(InvokeTransaction, TransactionHash)>,
     ) -> Result<StarknetOsRunnerOutput, RunnerError> {
-        let os_hints = self.create_os_hints(block_number, contract_class_manager, txs)?;
+        let os_hints = self.create_os_hints(block_number, contract_class_manager, txs).await?;
         let output = run_os_stateless(DEFAULT_OS_LAYOUT, os_hints)?;
         Ok(output)
     }

crates/starknet_os_runner/src/runner.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use async_trait::async_trait;
 use blockifier::state::cached_state::StateMaps;
 use starknet_api::block::BlockNumber;
 use starknet_api::core::{ClassHash, ContractAddress, Nonce};
@@ -32,8 +33,9 @@ use crate::virtual_block_executor::VirtualBlockExecutionData;
 /// The returned `StorageProofs` contains:
 /// - `proof_state`: The ambient state values (nonces, class hashes) discovered in the proof.
 /// - `commitment_infos`: The Patricia Merkle proof nodes for contracts, classes, and storage tries.
+#[async_trait]
 pub trait StorageProofProvider {
-    fn get_storage_proofs(
+    async fn get_storage_proofs(
         &self,
         block_number: BlockNumber,
         execution_data: &VirtualBlockExecutionData,
@@ -288,16 +290,16 @@ impl RpcStorageProofsProvider {
     }
 }
 
+#[async_trait]
 impl StorageProofProvider for RpcStorageProofsProvider {
-    fn get_storage_proofs(
+    async fn get_storage_proofs(
         &self,
         block_number: BlockNumber,
         execution_data: &VirtualBlockExecutionData,
     ) -> Result<StorageProofs, ProofProviderError> {
         let query = Self::prepare_query(execution_data);
 
-        let runtime = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
-        let rpc_proof = runtime.block_on(self.fetch_proofs(block_number, &query))?;
+        let rpc_proof = self.fetch_proofs(block_number, &query).await?;
 
         Self::to_storage_proofs(&rpc_proof, &query)
     }

crates/starknet_os_runner/src/storage_proofs.rs

@@ -50,7 +50,9 @@ fn test_get_storage_proofs_from_rpc(
         prev_base_block_hash: BlockHash::default(),
     };
 
-    let result = rpc_provider.get_storage_proofs(BlockNumber(block_number), &execution_data);
+    let result = runtime.block_on(async {
+        rpc_provider.get_storage_proofs(BlockNumber(block_number), &execution_data).await
+    });
     assert!(result.is_ok(), "Failed to get storage proofs: {:?}", result.err());
 
     let storage_proofs = result.unwrap();