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;
 
@@ -43,6 +46,37 @@ fn compiled_class_v1_to_casm(class: &CompiledClassV1) -> CasmContractClass {
     }
 }
 
+/// Fetch class from the state reader and contract manager.
+/// Returns error if the class is deprecated.
+fn fetch_class<S>(
+    state_reader_and_contract_manager: Arc<StateReaderAndContractManager<S>>,
+    class_hash: ClassHash,
+) -> Result<(CompiledClassHash, CasmContractClass), ClassesProviderError>
+where
+    S: FetchCompiledClasses + Send + Sync + 'static,
+{
+    let compiled_class = state_reader_and_contract_manager.get_compiled_class(class_hash)?;
+
+    let compiled_class_hash = state_reader_and_contract_manager
+        .get_compiled_class_hash_v2(class_hash, &compiled_class)?;
+
+    match compiled_class {
+        RunnableCompiledClass::V0(_v0) => {
+            Err(ClassesProviderError::DeprecatedContractError(class_hash))
+        }
+        RunnableCompiledClass::V1(compiled_class_v1) => {
+            let casm = compiled_class_v1_to_casm(&compiled_class_v1);
+            Ok((compiled_class_hash, casm))
+        }
+        #[cfg(feature = "cairo_native")]
+        RunnableCompiledClass::V1Native(compiled_class_v1_native) => {
+            let compiled_class_v1 = compiled_class_v1_native.casm();
+            let casm = compiled_class_v1_to_casm(&compiled_class_v1);
+            Ok((compiled_class_hash, casm))
+        }
+    }
+}
+
 /// The classes required for a Starknet OS run.
 /// Matches the fields in `StarknetOsInput`.
 pub struct ClassesInput {
@@ -51,53 +85,46 @@ pub struct ClassesInput {
     pub compiled_classes: BTreeMap<CompiledClassHash, CasmContractClass>,
 }
 
+#[async_trait]
 pub trait ClassesProvider {
     /// 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>;
+}
+
+#[async_trait]
+impl<S> ClassesProvider for Arc<StateReaderAndContractManager<S>>
+where
+    S: FetchCompiledClasses + Send + Sync + 'static,
+{
+    async fn get_classes(
         &self,
         executed_class_hashes: &HashSet<ClassHash>,
     ) -> Result<ClassesInput, ClassesProviderError> {
-        let mut compiled_classes = BTreeMap::new();
+        // clonning the arc to create new refference with static lifetime.
+        let shared_contract_class_manager = self.clone();
 
-        // 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 })
-    }
+        // Creating tasks to fetch classes in parallel.
+        let tasks = executed_class_hashes.iter().map(|&class_hash| {
+            let manager = shared_contract_class_manager.clone();
 
-    /// Fetches class by class hash.
-    fn fetch_class(
-        &self,
-        class_hash: ClassHash,
-    ) -> Result<(CompiledClassHash, CasmContractClass), ClassesProviderError>;
-}
+            tokio::task::spawn_blocking(move || fetch_class(manager, class_hash))
+        });
 
-impl<S: FetchCompiledClasses> ClassesProvider for StateReaderAndContractManager<S> {
-    /// Fetch class from the state reader and contract manager.
-    /// Returns error if the class is deprecated.
-    fn fetch_class(
-        &self,
-        class_hash: ClassHash,
-    ) -> Result<(CompiledClassHash, CasmContractClass), ClassesProviderError> {
-        let compiled_class = self.get_compiled_class(class_hash)?;
-        // TODO(Aviv): Make sure that the state reader is not returning dummy compiled class hash.
-        let compiled_class_hash = self.get_compiled_class_hash_v2(class_hash, &compiled_class)?;
-        match compiled_class {
-            RunnableCompiledClass::V0(_v0) => {
-                Err(ClassesProviderError::DeprecatedContractError(class_hash))
-            }
-            RunnableCompiledClass::V1(compiled_class_v1) => {
-                let casm = compiled_class_v1_to_casm(&compiled_class_v1);
-                Ok((compiled_class_hash, casm))
-            }
-            #[cfg(feature = "cairo_native")]
-            RunnableCompiledClass::V1Native(compiled_class_v1_native) => {
-                let compiled_class_v1 = compiled_class_v1_native.casm();
-                let casm = compiled_class_v1_to_casm(&compiled_class_v1);
-                Ok((compiled_class_hash, casm))
-            }
-        }
+        // Fetching classes in parallel.
+        let results = try_join_all(tasks)
+            .await
+            .map_err(|e| ClassesProviderError::GetClassesError(format!("Task join error: {e}")))?;
+
+        // Collecting results into a BTreeMap.
+        // results is Vec<Result<(CompiledClassHash, CasmContractClass), ClassesProviderError>>
+        let compiled_classes = results
+        .into_iter()
+        .collect::<Result<BTreeMap<CompiledClassHash, CasmContractClass>, ClassesProviderError>>()?;
+
+        Ok(ClassesInput { compiled_classes })
     }
 }

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();