feat: add SNOS processing timing instrumentation

crates/core/generate-pie/src/lib.rs

@@ -63,11 +63,13 @@
 // Standard library imports
 use std::path::Path;
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 // External crate imports
 use anyhow::bail;
 use cairo_vm::types::layout_name::LayoutName;
 use futures::future::join_all;
 use log::{info, warn};
+use rpc_client::utils::{reset_rpc_timing, rpc_timing_snapshot};
 use rpc_client::RpcClient;
 use starknet_api::core::OsChainInfo;
 use starknet_os::{
@@ -79,7 +81,7 @@ use tokio::sync::Semaphore;
 use crate::constants::{DEFAULT_MAX_PARALLEL_BLOCKS, MAX_EXECUTION_STEPS_WARNING_THRESHOLD};
 use block_processor::collect_single_block_info;
 use error::PieGenerationError;
-use types::{PieGenerationInput, PieGenerationResult};
+use types::{PieGenerationInput, PieGenerationResult, PieGenerationTiming};
 use utils::sort_abi_entries_for_deprecated_class;
 
 const MAX_PARALLEL_BLOCKS_ENV: &str = "SNOS_MAX_PARALLEL_BLOCKS";
@@ -154,6 +156,8 @@ pub mod types;
 /// }
 /// ```
 pub async fn generate_pie(input: PieGenerationInput) -> Result<PieGenerationResult, PieGenerationError> {
+    reset_rpc_timing();
+    let snos_started_at = Instant::now();
     info!("Starting PIE generation for {} blocks: {:?}", input.blocks.len(), input.blocks);
 
     // Validate input configuration
@@ -299,9 +303,43 @@ pub async fn generate_pie(input: PieGenerationInput) -> Result<PieGenerationResu
         info!("PIE written to file successfully: {}", output_path);
     }
 
+    let total_elapsed = snos_started_at.elapsed();
+    let rpc_timing = rpc_timing_snapshot();
+    let local_processing_elapsed = total_elapsed.saturating_sub(rpc_timing.wait_elapsed);
+    let mut rpc_calls_by_method = rpc_timing.calls_by_method;
+    rpc_calls_by_method.insert("total".to_string(), rpc_timing.calls);
+    let timing = PieGenerationTiming {
+        total_processing_time_ms: duration_millis(total_elapsed),
+        rpc_wait_time_ms: duration_millis(rpc_timing.wait_elapsed),
+        execution_time_ms: duration_millis(local_processing_elapsed),
+        rpc_calls_by_method,
+    };
+    info!(
+        "SNOS processing timing summary for blocks {:?}: total_elapsed={} rpc_wait_elapsed={} local_processing_elapsed={} rpc_calls={} cumulative_rpc_call_elapsed={} rpc_calls_by_method={:?}",
+        input.blocks,
+        format_duration(total_elapsed),
+        format_duration(rpc_timing.wait_elapsed),
+        format_duration(local_processing_elapsed),
+        rpc_timing.calls,
+        format_duration(rpc_timing.cumulative_call_elapsed),
+        timing.rpc_calls_by_method
+    );
     info!("PIE generation completed successfully for blocks {:?}", input.blocks);
 
-    Ok(PieGenerationResult { output, blocks_processed: input.blocks.clone(), output_path: input.output_path.clone() })
+    Ok(PieGenerationResult {
+        output,
+        blocks_processed: input.blocks.clone(),
+        output_path: input.output_path.clone(),
+        timing,
+    })
+}
+
+fn format_duration(duration: Duration) -> String {
+    format!("{:.3}s", duration.as_secs_f64())
+}
+
+fn duration_millis(duration: Duration) -> u64 {
+    duration.as_millis().min(u128::from(u64::MAX)) as u64
 }
 
 pub fn parse_layout(layout: &str) -> anyhow::Result<LayoutName> {

crates/core/generate-pie/src/types/pie.rs

@@ -2,6 +2,7 @@ use blockifier::blockifier_versioned_constants::VersionedConstants;
 use cairo_vm::types::layout_name::LayoutName;
 use starknet_os::io::os_output::StarknetOsRunnerOutput;
 use starknet_types_core::felt::Felt;
+use std::collections::HashMap;
 
 use crate::error::PieGenerationError;
 use crate::types::{ChainConfig, OsHintsConfiguration};
@@ -91,4 +92,19 @@ pub struct PieGenerationResult {
     pub blocks_processed: Vec<u64>,
     /// The output file path where the PIE was saved (if specified).
     pub output_path: Option<String>,
+    /// Timing details for the SNOS run.
+    pub timing: PieGenerationTiming,
+}
+
+/// Wall-clock timing details for a SNOS run.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct PieGenerationTiming {
+    /// Total wall-clock time spent in SNOS processing.
+    pub total_processing_time_ms: u64,
+    /// Wall-clock time spent with at least one RPC call in flight.
+    pub rpc_wait_time_ms: u64,
+    /// Wall-clock time spent on local execution/processing outside RPC waits.
+    pub execution_time_ms: u64,
+    /// RPC calls grouped by method name, including a `total` entry.
+    pub rpc_calls_by_method: HashMap<String, u64>,
 }

crates/core/generate-pie/src/types/proof.rs

@@ -112,6 +112,7 @@ impl ProofCollectionResult {
                     .chain(current_contract_commitment_facts)
                     .map(|(key, value)| (HashOutput(key), value))
                     .collect();
+            let merged_contract_fact_count = global_contract_commitment_facts.len();
 
             let contract_state_commitment_info = CommitmentInfo {
                 previous_root: HashOutput(previous_contract_storage_root),
@@ -120,6 +121,15 @@ impl ProofCollectionResult {
                 commitment_facts: global_contract_commitment_facts,
             };
 
+            info!(
+                "Prepared storage commitment info for block {:?}, contract {:#x}: previous_root={:#x} updated_root={:#x} merged_facts={}",
+                block_id,
+                contract_address,
+                previous_contract_storage_root,
+                current_contract_data.root,
+                merged_contract_fact_count
+            );
+
             address_to_storage_commitment_info.insert(
                 ContractAddress::try_from(contract_address)
                     .map_err(|e| BlockProcessingError::new_custom(format!("Invalid contract address: {:?}", e)))?,
@@ -197,6 +207,7 @@ impl ProofCollectionResult {
             .chain(current_state_commitment_facts)
             .map(|(k, v)| (HashOutput(k), v))
             .collect();
+        let merged_global_state_fact_count = global_state_commitment_facts.len();
 
         let contract_state_commitment_info = CommitmentInfo {
             previous_root: HashOutput(previous_contract_trie_root),
@@ -205,6 +216,18 @@ impl ProofCollectionResult {
             commitment_facts: global_state_commitment_facts,
         };
 
+        info!(
+            "Prepared global contract-state commitment for block {:?}: previous_root={:#x} updated_root={:#x} merged_facts={}",
+            block_id,
+            previous_contract_trie_root,
+            current_contract_trie_root,
+            merged_global_state_fact_count
+        );
+        info!(
+            "Prepared class commitment roots for block {:?}: previous_root={:#x} updated_root={:#x}",
+            block_id, previous_root, updated_root
+        );
+
         // Compute class commitment
         let contract_class_commitment_info =
             compute_class_commitment(&self.previous_class_proofs, &self.class_proofs, previous_root, updated_root);

crates/core/generate-pie/src/utils/rpc.rs

@@ -4,7 +4,7 @@ use blockifier::execution::call_info::CallInfo;
 use blockifier::state::cached_state::StateMaps;
 use blockifier::transaction::objects::TransactionExecutionInfo;
 use cairo_vm::Felt252;
-use log::info;
+use log::{info, warn};
 use rpc_client::client::ProofClient;
 use rpc_client::error::ClientError;
 use rpc_client::types::{ClassProof, ContractData, ContractProof};
@@ -16,6 +16,14 @@ use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 use std::collections::{HashMap, HashSet};
 
+fn summarize_felts(values: &[Felt], limit: usize) -> String {
+    let mut summary: Vec<String> = values.iter().take(limit).map(|felt| format!("{:#x}", felt)).collect();
+    if values.len() > limit {
+        summary.push(format!("... +{}", values.len() - limit));
+    }
+    summary.join(", ")
+}
+
 /// Comprehensive structure that captures all access information from transaction execution
 #[derive(Debug, Clone)]
 pub struct BlockAccessInfo {
@@ -148,7 +156,12 @@ pub(crate) async fn get_class_proofs(
         let proof =
             execute_with_retry(&operation_name, || rpc_client.starknet_rpc().get_class_proof(block_number, class_hash))
                 .await
-                .map_err(ClientError::ProviderError)?;
+                .map_err(|e| {
+                    let message =
+                        format!("class proof request failed for block {block_number} class_hash {class_hash:#x}: {e}");
+                    warn!("{message}");
+                    ClientError::CustomError(message)
+                })?;
         // TODO: need to combine these, similar to merge_chunked_storage_proofs above?
         proofs.insert(**class_hash, proof);
     }
@@ -188,18 +201,37 @@ async fn get_storage_proof_for_contract<KeyIter: Iterator<Item = StorageKey>>(
 
     let contract_data = match &contract_proof.contract_data {
         None => {
+            warn!(
+                "Storage proof for contract {} at block {} returned no contract_data",
+                contract_address, block_number
+            );
             return Ok(contract_proof);
         }
         Some(contract_data) => contract_data,
     };
 
+    info!(
+        "Fetched initial storage proof for contract {} at block {}: root={:#x} storage_proof_sets={} contract_nodes={} requested_keys=[{}]",
+        contract_address,
+        block_number,
+        contract_data.root,
+        contract_data.storage_proofs.len(),
+        contract_proof.contract_proof.len(),
+        summarize_felts(&keys, 8)
+    );
+
     let additional_keys = if contract_data.root != Felt::ZERO {
         contract_data.get_additional_keys(&keys).map_err(|e| ClientError::CustomError(format!("{}", e)))?
     } else {
         vec![]
     };
 
-    info!("Got {} additional keys for contract {}", additional_keys.len(), contract_address);
+    info!(
+        "Got {} additional keys for contract {} [{}]",
+        additional_keys.len(),
+        contract_address,
+        summarize_felts(&additional_keys, 8)
+    );
 
     // Fetch additional proofs required to fill gaps in the storage trie that could make
     // the OS crash otherwise.
@@ -210,7 +242,12 @@ async fn get_storage_proof_for_contract<KeyIter: Iterator<Item = StorageKey>>(
         // Combine all storage proofs into a single vector
         match &additional_proof.contract_data {
             None => {
-                panic!("Failed to fetch additional proof for contract {}", contract_address)
+                let message = format!(
+                    "Additional storage proof for contract {} at block {} returned no contract_data",
+                    contract_address, block_number
+                );
+                warn!("{message}");
+                return Err(ClientError::CustomError(message));
             }
             Some(contract_data) => {
                 additional_proof.contract_data = Some(ContractData {
@@ -223,6 +260,17 @@ async fn get_storage_proof_for_contract<KeyIter: Iterator<Item = StorageKey>>(
         contract_proof = merge_storage_proofs(vec![contract_proof.clone(), additional_proof]);
     }
 
+    if let Some(contract_data) = &contract_proof.contract_data {
+        info!(
+            "Final merged storage proof for contract {} at block {}: root={:#x} storage_proof_sets={} contract_nodes={}",
+            contract_address,
+            block_number,
+            contract_data.root,
+            contract_data.storage_proofs.len(),
+            contract_proof.contract_proof.len()
+        );
+    }
+
     Ok(contract_proof)
 }
 
@@ -235,8 +283,32 @@ async fn fetch_storage_proof_for_contract(
     keys: &[Felt],
     block_number: u64,
 ) -> Result<ContractProof, ClientError> {
-    info!("Fetching storage proof for contract {} with {} keys", contract_address, keys.len());
-    rpc_client.starknet_rpc().get_proof(block_number, contract_address, keys).await.map_err(ClientError::ProviderError)
+    info!(
+        "Fetching storage proof for contract {} with {} keys [{}]",
+        contract_address,
+        keys.len(),
+        summarize_felts(keys, 8)
+    );
+
+    let operation_name = format!(
+        "get_proof(block_number: {block_number}, contract_address: {contract_address:#x}, keys: {})",
+        keys.len()
+    );
+
+    execute_with_retry(&operation_name, || rpc_client.starknet_rpc().get_proof(block_number, contract_address, keys))
+        .await
+        .map_err(|e| {
+            let message = format!(
+                "storage proof request failed for block {} contract {:#x} keys={} [{}]: {}",
+                block_number,
+                contract_address,
+                keys.len(),
+                summarize_felts(keys, 8),
+                e
+            );
+            warn!("{message}");
+            ClientError::CustomError(message)
+        })
 }
 
 /// Merges the storage proofs of the SAME contract.

crates/rpc-client/src/state_reader/mod.rs

@@ -2,7 +2,7 @@ use blockifier::execution::contract_class::{CompiledClassV0, CompiledClassV1, Ru
 use blockifier::state::errors::StateError;
 use blockifier::state::state_api::{StateReader, StateResult};
 use cairo_lang_starknet_classes::contract_class::version_id_from_serialized_sierra_program;
-use log::{debug, warn};
+use log::{debug, info, warn};
 use starknet::core::types::{BlockId, Felt, StarknetError};
 use starknet::providers::{Provider, ProviderError};
 use starknet_api::contract_class::compiled_class_hash::{HashVersion, HashableCompiledClass};
@@ -50,6 +50,26 @@ impl AsyncRpcStateReader {
     }
 }
 
+fn log_cached_state_zero_fallback(
+    field_name: &str,
+    block_id: BlockId,
+    contract_address: ContractAddress,
+    key: Option<StorageKey>,
+    error: &ProviderError,
+) {
+    let key_suffix = key.map(|key| format!(", key={:#x}", Felt::from(*key.0.key()))).unwrap_or_default();
+    let message = format!(
+        "Cached state {field_name} fallback to zero for block {block_id:?}, contract={:#x}{key_suffix}: {error}",
+        Felt::from(*contract_address.key())
+    );
+
+    match error {
+        ProviderError::StarknetError(StarknetError::ContractNotFound) => info!("{message}"),
+        ProviderError::StarknetError(StarknetError::ClassHashNotFound) => info!("{message}"),
+        _ => warn!("{message}"),
+    }
+}
+
 // Helper function to convert provider error to state error
 fn provider_error_to_state_error(provider_error: ProviderError) -> StateError {
     StateError::StateReadError(provider_error.to_string())
@@ -80,7 +100,10 @@ impl AsyncRpcStateReader {
         .await
         {
             Ok(value) => Ok(value.value()),
-            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => Ok(Felt::ZERO),
+            Err(err @ ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                log_cached_state_zero_fallback("storage", block_id, contract_address, Some(key), &err);
+                Ok(Felt::ZERO)
+            }
             Err(e) => Err(provider_error_to_state_error(e)),
         }?;
 
@@ -103,7 +126,10 @@ impl AsyncRpcStateReader {
         .await
         {
             Ok(value) => Ok(value),
-            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => Ok(Felt::ZERO),
+            Err(err @ ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                log_cached_state_zero_fallback("nonce", block_id, contract_address, None, &err);
+                Ok(Felt::ZERO)
+            }
             Err(e) => Err(provider_error_to_state_error(e)),
         }?;
         Ok(Nonce(nonce))
@@ -125,7 +151,10 @@ impl AsyncRpcStateReader {
         .await
         {
             Ok(class_hash) => Ok(class_hash),
-            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => Ok(ClassHash::default().0),
+            Err(err @ ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                log_cached_state_zero_fallback("class_hash", block_id, contract_address, None, &err);
+                Ok(ClassHash::default().0)
+            }
             Err(e) => Err(provider_error_to_state_error(e)),
         }?;