apollo_proof_manager: stress test changes (do not merge)

Author: einat-starkware

Cargo.lock

@@ -1050,6 +1050,15 @@
             "unit": "s"
           }
         },
+        {
+          "title": "Proof Verifications by Component",
+          "description": "The number of proof verifications by component (gateway vs consensus) (over the selected time range)",
+          "type": "stat",
+          "exprs": [
+            "sum by (component) (increase(proof_verification_count{cluster=~\"$cluster\", namespace=~\"$namespace\", pod=~\"$pod\"}[$__range]))"
+          ],
+          "extra_params": {}
+        },
         {
           "title": "Proof Manager Store Latency",
           "description": "The time it takes to store a proof in the proof manager",

crates/apollo_dashboard/resources/dev_grafana.json

@@ -13,7 +13,11 @@ use apollo_gateway::metrics::{
     LABEL_NAME_TX_TYPE as GATEWAY_LABEL_NAME_TX_TYPE,
 };
 use apollo_metrics::metrics::MetricQueryName;
-use apollo_transaction_converter::metrics::PROOF_VERIFICATION_LATENCY;
+use apollo_transaction_converter::metrics::{
+    LABEL_NAME_COMPONENT,
+    PROOF_VERIFICATION_COUNT,
+    PROOF_VERIFICATION_LATENCY,
+};
 
 use crate::dashboard::Row;
 use crate::panel::{Panel, PanelType, Unit};
@@ -148,6 +152,21 @@ fn get_panel_proof_verification_latency() -> Panel {
     .with_unit(Unit::Seconds)
 }
 
+fn get_panel_proof_verification_count_by_component() -> Panel {
+    Panel::new(
+        "Proof Verifications by Component",
+        "The number of proof verifications by component (gateway vs consensus) (over the selected \
+         time range)",
+        sum_by_label(
+            &PROOF_VERIFICATION_COUNT,
+            LABEL_NAME_COMPONENT,
+            DisplayMethod::Increase(RANGE_DURATION),
+            false,
+        ),
+        PanelType::Stat,
+    )
+}
+
 fn get_panel_gateway_proof_manager_store_latency() -> Panel {
     Panel::from_hist(
         &GATEWAY_PROOF_MANAGER_STORE_LATENCY,
@@ -180,6 +199,7 @@ pub(crate) fn get_gateway_row() -> Row {
             get_panel_gateway_validate_stateful_tx_storage_time(),
             get_panel_gateway_validate_stateful_tx_storage_operations(),
             get_panel_proof_verification_latency(),
+            get_panel_proof_verification_count_by_component(),
             get_panel_gateway_proof_manager_store_latency(),
         ],
     )

crates/apollo_dashboard/src/panels/gateway.rs

@@ -249,7 +249,7 @@ impl<
         // validation, to avoid storing proofs for rejected transactions.
         let store_result = self
             .transaction_converter
-            .store_proof_in_proof_manager(proof_facts.clone(), proof.clone())
+            .store_proof_in_proof_manager(proof_facts.clone(), tx_hash, proof.clone())
             .await;
         match store_result {
             Ok(proof_manager_store_duration) => {

crates/apollo_gateway/src/gateway.rs

@@ -192,8 +192,8 @@ impl MockDependencies {
         self.mock_transaction_converter
             .expect_store_proof_in_proof_manager()
             .once()
-            .with(eq(proof_facts), eq(proof))
-            .returning(|_, _| Ok(std::time::Duration::ZERO));
+            .withf(move |pf, _tx_hash, p| *pf == proof_facts && *p == proof)
+            .returning(|_, _, _| Ok(std::time::Duration::ZERO));
     }
 }
 

crates/apollo_gateway/src/gateway_test.rs

@@ -13,21 +13,25 @@ pub type RemoteProofManagerServer =
 impl ComponentRequestHandler<ProofManagerRequest, ProofManagerResponse> for ProofManager {
     async fn handle_request(&mut self, request: ProofManagerRequest) -> ProofManagerResponse {
         match request {
-            ProofManagerRequest::SetProof(proof_facts, proof) => ProofManagerResponse::SetProof(
-                self.set_proof(proof_facts, proof)
-                    .await
-                    .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
-            ),
-            ProofManagerRequest::GetProof(proof_facts) => ProofManagerResponse::GetProof(
-                self.get_proof(proof_facts)
-                    .await
-                    .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
-            ),
-            ProofManagerRequest::ContainsProof(proof_facts) => ProofManagerResponse::ContainsProof(
-                self.contains_proof(proof_facts)
+            ProofManagerRequest::SetProof(proof_facts, tx_hash, proof) => {
+                ProofManagerResponse::SetProof(
+                    self.set_proof(proof_facts, tx_hash, proof)
+                        .await
+                        .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
+                )
+            }
+            ProofManagerRequest::GetProof(proof_facts, tx_hash) => ProofManagerResponse::GetProof(
+                self.get_proof(proof_facts, tx_hash)
                     .await
                     .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
             ),
+            ProofManagerRequest::ContainsProof(proof_facts, tx_hash) => {
+                ProofManagerResponse::ContainsProof(
+                    self.contains_proof(proof_facts, tx_hash)
+                        .await
+                        .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
+                )
+            }
         }
     }
 }

crates/apollo_proof_manager/src/communication.rs

@@ -6,34 +6,35 @@ use apollo_proof_manager_config::config::ProofManagerConfig;
 use async_trait::async_trait;
 use lru::LruCache;
 use starknet_api::transaction::fields::{Proof, ProofFacts};
+use starknet_api::transaction::TransactionHash;
 use starknet_types_core::felt::Felt;
 
 use crate::proof_storage::{FsProofStorage, FsProofStorageError, ProofStorage};
 
-/// In-memory LRU cache for proofs, keyed by the hash of the proof facts.
+/// In-memory LRU cache for proofs, keyed by the (proof facts hash, tx hash) pair.
 #[derive(Clone)]
 pub struct ProofCache {
-    cache: Arc<Mutex<LruCache<Felt, Proof>>>,
+    cache: Arc<Mutex<LruCache<(Felt, TransactionHash), Proof>>>,
 }
 
 impl ProofCache {
     pub fn new(capacity: NonZeroUsize) -> Self {
         Self { cache: Arc::new(Mutex::new(LruCache::new(capacity))) }
     }
 
-    pub fn get(&self, facts_hash: &Felt) -> Option<Proof> {
+    pub fn get(&self, key: &(Felt, TransactionHash)) -> Option<Proof> {
         let mut guard = self.cache.lock().expect("Failed to lock proof cache.");
-        guard.get(facts_hash).cloned()
+        guard.get(key).cloned()
     }
 
-    pub fn insert(&self, facts_hash: Felt, proof: Proof) {
+    pub fn insert(&self, key: (Felt, TransactionHash), proof: Proof) {
         let mut guard = self.cache.lock().expect("Failed to lock proof cache.");
-        guard.put(facts_hash, proof);
+        guard.put(key, proof);
     }
 
-    pub fn contains(&self, facts_hash: &Felt) -> bool {
+    pub fn contains(&self, key: &(Felt, TransactionHash)) -> bool {
         let guard = self.cache.lock().expect("Failed to lock proof cache.");
-        guard.contains(facts_hash)
+        guard.contains(key)
     }
 }
 
@@ -54,45 +55,48 @@ impl ProofManager {
     pub async fn set_proof(
         &self,
         proof_facts: ProofFacts,
+        tx_hash: TransactionHash,
         proof: Proof,
     ) -> Result<(), FsProofStorageError> {
-        if self.contains_proof(proof_facts.clone()).await? {
+        if self.contains_proof(proof_facts.clone(), tx_hash).await? {
             return Ok(());
         }
-        let facts_hash = proof_facts.hash();
-        self.proof_storage.set_proof(facts_hash, proof.clone()).await?;
-        self.cache.insert(facts_hash, proof);
+        let key = (proof_facts.hash(), tx_hash);
+        self.proof_storage.set_proof(key.0, tx_hash, proof.clone()).await?;
+        self.cache.insert(key, proof);
         Ok(())
     }
 
     pub async fn get_proof(
         &self,
         proof_facts: ProofFacts,
+        tx_hash: TransactionHash,
     ) -> Result<Option<Proof>, FsProofStorageError> {
-        let facts_hash = proof_facts.hash();
+        let key = (proof_facts.hash(), tx_hash);
         // Check cache first.
-        if let Some(proof) = self.cache.get(&facts_hash) {
+        if let Some(proof) = self.cache.get(&key) {
             return Ok(Some(proof));
         }
         // Fallback to filesystem.
-        let proof = self.proof_storage.get_proof(facts_hash).await?;
+        let proof = self.proof_storage.get_proof(key.0, tx_hash).await?;
         if let Some(proof) = &proof {
-            self.cache.insert(facts_hash, proof.clone());
+            self.cache.insert(key, proof.clone());
         }
         Ok(proof)
     }
 
     pub async fn contains_proof(
         &self,
         proof_facts: ProofFacts,
+        tx_hash: TransactionHash,
     ) -> Result<bool, FsProofStorageError> {
-        let facts_hash = proof_facts.hash();
+        let key = (proof_facts.hash(), tx_hash);
         // Check cache first.
-        if self.cache.contains(&facts_hash) {
+        if self.cache.contains(&key) {
             return Ok(true);
         }
         // Fallback to filesystem.
-        self.proof_storage.contains_proof(facts_hash).await
+        self.proof_storage.contains_proof(key.0, tx_hash).await
     }
 }
 

crates/apollo_proof_manager/src/proof_manager.rs

@@ -2,6 +2,7 @@ use std::error::Error;
 use std::path::{Path, PathBuf};
 
 use starknet_api::transaction::fields::Proof;
+use starknet_api::transaction::TransactionHash;
 use starknet_types_core::felt::Felt;
 use thiserror::Error;
 use tokio::io::AsyncWriteExt;
@@ -13,9 +14,22 @@ mod proof_storage_test;
 #[async_trait::async_trait]
 pub trait ProofStorage: Send + Sync {
     type Error: Error;
-    async fn set_proof(&self, facts_hash: Felt, proof: Proof) -> Result<(), Self::Error>;
-    async fn get_proof(&self, facts_hash: Felt) -> Result<Option<Proof>, Self::Error>;
-    async fn contains_proof(&self, facts_hash: Felt) -> Result<bool, Self::Error>;
+    async fn set_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+        proof: Proof,
+    ) -> Result<(), Self::Error>;
+    async fn get_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+    ) -> Result<Option<Proof>, Self::Error>;
+    async fn contains_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+    ) -> Result<bool, Self::Error>;
 }
 
 #[derive(Debug, Error)]
@@ -40,27 +54,30 @@ impl FsProofStorage {
         Ok(Self { persistent_root })
     }
 
-    /// Returns the directory that will hold the proof of a certain proof facts hash.
-    /// For a proof facts hash: 0xa1b2c3d4... (rest of hash), the structure is:
+    /// Returns the directory that will hold the proof of a certain proof facts hash and tx hash.
+    /// For a proof facts hash: 0xa1b2c3d4... and tx hash, the structure is:
     /// a1/
     /// └── b2/
     ///     └── a1b2c3d4.../
-    fn get_proof_dir(&self, facts_hash: Felt) -> PathBuf {
-        let facts_hash = hex::encode(facts_hash.to_bytes_be());
-        let (first_msb_byte, second_msb_byte, _rest_of_bytes) =
-            (&facts_hash[..2], &facts_hash[2..4], &facts_hash[4..]);
-        PathBuf::from(first_msb_byte).join(second_msb_byte).join(facts_hash)
+    ///         └── <tx_hash_hex>/
+    fn get_proof_dir(&self, facts_hash: Felt, tx_hash: TransactionHash) -> PathBuf {
+        let facts_hash_hex = hex::encode(facts_hash.to_bytes_be());
+        let (first_msb_byte, second_msb_byte, _rest) =
+            (&facts_hash_hex[..2], &facts_hash_hex[2..4], &facts_hash_hex[4..]);
+        let tx_hash_hex = hex::encode(tx_hash.0.to_bytes_be());
+        PathBuf::from(first_msb_byte).join(second_msb_byte).join(&facts_hash_hex).join(tx_hash_hex)
     }
 
-    fn get_persistent_dir(&self, facts_hash: Felt) -> PathBuf {
-        self.persistent_root.join(self.get_proof_dir(facts_hash))
+    fn get_persistent_dir(&self, facts_hash: Felt, tx_hash: TransactionHash) -> PathBuf {
+        self.persistent_root.join(self.get_proof_dir(facts_hash, tx_hash))
     }
 
     async fn get_persistent_dir_with_create(
         &self,
         facts_hash: Felt,
+        tx_hash: TransactionHash,
     ) -> FsProofStorageResult<PathBuf> {
-        let path = self.get_persistent_dir(facts_hash);
+        let path = self.get_persistent_dir(facts_hash, tx_hash);
         if let Some(parent) = path.parent() {
             tokio::fs::create_dir_all(parent).await?;
         }
@@ -71,9 +88,10 @@ impl FsProofStorage {
     async fn create_tmp_dir(
         &self,
         facts_hash: Felt,
+        tx_hash: TransactionHash,
     ) -> FsProofStorageResult<(tempfile::TempDir, PathBuf)> {
-        // Compute the final persistent directory for this `facts_hash`
-        let persistent_dir = self.get_persistent_dir(facts_hash);
+        // Compute the final persistent directory for this `facts_hash` and `tx_hash`
+        let persistent_dir = self.get_persistent_dir(facts_hash, tx_hash);
         let parent_dir = persistent_dir
             .parent()
             .expect("Proof persistent dir should have a parent")
@@ -106,26 +124,32 @@ impl FsProofStorage {
     }
 
     /// Reads a proof from a file in binary format.
-    async fn read_proof_from_file(&self, facts_hash: Felt) -> FsProofStorageResult<Proof> {
-        let file_path = self.get_persistent_dir(facts_hash).join("proof");
+    async fn read_proof_from_file(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+    ) -> FsProofStorageResult<Proof> {
+        let file_path = self.get_persistent_dir(facts_hash, tx_hash).join("proof");
         let buffer = tokio::fs::read(&file_path).await?;
         Ok(Proof::from(buffer))
     }
 
     async fn write_proof_atomically(
         &self,
         facts_hash: Felt,
+        tx_hash: TransactionHash,
         proof: Proof,
     ) -> FsProofStorageResult<()> {
         // Write proof to a temporary directory.
-        let (_tmp_root, tmp_dir) = self.create_tmp_dir(facts_hash).await?;
+        let (_tmp_root, tmp_dir) = self.create_tmp_dir(facts_hash, tx_hash).await?;
         self.write_proof_to_file(&tmp_dir, &proof).await?;
 
         // Atomically rename directory to persistent one.
         // If a concurrent write already placed the proof at the persistent path, the rename
         // will fail (e.g. ENOTEMPTY on Linux). Since proofs are deterministic for a given
-        // facts_hash, the existing proof is identical and we can safely treat this as success.
-        let persistent_dir = self.get_persistent_dir_with_create(facts_hash).await?;
+        // facts_hash and tx_hash, the existing proof is identical and we can safely treat this as
+        // success.
+        let persistent_dir = self.get_persistent_dir_with_create(facts_hash, tx_hash).await?;
         match tokio::fs::rename(&tmp_dir, &persistent_dir).await {
             Ok(()) => Ok(()),
             Err(_)
@@ -142,16 +166,25 @@ impl FsProofStorage {
 impl ProofStorage for FsProofStorage {
     type Error = FsProofStorageError;
 
-    async fn set_proof(&self, facts_hash: Felt, proof: Proof) -> Result<(), Self::Error> {
-        self.write_proof_atomically(facts_hash, proof).await
+    async fn set_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+        proof: Proof,
+    ) -> Result<(), Self::Error> {
+        self.write_proof_atomically(facts_hash, tx_hash, proof).await
     }
 
-    async fn get_proof(&self, facts_hash: Felt) -> Result<Option<Proof>, Self::Error> {
-        if !self.contains_proof(facts_hash).await? {
+    async fn get_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+    ) -> Result<Option<Proof>, Self::Error> {
+        if !self.contains_proof(facts_hash, tx_hash).await? {
             return Ok(None);
         }
 
-        match self.read_proof_from_file(facts_hash).await {
+        match self.read_proof_from_file(facts_hash, tx_hash).await {
             Ok(proof) => Ok(Some(proof)),
             Err(FsProofStorageError::IoError(e)) if e.kind() == std::io::ErrorKind::NotFound => {
                 Ok(None)
@@ -160,7 +193,11 @@ impl ProofStorage for FsProofStorage {
         }
     }
 
-    async fn contains_proof(&self, facts_hash: Felt) -> Result<bool, Self::Error> {
-        Ok(tokio::fs::try_exists(self.get_persistent_dir(facts_hash)).await?)
+    async fn contains_proof(
+        &self,
+        facts_hash: Felt,
+        tx_hash: TransactionHash,
+    ) -> Result<bool, Self::Error> {
+        Ok(tokio::fs::try_exists(self.get_persistent_dir(facts_hash, tx_hash)).await?)
     }
 }