Improve benchmark

entropyxyz · Aug 15, 2024 · 69dd18e · 69dd18e
1 parent fb5a4b7
commit 69dd18e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/pallets/attestation/Cargo.toml b/pallets/attestation/Cargo.toml
@@ -20,6 +20,7 @@ sp-runtime        ={ version="32.0.0", default-features=false }
 sp-staking        ={ version="27.0.0", default-features=false }
 frame-benchmarking={ version="29.0.0", default-features=false, optional=true }
 sp-std            ={ version="14.0.0", default-features=false }
+pallet-session    ={ version="29.0.0", default-features=false, optional=true }
 
 entropy-shared={ version="0.2.0", path="../../crates/shared", features=[
   "wasm-no-std",
@@ -28,8 +29,8 @@ pallet-staking-extension={ version="0.2.0", path="../staking", default-features=
 tdx-quote={ git="https://github.com/entropyxyz/tdx-quote" }
 
 [dev-dependencies]
-pallet-staking                 ={ version="29.0.0", default-features=false }
 pallet-session                 ={ version="29.0.0", default-features=false }
+pallet-staking                 ={ version="29.0.0", default-features=false }
 pallet-balances                ={ version="29.0.0", default-features=false }
 pallet-bags-list               ={ version="28.0.0", default-features=false }
 pallet-timestamp               ={ version="28.0.0", default-features=false }
@@ -42,7 +43,7 @@ rand_core                      ="0.6.4"
 
 [features]
 default=['std']
-runtime-benchmarks=['frame-benchmarking', 'tdx-quote/mock']
+runtime-benchmarks=['frame-benchmarking', 'tdx-quote/mock', 'pallet-session']
 std=[
   'frame-benchmarking/std',
   'frame-support/std',

diff --git a/pallets/attestation/src/benchmarking.rs b/pallets/attestation/src/benchmarking.rs
@@ -22,6 +22,7 @@ const ENCLAVE_SIGNING_KEY: [u8; 32] = [
 use entropy_shared::QuoteInputData;
 use frame_benchmarking::{benchmarks, impl_benchmark_test_suite, whitelisted_caller};
 use frame_system::{EventRecord, RawOrigin};
+use pallet_staking_extension::{ServerInfo, ThresholdServers, ThresholdToStash};
 
 use super::*;
 #[allow(unused)]
@@ -52,6 +53,18 @@ benchmarks! {
 
     // Insert a pending attestation so that this quote is expected
     <PendingAttestations<T>>::insert(attestee.clone(), nonce);
+
+    let stash_account = <T as pallet_session::Config>::ValidatorId::try_from(attestee.clone())
+        .or(Err(()))
+        .unwrap();
+
+    <ThresholdToStash<T>>::insert(attestee.clone(), stash_account.clone());
+    <ThresholdServers<T>>::insert(stash_account.clone(), ServerInfo {
+        tss_account: attestee.clone(),
+        x25519_public_key: [0; 32],
+        endpoint: b"http://localhost:3001".to_vec(),
+    });
+
   }: _(RawOrigin::Signed(attestee.clone()), quote.clone())
   verify {
     assert_last_event::<T>(

diff --git a/pallets/attestation/src/lib.rs b/pallets/attestation/src/lib.rs
@@ -40,6 +40,7 @@ pub mod pallet {
     // pub use crate::weights::WeightInfo;
 
     #[pallet::pallet]
+    #[pallet::without_storage_info]
     pub struct Pallet<T>(_);
 
     #[pallet::config]
@@ -69,6 +70,13 @@ pub mod pallet {
     pub type PendingAttestations<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, [u8; 32], OptionQuery>;
 
+    /// A mapping between block numbers and TSS nodes for who we want to make a request for
+    /// attestation, used to make attestation requests via an offchain worker
+    #[pallet::storage]
+    #[pallet::getter(fn attestation_requests)]
+    pub type AttestationRequests<T: Config> =
+        StorageMap<_, Blake2_128Concat, BlockNumberFor<T>, Vec<Vec<u8>>, OptionQuery>;
+
     #[pallet::event]
     #[pallet::generate_deposit(pub(super) fn deposit_event)]
     pub enum Event<T: Config> {