Repository: tangle-network/gadget (language: rust)

Issue:
[TASK] Update the manager to handle benchmark mode.



Discussion / hints:
Can you include #138 to this @shekohex? After a benchmark it would be nice to also output what the hardware specs of the machine are.

Fix patch:
diff --git a/blueprint-manager/src/executor/event_handler.rs b/blueprint-manager/src/executor/event_handler.rs
index f137104c5..e5bcc75cb 100644
--- a/blueprint-manager/src/executor/event_handler.rs
+++ b/blueprint-manager/src/executor/event_handler.rs
@@ -379,9 +379,7 @@ pub(crate) async fn handle_tangle_event(
     {
         let mut services_for_this_blueprint = vec![];
         if let Gadget::Native(gadget) = &blueprint.gadget {
-            // TODO: fix typo in soruces -> sources
-            // needs to update the tangle-subxt to fix the typo
-            let gadget_source = &gadget.soruces.0[0];
+            let gadget_source = &gadget.sources.0[0];
             if let gadget_common::tangle_runtime::api::runtime_types::tangle_primitives::services::GadgetSourceFetcher::Github(gh) = &gadget_source.fetcher {
                 let metadata = github_fetcher_to_native_github_metadata(gh, blueprint.blueprint_id);
                 onchain_services.push(metadata);
diff --git a/blueprints/incredible-squaring/src/lib.rs b/blueprints/incredible-squaring/src/lib.rs
index 53a0b5170..6ba2abda7 100644
--- a/blueprints/incredible-squaring/src/lib.rs
+++ b/blueprints/incredible-squaring/src/lib.rs
@@ -1,6 +1,6 @@
 use alloy_primitives::{Bytes, U256};
 use alloy_sol_types::sol;
-use gadget_sdk::job;
+use gadget_sdk::{benchmark, job};
 use std::convert::Infallible;
 use IncredibleSquaringTaskManager::{
     respondToTaskCall, G1Point, G2Point, NonSignerStakesAndSignature, Task, TaskResponse,
@@ -95,3 +95,8 @@ pub fn convert_event_to_inputs(
         quorum_threshold_percentage,
     )
 }
+
+#[benchmark(cores = 1, job_id = 0)]
+fn xsquare() {
+    let _ = xsquare(10).unwrap();
+}
diff --git a/blueprints/incredible-squaring/src/main.rs b/blueprints/incredible-squaring/src/main.rs
index 157d1e914..b8f7cf6e6 100644
--- a/blueprints/incredible-squaring/src/main.rs
+++ b/blueprints/incredible-squaring/src/main.rs
@@ -3,15 +3,16 @@ use alloy_primitives::{Address, FixedBytes};
 use alloy_provider::ProviderBuilder;
 use alloy_signer::k256::ecdsa::SigningKey;
 use alloy_signer_local::PrivateKeySigner;
-use color_eyre::{eyre::eyre, eyre::OptionExt, Result};
+use color_eyre::{eyre, eyre::eyre, eyre::OptionExt, Result};
 use gadget_sdk::{
-    env::Protocol,
+    env::{Protocol, StdGadgetConfiguration},
     events_watcher::{
         evm::{Config, EventWatcher},
+        substrate::SubstrateEventWatcher,
         tangle::TangleEventsWatcher,
-        SubstrateEventWatcher,
     },
     keystore::Backend,
+    run::GadgetRunner,
     tangle_subxt::tangle_testnet_runtime::api::{
         self,
         runtime_types::{sp_core::ecdsa, tangle_primitives::services},
@@ -23,24 +24,25 @@ use incredible_squaring_blueprint::{self as blueprint, IncredibleSquaringTaskMan
 
 use std::sync::Arc;
 
-#[async_trait::async_trait]
-trait GadgetRunner {
-    async fn register(&self) -> Result<()>;
-    async fn run(&self) -> Result<()>;
-}
-
 struct TangleGadgetRunner {
-    env: gadget_sdk::env::GadgetConfiguration,
+    env: StdGadgetConfiguration,
 }
 
 #[async_trait::async_trait]
 impl GadgetRunner for TangleGadgetRunner {
+    type Error = eyre::Report;
+
+    fn env(&self) -> &StdGadgetConfiguration {
+        &self.env
+    }
     async fn register(&self) -> Result<()> {
         let client = self.env.client().await?;
         let signer = self.env.first_signer()?;
+        let ecdsa_signer = self.env.first_ecdsa_signer()?;
+        let ecdsa_key = ecdsa_signer.public_key().to_bytes();
 
         let xt = api::tx().services().register(
-            env.blueprint_id,
+            self.env.blueprint_id,
             services::OperatorPreferences {
                 key: ecdsa::Public(ecdsa_key),
                 approval: services::ApprovalPrefrence::None,
@@ -75,26 +77,40 @@ impl GadgetRunner for TangleGadgetRunner {
 
         Ok(())
     }
+
+    async fn benchmark(&self) -> Result<()> {
+        let env = gadget_sdk::env::load(None)?;
+        let client = env.client().await?;
+        let signer = env.first_signer()?;
+        let xsquare_summary = blueprint::xsquare_benchmark();
+        Ok(())
+    }
 }
 
 struct EigenlayerGadgetRunner {
-    env: gadget_sdk::env::GadgetConfiguration,
+    env: StdGadgetConfiguration,
 }
 
-struct EigenlayerEventWatcher<T> {
-    _phantom: std::marker::PhantomData<T>,
+#[derive(Debug, Default, Clone)]
+struct EigenlayerEventWatcher<C> {
+    _phantom: std::marker::PhantomData<C>,
 }
 
-impl<T: Config> EventWatcher for EigenlayerEventWatcher<T> {
+impl<C: Config> EventWatcher<C> for EigenlayerEventWatcher<C> {
     const TAG: &'static str = "eigenlayer";
     type Contract =
-        IncredibleSquaringTaskManager::IncredibleSquaringTaskManagerInstance<T::T, T::P, T::N>;
+        IncredibleSquaringTaskManager::IncredibleSquaringTaskManagerInstance<C::T, C::P, C::N>;
     type Event = IncredibleSquaringTaskManager::NewTaskCreated;
-    const GENESIS_TX_HASH: FixedBytes<32> = FixedBytes::from_slice(&[0; 32]);
+    const GENESIS_TX_HASH: FixedBytes<32> = FixedBytes([0; 32]);
 }
 
 #[async_trait::async_trait]
 impl GadgetRunner for EigenlayerGadgetRunner {
+    type Error = eyre::Report;
+
+    fn env(&self) -> &StdGadgetConfiguration {
+        &self.env
+    }
     async fn register(&self) -> Result<()> {
         let keystore = self.env.keystore()?;
         let ecdsa_key = keystore
@@ -129,24 +145,23 @@ impl GadgetRunner for EigenlayerGadgetRunner {
             .ok_or_eyre("No ECDSA key found")?;
         let signing_key: SigningKey = SigningKey::from_bytes(&ecdsa_key)?;
         let priv_key_signer: PrivateKeySigner = PrivateKeySigner::from_signing_key(signing_key);
-        let wallet = EthereumWallet::from(signer.clone());
+        let wallet = EthereumWallet::from(priv_key_signer.clone());
         // Set up eignelayer AVS
         let contract_address = Address::from_slice(&[0; 20]);
         // Set up the HTTP provider with the `reqwest` crate.
         let provider = ProviderBuilder::new()
             .with_recommended_fillers()
             .wallet(wallet)
-            .on_http(env.rpc_endpoint);
+            .on_http(env.rpc_endpoint.parse()?);
         let x_square_eigen = blueprint::XsquareEigenEventHandler {};
 
-        let contract = IncredibleSquaringTaskManager::IncredibleSquaringTaskManagerInstance::<
-            T::T,
-            T::P,
-            T::N,
-        >::new(contract_address, provider);
+        let contract = IncredibleSquaringTaskManager::IncredibleSquaringTaskManagerInstance::new(
+            contract_address,
+            provider,
+        );
 
         EventWatcher::run(
-            &EigenlayerEventWatcher,
+            &EigenlayerEventWatcher::default(),
             contract,
             // Add more handler here if we have more functions.
             vec![Box::new(x_square_eigen)],
@@ -155,13 +170,21 @@ impl GadgetRunner for EigenlayerGadgetRunner {
 
         Ok(())
     }
+
+    async fn benchmark(&self) -> Result<()> {
+        let env = gadget_sdk::env::load(Some(Protocol::Eigenlayer))?;
+        let client = env.client().await?;
+        let signer = env.first_signer()?;
+        let xsquare_summary = blueprint::xsquare_benchmark();
+        Ok(())
+    }
 }
 
-fn create_gadget_runner(protocol: Protocol) -> (GadgetConfiguration, Arc<dyn GadgetRunner>) {
+fn create_gadget_runner(protocol: Protocol) -> Arc<dyn GadgetRunner> {
     let env = gadget_sdk::env::load(Some(protocol)).expect("Failed to load environment");
     match protocol {
-        Protocol::Tangle => (env, Arc::new(TangleGadgetRunner { env })),
-        Protocol::Eigenlayer => (env, Arc::new(EigenlayerGadgetRunner { env })),
+        Protocol::Tangle => Arc::new(TangleGadgetRunner { env }),
+        Protocol::Eigenlayer => Arc::new(EigenlayerGadgetRunner { env }),
     }
 }
 
@@ -176,12 +199,16 @@ async fn main() -> Result<()> {
     logger.init();
 
     // Load the environment and create the gadget runner
-    let protocol = Protocol::from_env().unwrap_or(Protocol::Tangle);
-    let (env, runner) = create_gadget_runner(protocol);
+    let protocol = Protocol::from_env();
+    let runner = create_gadget_runner(protocol);
     // Register the operator if needed
-    if env.should_run_registration() {
+    if runner.env().should_run_registration() {
         runner.register().await?;
     }
+    // Run benchmark if needed
+    if runner.env().should_run_benchmarks() {
+        runner.benchmark().await?;
+    }
 
     // Run the gadget / AVS
     runner.run().await?;
diff --git a/flake.lock b/flake.lock
index f65ad7260..31db08b51 100644
--- a/flake.lock
+++ b/flake.lock
@@ -44,11 +44,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1722730825,
-        "narHash": "sha256-X6U+w8qFBuGPCYrZzc9mpN34aRjQ8604MonpBUkj908=",
+        "lastModified": 1724748588,
+        "narHash": "sha256-NlpGA4+AIf1dKNq76ps90rxowlFXUsV9x7vK/mN37JM=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "f3834de3782b82bfc666abf664f946d0e7d1f116",
+        "rev": "a6292e34000dc93d43bccf78338770c1c5ec8a99",
         "type": "github"
       },
       "original": {
@@ -73,11 +73,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1722824458,
-        "narHash": "sha256-2k3/geD5Yh8JT1nrGaRycje5kB0DkvQA/OUZoel1bIU=",
+        "lastModified": 1724811750,
+        "narHash": "sha256-PvhVgQ1rm3gfhK7ts4emprhh/KMkFwXogmgsQ3srR7g=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "a8a937c304e62a5098c6276c9cdf65c19a43b1a5",
+        "rev": "6a1c4915dca7149e7258d8c7f3ac634d8c65f6c6",
         "type": "github"
       },
       "original": {
diff --git a/macros/blueprint-proc-macro/src/benchmark.rs b/macros/blueprint-proc-macro/src/benchmark.rs
index b25582058..12c6601ab 100644
--- a/macros/blueprint-proc-macro/src/benchmark.rs
+++ b/macros/blueprint-proc-macro/src/benchmark.rs
@@ -29,20 +29,24 @@ pub(crate) fn benchmark_impl(args: &BenchmarkArgs, input: &ItemFn) -> syn::Resul
     let block = &input.block;
     let expanded = quote! {
         #[doc(hidden)]
-        pub fn #name() {
+        #[cfg(feature = "gadget")]
+        pub fn #name() -> gadget_sdk::benchmark::BenchmarkSummary {
+            use gadget_sdk::benchmark::*;
+
             let cores: usize = #cores;
-            let rt = gadget_sdk::benchmark::tokio::runtime::Builder::new_multi_thread()
+            let rt = tokio::runtime::Builder::new_multi_thread()
                 .worker_threads(cores)
                 .max_blocking_threads(cores)
                 .enable_all()
                 .build()
-                .unwrap();
+                .expect("build tokio runtime");
             let _guard = rt.enter();
-            let b = gadget_sdk::benchmark::Bencher::new(cores, gadget_sdk::benchmark::TokioRuntime);
-            b.block_on(async move { #block });
+            let b = Bencher::new(cores, TokioRuntime);
+            let _ = b.block_on(async move { #block });
+            // shutdown the runtime.
+            rt.shutdown_background();
             let summary = b.stop(stringify!(#original_name), #job_id);
-            eprintln!("{}", summary);
-            return;
+            return summary;
         }
     };
     Ok(expanded.into())
diff --git a/sdk/src/env.rs b/sdk/src/env.rs
index 761ee8d2f..2c2be8c67 100644
--- a/sdk/src/env.rs
+++ b/sdk/src/env.rs
@@ -9,23 +9,67 @@ pub enum Protocol {
     Eigenlayer,
 }
 
+impl Protocol {
+    /// Returns the protocol from the environment variable `PROTOCOL`.
+    #[cfg(feature = "std")]
+    pub fn from_env() -> Self {
+        std::env::var("PROTOCOL")
+            .map(|v| v.parse::<Protocol>().unwrap_or_default())
+            .unwrap_or_default()
+    }
+
+    /// Returns the protocol from the environment variable `PROTOCOL`.
+    #[cfg(not(feature = "std"))]
+    pub fn from_env() -> Self {
+        Self::Tangle
+    }
+}
+
+impl core::fmt::Display for Protocol {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        match self {
+            Self::Tangle => write!(f, "tangle"),
+            Self::Eigenlayer => write!(f, "eigenlayer"),
+        }
+    }
+}
+
+impl core::str::FromStr for Protocol {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "tangle" => Ok(Self::Tangle),
+            "eigenlayer" => Ok(Self::Eigenlayer),
+            _ => Err(()),
+        }
+    }
+}
+
+/// Gadget environment using the `parking_lot` RwLock.
+#[cfg(feature = "std")]
+pub type StdGadgetConfiguration = GadgetConfiguration<parking_lot::RawRwLock>;
+
 /// Gadget environment.
 #[derive(Debug, Clone)]
 #[non_exhaustive]
 pub struct GadgetConfiguration<RwLock: lock_api::RawRwLock> {
-    /// Tangle RPC endpoint.
+    /// RPC endpoint.
     pub rpc_endpoint: String,
     /// Keystore URI
     ///
     /// * In Memory: `file::memory:` or `:memory:`
     /// * Filesystem: `file:/path/to/keystore` or `file:///path/to/keystore`
     pub keystore_uri: String,
+
     /// Data directory path.
     /// This is used to store the data for the gadget.
     /// If not provided, the gadget is expected to store the data in memory.
     pub data_dir_path: Option<String>,
+
     /// Blueprint ID for this gadget.
     pub blueprint_id: u64,
+
     /// Service ID for this gadget.
     ///
     /// This is only set to `None` when the gadget is in the registration mode.
@@ -40,6 +84,14 @@ pub struct GadgetConfiguration<RwLock: lock_api::RawRwLock> {
     /// If this is set to true, the gadget should do some work and register the operator on the blueprint.
     pub is_registration: bool,
 
+    /// The Current Environment is for the `Benchmarking` of the Gadget
+    ///
+    /// The gadget will now start in the Benchmarking mode and will try to benchmark the current operator on that blueprint
+    /// There is no Service ID for this mode, since we need to benchmark the operator first on the blueprint.
+    ///
+    /// If this is set to true, the gadget should do some work and benchmark the operator on the blueprint.
+    pub is_benchmarking: bool,
+
     /// The type of protocol the gadget is executing on.
     pub protocol: Protocol,
 
@@ -88,7 +140,6 @@ pub enum Error {
     /// No ECDSA keypair found in the keystore.
     #[error("No ECDSA keypair found in the keystore")]
     NoEcdsaKeypair,
-
     /// Invalid ECDSA keypair found in the keystore.
     #[error("Invalid ECDSA keypair found in the keystore")]
     InvalidEcdsaKeypair,
@@ -123,6 +174,7 @@ fn load_inner<RwLock: lock_api::RawRwLock>(
     protocol: Option<Protocol>,
 ) -> Result<GadgetConfiguration<RwLock>, Error> {
     let is_registration = std::env::var("REGISTRATION_MODE_ON").is_ok();
+    let is_benchmarking = std::env::var("BENCHMARKING_MODE_ON").is_ok();
     Ok(GadgetConfiguration {
         rpc_endpoint: std::env::var("RPC_URL").map_err(|_| Error::MissingTangleRpcEndpoint)?,
         keystore_uri: std::env::var("KEYSTORE_URI").map_err(|_| Error::MissingKeystoreUri)?,
@@ -131,8 +183,8 @@ fn load_inner<RwLock: lock_api::RawRwLock>(
             .map_err(|_| Error::MissingBlueprintId)?
             .parse()
             .map_err(Error::MalformedBlueprintId)?,
-        // If the registration mode is on, we don't need the service ID
-        service_id: if is_registration {
+        // If the registration mode is on, or benchmarking mode, we don't need the service ID
+        service_id: if is_registration || is_benchmarking {
             None
         } else {
             Some(
@@ -143,6 +195,7 @@ fn load_inner<RwLock: lock_api::RawRwLock>(
             )
         },
         is_registration,
+        is_benchmarking,
         protocol: protocol.unwrap_or(Protocol::Tangle),
         _lock: core::marker::PhantomData,
     })
@@ -235,6 +288,12 @@ impl<RwLock: lock_api::RawRwLock> GadgetConfiguration<RwLock> {
         self.is_registration
     }
 
+    /// Returns whether the gadget should run in benchmarking mode.
+    #[must_use]
+    pub const fn should_run_benchmarks(&self) -> bool {
+        self.is_benchmarking
+    }
+
     /// Returns a new [`subxt::OnlineClient`] for the Tangle.
     ///
     /// # Errors
diff --git a/sdk/src/keystore/backend/mem.rs b/sdk/src/keystore/backend/mem.rs
index 0137a1691..0f40ffc2f 100644
--- a/sdk/src/keystore/backend/mem.rs
+++ b/sdk/src/keystore/backend/mem.rs
@@ -50,9 +50,8 @@ impl Ord for Bls381PublicWrapper {
 
 /// The In-Memory Keystore Backend.
 ///
-/// It stores everything in memory and does not persist anything. Once dropped,
-/// all the keys are lost, so use with caution.
-///
+/// It stores everything in memory and does not persist anything,
+/// once dropped all the keys are lost, so use with caution.
 /// This is useful for testing and development purposes.
 ///
 /// Note: Cloning this backend is cheap, as it uses [`Arc`] and [`RwLock`] internally.
diff --git a/sdk/src/keystore/mod.rs b/sdk/src/keystore/mod.rs
index 60f49c9dc..7cfb962d3 100644
--- a/sdk/src/keystore/mod.rs
+++ b/sdk/src/keystore/mod.rs
@@ -42,9 +42,10 @@ pub mod sr25519;
 
 pub use error::Error;
 
-/// Functions that a keystore backend must implement
+/// The Keystore [`Backend`] trait.
 ///
-/// This adds support for various cryptographic key operations such as key generation,
+/// it does defines the necessary functions that a keystore backend
+/// must implement to support various cryptographic key operations such as key generation,
 /// signing, and public key retrieval.
 pub trait Backend {
     /// Generate a new sr25519 key pair with an optional seed.
diff --git a/sdk/src/lib.rs b/sdk/src/lib.rs
index 598d2e094..b71e6ebfb 100644
--- a/sdk/src/lib.rs
+++ b/sdk/src/lib.rs
@@ -34,11 +34,17 @@ pub mod tx;
 
 pub use tangle_subxt;
 
+/// Network Module
 #[cfg(feature = "std")]
 pub mod network;
 
 pub mod slashing;
 
+/// Benchmark Module
+#[cfg(feature = "std")]
 pub mod benchmark;
 
+/// Gadget Runner Module
+pub mod run;
+
 pub use gadget_blueprint_proc_macro::*;
diff --git a/sdk/src/network/setup.rs b/sdk/src/network/setup.rs
index baa43808a..df52878ee 100644
--- a/sdk/src/network/setup.rs
+++ b/sdk/src/network/setup.rs
@@ -31,10 +31,10 @@ pub const AGENT_VERSION: &str = "tangle/gadget-sdk/1.0.0";
 /// The version of the client
 pub const CLIENT_VERSION: &str = "1.0.0";
 
-/// The base network configuration for a blueprint's `libp2p` network.
+/// The base network configuration for a blueprint's libp2p network.
 ///
-/// This configuration is used to setup the `libp2p` network for a blueprint.
-/// Construct using [`NetworkConfig::new`] for advanced users or [`NetworkConfig::new_service_network`] ordinarily.
+/// This configuration is used to setup the libp2p network for a blueprint.
+/// Construct using `NetworkConfig::new` for advanced users or `NetworkConfig::new_service_network` ordinarily.
 pub struct NetworkConfig {
     pub identity: libp2p::identity::Keypair,
     pub role_key: ecdsa::Pair,
diff --git a/sdk/src/run.rs b/sdk/src/run.rs
new file mode 100644
index 000000000..28198730f
--- /dev/null
+++ b/sdk/src/run.rs
@@ -0,0 +1,37 @@
+use crate::env::StdGadgetConfiguration;
+
+/// A trait defining the interface for running a gadget.
+///
+/// This trait provides methods for managing the lifecycle of a gadget,
+/// including registration, benchmarking, and execution.
+#[async_trait::async_trait]
+pub trait GadgetRunner {
+    /// The type of error that can be returned by the gadget runner methods.
+    type Error;
+
+    /// Returns a reference to the gadget's standard configuration.
+    ///
+    /// # Returns
+    /// A reference to the `StdGadgetConfiguration` for this gadget.
+    fn env(&self) -> &StdGadgetConfiguration;
+
+    /// Registers the current operator with the the blueprint.
+    ///
+    /// This method should be called only during the registration mode.
+    ///
+    /// # Returns
+    /// A `Result` indicating success or an error if registration fails.
+    async fn register(&self) -> Result<(), Self::Error>;
+
+    /// Performs a benchmark of the gadget's performance.
+    ///
+    /// # Returns
+    /// A `Result` indicating success or an error if benchmarking fails.
+    async fn benchmark(&self) -> Result<(), Self::Error>;
+
+    /// Executes the gadget's main functionality.
+    ///
+    /// # Returns
+    /// A `Result` indicating success or an error if the execution fails.
+    async fn run(&self) -> Result<(), Self::Error>;
+}


Test patch: