Merge pull request risc0#82 from taikoxyz/performance-metrics

Performance metrics

.github/workflows/ci.yml

@@ -58,6 +58,24 @@ jobs:
  cargo test -p sp1-prover --features enable
  shell: bash
 
+ build-test-sgx:
+ name: Build and test sgx
+ runs-on: ubuntu-latest
+ timeout-minutes: 60
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions-rs/toolchain@v1
+ with:
+ profile: minimal
+ toolchain: nightly-2024-01-25
+ override: true
+ components: clippy, rustfmt, rust-src
+ - name: Setup and build
+ run: |
+ cargo build --features sgx
+ cargo test -p sgx-prover --features enable
+ shell: bash
+
  test-lib:
  name: Test raiko-lib
  runs-on: ubuntu-latest

host/src/host/host.rs

@@ -23,6 +23,7 @@ use raiko_lib::{
  BlockProposed, GuestInput, TaikoGuestInput, TaikoProverData,
  },
  taiko_utils::{generate_transactions, to_header},
+ Measurement,
 };
 use raiko_primitives::{
  eip4844::{kzg_to_versioned_hash, MAINNET_KZG_TRUSTED_SETUP},
@@ -31,7 +32,7 @@ use raiko_primitives::{
 use serde::{Deserialize, Serialize};
 use url::Url;
 
-use crate::host::provider_db::ProviderDb;
+use crate::host::provider_db::{MeasuredProviderDb, ProviderDb};
 
 pub fn preflight(
  rpc_url: Option<String>,
@@ -44,12 +45,15 @@ pub fn preflight(
  let http = Http::new(Url::parse(&rpc_url.clone().unwrap()).expect("invalid rpc url"));
  let provider = ProviderBuilder::new().provider(RootProvider::new(RpcClient::new(http, true)));
 
+ let measurement = Measurement::start("Fetching block data...", true);
+
  let block = get_block(&provider, block_no, true).unwrap();
  let parent_block = get_block(&provider, block_no - 1, false).unwrap();
 
  println!("block.hash: {:?}", block.header.hash.unwrap());
+ println!("block header: {:?}", block.header);
  println!("block.parent_hash: {:?}", block.header.parent_hash);
- println!("Block transactions: {:?}", block.transactions.len());
+ println!("block transactions: {:?}", block.transactions.len());
 
  let taiko_guest_input = if network != Network::Ethereum {
  let http_l1 = Http::new(Url::parse(&l1_rpc_url.clone().unwrap()).expect("invalid rpc url"));
@@ -142,14 +146,14 @@ pub fn preflight(
  prover_data,
  }
  } else {
- println!("block header: {:?}", block.header);
  // For Ethereum blocks we just convert the block transactions in a tx_list
  // so that we don't have to supports separate paths.
  TaikoGuestInput {
  tx_list: alloy_rlp::encode(&get_transactions_from_block(&block)),
  ..Default::default()
  }
  };
+ measurement.stop();
 
  let input = GuestInput {
  network,
@@ -186,22 +190,31 @@ pub fn preflight(
  parent_block.header.number.unwrap().try_into().unwrap(),
  );
  let mut builder = BlockBuilder::new(&input)
- .with_db(provider_db)
+ .with_db(MeasuredProviderDb::new(provider_db))
  .prepare_header::<TaikoHeaderPrepStrategy>()?
  .execute_transactions::<TkoTxExecStrategy>()?;
- let provider_db: &mut ProviderDb = builder.mut_db().unwrap();
+ let provider_db = builder.mut_db().unwrap();
+ provider_db.print_report();
+ let provider_db = provider_db.db();
 
- // Construct the state trie and storage from the storage proofs.
  // Gather inclusion proofs for the initial and final state
- let parent_proofs = provider_db.get_initial_proofs()?;
- let proofs = provider_db.get_latest_proofs()?;
+ let measurement = Measurement::start("Fetching storage proofs...", true);
+ let (parent_proofs, proofs) = provider_db.get_proofs()?;
+ measurement.stop();
+
+ // Construct the state trie and storage from the storage proofs.
+ let measurement = Measurement::start("Constructing MPT...", true);
  let (state_trie, storage) =
  proofs_to_tries(input.parent_header.state_root, parent_proofs, proofs)?;
+ measurement.stop();
 
  // Gather proofs for block history
+ let measurement = Measurement::start("Fetching historical block headers...", true);
  let ancestor_headers = provider_db.get_ancestor_headers()?;
+ measurement.stop();
 
  // Get the contracts from the initial db.
+ let measurement = Measurement::start("Fetching contract code...", true);
  let mut contracts = HashSet::new();
  let initial_db = &provider_db.initial_db;
  for account in initial_db.accounts.values() {
@@ -210,6 +223,7 @@ pub fn preflight(
  contracts.insert(code.bytecode.0.clone());
  }
  }
+ measurement.stop();
 
  // Add the collected data to the input
  Ok(GuestInput {

host/src/host/provider_db.rs

@@ -11,10 +11,15 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+use std::{
+ ops::AddAssign,
+ time::{Duration, Instant},
+};
+
 use alloy_consensus::Header as AlloyConsensusHeader;
 use alloy_provider::{Provider, ReqwestProvider};
 use alloy_rpc_types::{BlockId, EIP1186AccountProofResponse};
-use raiko_lib::{mem_db::MemDb, taiko_utils::to_header};
+use raiko_lib::{clear_line, inplace_print, mem_db::MemDb, taiko_utils::to_header};
 use raiko_primitives::{Address, B256, U256};
 use revm::{
  primitives::{Account, AccountInfo, Bytecode, HashMap},
@@ -51,33 +56,44 @@ impl ProviderDb {
  &self.current_db
  }
 
- fn get_proofs(
+ fn get_storage_proofs(
  &mut self,
  block_number: u64,
  storage_keys: HashMap<Address, Vec<U256>>,
+ offset: usize,
+ num_storage_proofs: usize,
  ) -> Result<HashMap<Address, EIP1186AccountProofResponse>, anyhow::Error> {
  let mut storage_proofs = HashMap::new();
+ let mut idx = offset;
  for (address, keys) in storage_keys {
- let indices = keys.into_iter().map(|x| x.to_be_bytes().into()).collect();
+ inplace_print(&format!(
+ "fetching storage proof {idx}/{num_storage_proofs}..."
+ ));
+
+ let indices = keys.iter().map(|x| x.to_be_bytes().into()).collect();
  let proof = self.async_executor.block_on(async {
  self.provider
  .get_proof(address, indices, Some(BlockId::from(block_number)))
  .await
  })?;
  storage_proofs.insert(address, proof);
+ idx += keys.len();
  }
- Ok(storage_proofs)
- }
+ clear_line();
 
- pub fn get_initial_proofs(
- &mut self,
- ) -> Result<HashMap<Address, EIP1186AccountProofResponse>, anyhow::Error> {
- self.get_proofs(self.block_number, self.initial_db.storage_keys())
+ Ok(storage_proofs)
  }
 
- pub fn get_latest_proofs(
+ pub fn get_proofs(
  &mut self,
- ) -> Result<HashMap<Address, EIP1186AccountProofResponse>, anyhow::Error> {
+ ) -> Result<
+ (
+ HashMap<Address, EIP1186AccountProofResponse>,
+ HashMap<Address, EIP1186AccountProofResponse>,
+ ),
+ anyhow::Error,
+ > {
+ // Latest proof keys
  let mut storage_keys = self.initial_db.storage_keys();
  for (address, mut indices) in self.current_db.storage_keys() {
  match storage_keys.get_mut(&address) {
@@ -87,7 +103,32 @@ impl ProviderDb {
  }
  }
  }
- self.get_proofs(self.block_number + 1, storage_keys)
+
+ // Calculate how many storage proofs we need
+ let num_initial_values: usize = self
+ .initial_db
+ .storage_keys()
+ .iter()
+ .map(|(_address, keys)| keys.len())
+ .sum();
+ let num_latest_values: usize = storage_keys.iter().map(|(_address, keys)| keys.len()).sum();
+ let num_storage_proofs = num_initial_values + num_latest_values;
+
+ // Initial proofs
+ let initial_proofs = self.get_storage_proofs(
+ self.block_number,
+ self.initial_db.storage_keys(),
+ 0,
+ num_storage_proofs,
+ )?;
+ let latest_proofs = self.get_storage_proofs(
+ self.block_number + 1,
+ storage_keys,
+ num_initial_values,
+ num_storage_proofs,
+ )?;
+
+ Ok((initial_proofs, latest_proofs))
  }
 
  pub fn get_ancestor_headers(&mut self) -> Result<Vec<AlloyConsensusHeader>, anyhow::Error> {
@@ -205,3 +246,98 @@ impl DatabaseCommit for ProviderDb {
  self.current_db.commit(changes)
  }
 }
+
+pub struct MeasuredProviderDb {
+ pub provider: ProviderDb,
+ pub num_basic: u64,
+ pub time_basic: Duration,
+ pub num_storage: u64,
+ pub time_storage: Duration,
+ pub num_block_hash: u64,
+ pub time_block_hash: Duration,
+ pub num_code_by_hash: u64,
+}
+
+impl MeasuredProviderDb {
+ pub fn new(provider: ProviderDb) -> Self {
+ MeasuredProviderDb {
+ provider,
+ num_basic: 0,
+ time_basic: Duration::default(),
+ num_storage: 0,
+ time_storage: Duration::default(),
+ num_block_hash: 0,
+ time_block_hash: Duration::default(),
+ num_code_by_hash: 0,
+ }
+ }
+
+ pub fn db(&mut self) -> &mut ProviderDb {
+ &mut self.provider
+ }
+
+ pub fn print_report(&self) {
+ println!("db accesses: ");
+ println!(
+ "- account: {}.{} seconds ({} ops)",
+ self.time_basic.as_secs(),
+ self.time_basic.subsec_millis(),
+ self.num_basic
+ );
+ println!(
+ "- storage: {}.{} seconds ({} ops)",
+ self.time_storage.as_secs(),
+ self.time_storage.subsec_millis(),
+ self.num_storage
+ );
+ println!(
+ "- block_hash: {}.{} seconds ({} ops)",
+ self.time_block_hash.as_secs(),
+ self.time_block_hash.subsec_millis(),
+ self.num_block_hash
+ );
+ println!("- code_by_hash: {}", self.num_code_by_hash);
+ }
+}
+
+impl Database for MeasuredProviderDb {
+ type Error = anyhow::Error;
+
+ fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
+ self.num_basic += 1;
+ let start = Instant::now();
+ let res = self.provider.basic(address);
+ self.time_basic
+ .add_assign(Instant::now().duration_since(start));
+ res
+ }
+
+ fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
+ self.num_storage += 1;
+ let start = Instant::now();
+ let res = self.provider.storage(address, index);
+ self.time_storage
+ .add_assign(Instant::now().duration_since(start));
+ res
+ }
+
+ fn block_hash(&mut self, number: U256) -> Result<B256, Self::Error> {
+ self.num_block_hash += 1;
+ let start = Instant::now();
+ let res = self.provider.block_hash(number);
+ self.time_block_hash
+ .add_assign(Instant::now().duration_since(start));
+ res
+ }
+
+ fn code_by_hash(&mut self, _code_hash: B256) -> Result<Bytecode, Self::Error> {
+ self.num_code_by_hash += 1;
+ self.provider.code_by_hash(_code_hash)
+ }
+}
+
+impl DatabaseCommit for MeasuredProviderDb {
+ fn commit(&mut self, changes: HashMap<Address, Account>) {
+ self.provider.commit(changes)
+ }
+}

host/src/metrics.rs

@@ -54,6 +54,6 @@ pub fn _inc_sgx_error(block: u64) {
  SGX_PROOF_ERROR_COUNTER.with(&label).inc();
 }
 
-pub fn observe_input(time: i64) {
+pub fn _observe_input(time: i64) {
  PREPARE_INPUT_TIME.set(time);
 }

host/src/prover/execution.rs

@@ -1,4 +1,4 @@
-use std::{str::FromStr, time::Instant};
+use std::str::FromStr;
 
 use raiko_lib::{
  builder::{BlockBuilderStrategy, TaikoStrategy},
@@ -7,25 +7,26 @@ use raiko_lib::{
  protocol_instance::{assemble_protocol_instance, ProtocolInstance},
  prover::{Prover, ProverResult},
  taiko_utils::HeaderHasher,
+ Measurement,
 };
 use raiko_primitives::B256;
 use serde::{Deserialize, Serialize};
 use tracing::{info, warn};
 
 use super::{context::Context, error::Result, request::ProofRequest};
-use crate::{host::host::preflight, metrics::observe_input, prover::error::HostError};
+use crate::{host::host::preflight, prover::error::HostError};
 
 pub async fn execute<D: Prover>(
  ctx: &mut Context,
  req: &ProofRequest<D::ProofParam>,
 ) -> Result<D::ProofResponse> {
  println!("- {:?}", req);
- // 1. load input data into cache path
- let start = Instant::now();
+ // Generate the witness
+ let measurement = Measurement::start("Generating witness...", false);
  let input = prepare_input(ctx, req).await?;
- let elapsed = Instant::now().duration_since(start).as_millis() as i64;
- observe_input(elapsed);
- // 2. pre-build the block
+ measurement.stop_with("=> Witness generated");
+
+ // 2. Test run the block
  let build_result = TaikoStrategy::build_from(&input);
  let output = match &build_result {
  Ok((header, _mpt_node)) => {
@@ -47,12 +48,15 @@ pub async fn execute<D: Prover>(
  GuestOutput::Failure
  }
  };
- let elapsed = Instant::now().duration_since(start).as_millis() as i64;
- observe_input(elapsed);
 
- D::run(input, output, req.proof_param.clone())
+ // Prove
+ let measurement = Measurement::start("Generating proof...", false);
+ let res = D::run(input, output, req.proof_param.clone())
  .await
- .map_err(|e| HostError::GuestError(e.to_string()))
+ .map_err(|e| HostError::GuestError(e.to_string()));
+ measurement.stop_with("=> Proof generated");
+
+ res
 }
 
 /// prepare input data for provers