feat: profiling/flamegraph

.github/workflows/integrations.yml

@@ -60,10 +60,6 @@ jobs:
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Build MPC Recovery Binary Locally
-        run: |
-          cargo build -p mpc-recovery --release
-
       - name: Test
         run: cargo test -p mpc-recovery-integration-tests --jobs 1 -- --test-threads 1
         env:

Cargo.toml

@@ -5,3 +5,7 @@ members = [
     "load-tests",
     "test-oidc-provider",
 ]
+
+[profile.flamegraph]
+inherits = "release"
+debug = true

integration-tests/Cargo.toml

@@ -22,14 +22,15 @@ near-crypto = "0.17"
 near-jsonrpc-client = "0.6"
 near-primitives = "0.17"
 near-units = "0.2.0"
+nix = { version = "0.27", features = ["signal"] }
 once_cell = "1"
 rand = "0.7"
 serde = "1"
 serde_json = "1"
 testcontainers = { version = "0.14", features = ["experimental"] }
 tokio = { version = "1.28", features = ["full"] }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-near-workspaces =  "0.8.0"
+near-workspaces = "0.8.0"
 toml = "0.8.1"
 
 [dev-dependencies]
@@ -44,3 +45,4 @@ reqwest = "0.11.16"
 [features]
 default = []
 docker-test = []
+flamegraph = ["mpc-recovery/disable-open-telemetry"]

integration-tests/README.md

@@ -16,12 +16,6 @@ Build OIDC Provider test image
 docker build -t near/test-oidc-provider ./test-oidc-provider
 ```
 
-Now, build mpc-recovery from the project's root:
-
-```BASH
-cargo build --release
-```
-
 Then run the integration tests:
 
 ```BASH
@@ -44,6 +38,22 @@ Finally, run the integration tests with the built docker image:
 cargo test -p mpc-recovery-integration-tests --features docker-test
 ```
 
+## Profiling: Flamegraphs
+
+To profile code and get a flamegraph, run the following:
+
+```sh
+cargo flamegraph --root --profile flamegraph --test lib
+```
+
+Or for a singular test like `test_basic_action`:
+
+```sh
+cargo flamegraph --root --profile flamegraph --test lib -- test_basic_action
+```
+
+This will generate a `flamegraph.svg`. Open this on a browser and inspect each of the callstacks.
+
 ## FAQ
 
 ### I want to run a test, but keep the docker containers from being destroyed

integration-tests/src/env/local.rs

@@ -1,12 +1,12 @@
 use aes_gcm::aead::consts::U32;
 use aes_gcm::aead::generic_array::GenericArray;
-use async_process::Child;
 use mpc_recovery::firewall::allowed::DelegateActionRelayer;
 use mpc_recovery::logging;
 use mpc_recovery::relayer::NearRpcAndRelayerClient;
 use multi_party_eddsa::protocols::ExpandedKeyPair;
 
 use crate::env::{LeaderNodeApi, SignerNodeApi};
+use crate::mpc::{self, NodeProcess};
 use crate::util;
 
 pub struct SignerNode {
@@ -19,8 +19,7 @@ pub struct SignerNode {
     gcp_datastore_url: String,
 
     // process held so it's not dropped. Once dropped, process will be killed.
-    #[allow(unused)]
-    process: Child,
+    _process: NodeProcess,
 }
 
 impl SignerNode {
@@ -51,11 +50,10 @@ impl SignerNode {
             gcp_datastore_url: Some(ctx.datastore.local_address.clone()),
             jwt_signature_pk_url: ctx.oidc_provider.jwt_pk_local_url.clone(),
             logging_options: logging::Options::default(),
-        }
-        .into_str_args();
+        };
 
-        let sign_node_id = format!("sign/{node_id}");
-        let process = util::spawn_mpc(ctx.release, &sign_node_id, &args)?;
+        let sign_node_id = format!("sign-{node_id}");
+        let process = mpc::spawn(ctx.release, &sign_node_id, args).await?;
         let address = format!("http://127.0.0.1:{web_port}");
         tracing::info!("Signer node is starting at {}", address);
         util::ping_until_ok(&address, 60).await?;
@@ -69,7 +67,7 @@ impl SignerNode {
             cipher_key: *cipher_key,
             gcp_project_id: ctx.gcp_project_id.clone(),
             gcp_datastore_url: ctx.datastore.local_address.clone(),
-            process,
+            _process: process,
         })
     }
 
@@ -92,8 +90,7 @@ pub struct LeaderNode {
     relayer_url: String,
 
     // process held so it's not dropped. Once dropped, process will be killed.
-    #[allow(unused)]
-    process: Child,
+    _process: NodeProcess,
 }
 
 impl LeaderNode {
@@ -134,10 +131,9 @@ impl LeaderNode {
             gcp_datastore_url: Some(ctx.datastore.local_address.clone()),
             jwt_signature_pk_url: ctx.oidc_provider.jwt_pk_local_url.clone(),
             logging_options: logging::Options::default(),
-        }
-        .into_str_args();
+        };
 
-        let process = util::spawn_mpc(ctx.release, "leader", &args)?;
+        let process = mpc::spawn(ctx.release, "leader", args).await?;
         let address = format!("http://127.0.0.1:{web_port}");
         tracing::info!("Leader node container is starting at {}", address);
         util::ping_until_ok(&address, 60).await?;
@@ -147,7 +143,7 @@ impl LeaderNode {
             address,
             near_rpc: ctx.relayer_ctx.sandbox.local_address.clone(),
             relayer_url: ctx.relayer_ctx.relayer.local_address.clone(),
-            process,
+            _process: process,
         })
     }
 

integration-tests/src/lib.rs

@@ -11,6 +11,7 @@ use near_workspaces::{
 use crate::env::containers;
 
 pub mod env;
+pub mod mpc;
 pub mod sandbox;
 pub mod util;
 

integration-tests/src/mpc.rs

@@ -0,0 +1,87 @@
+use std::path::{Path, PathBuf};
+
+use anyhow::Context;
+use async_process::{Command, ExitStatus, Stdio};
+use tokio::runtime::Runtime;
+
+use mpc_recovery::Cli;
+
+const PACKAGE: &str = "mpc-recovery";
+
+/// NodeProcess holds onto the respective handles such that on drop, it will clean
+/// the running process, task, or thread.
+pub enum NodeProcess {
+    Subprocess(async_process::Child),
+    Threaded(std::thread::JoinHandle<anyhow::Result<()>>),
+}
+
+fn executable(release: bool) -> Option<PathBuf> {
+    let executable = target_dir()?
+        .join(if release { "release" } else { "debug" })
+        .join(PACKAGE);
+    Some(executable)
+}
+
+fn target_dir() -> Option<PathBuf> {
+    let mut out_dir = Path::new(std::env!("OUT_DIR"));
+    loop {
+        if out_dir.ends_with("target") {
+            break Some(out_dir.to_path_buf());
+        }
+
+        match out_dir.parent() {
+            Some(parent) => out_dir = parent,
+            None => break None, // We've reached the root directory and didn't find "target"
+        }
+    }
+}
+
+async fn build(release: bool) -> anyhow::Result<ExitStatus> {
+    let mut cmd = Command::new("cargo");
+    cmd.arg("build")
+        .arg("--package")
+        .arg(PACKAGE)
+        .envs(std::env::vars())
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit());
+
+    if release {
+        cmd.arg("--release");
+    }
+
+    Ok(cmd.spawn()?.status().await?)
+}
+
+pub async fn spawn(release: bool, node: &str, cli: Cli) -> anyhow::Result<NodeProcess> {
+    if cfg!(feature = "flamegraph") {
+        let handle: std::thread::JoinHandle<anyhow::Result<()>> = std::thread::spawn(|| {
+            let rt = Runtime::new()?;
+            rt.block_on(async move {
+                mpc_recovery::run(cli).await?;
+                anyhow::Result::<(), anyhow::Error>::Ok(())
+            })
+            .unwrap();
+            Ok(())
+        });
+
+        return Ok(NodeProcess::Threaded(handle));
+    }
+
+    if !build(release).await?.success() {
+        anyhow::bail!("failed to prebuild MPC service for {node} node");
+    }
+
+    let executable = executable(release)
+        .with_context(|| format!("could not find target dir while starting {node} node"))?;
+    let child = Command::new(executable)
+        .args(cli.into_str_args())
+        .env("RUST_LOG", "mpc_recovery=INFO")
+        .envs(std::env::vars())
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit())
+        .kill_on_drop(true)
+        .spawn()
+        .with_context(|| format!("failed to execute {node} node"))?;
+
+    Ok(NodeProcess::Subprocess(child))
+}

integration-tests/src/util.rs

@@ -1,20 +1,16 @@
 use std::{
     fs::{self, File},
     io::Write,
-    path::{Path, PathBuf},
 };
 
 use anyhow::Context;
-use async_process::{Child, Command, Stdio};
 use hyper::{Body, Client, Method, Request, StatusCode, Uri};
 use near_workspaces::{types::SecretKey, AccountId};
 use serde::{Deserialize, Serialize};
 use toml::Value;
 
 use crate::containers::RelayerConfig;
 
-const EXECUTABLE: &str = "mpc-recovery";
-
 pub async fn post<U, Req: Serialize, Resp>(
     uri: U,
     request: Req,
@@ -232,39 +228,3 @@ pub async fn ping_until_ok(addr: &str, timeout: u64) -> anyhow::Result<()> {
     .await?;
     Ok(())
 }
-
-pub fn target_dir() -> Option<PathBuf> {
-    let mut out_dir = Path::new(std::env!("OUT_DIR"));
-    loop {
-        if out_dir.ends_with("target") {
-            break Some(out_dir.to_path_buf());
-        }
-
-        match out_dir.parent() {
-            Some(parent) => out_dir = parent,
-            None => break None, // We've reached the root directory and didn't find "target"
-        }
-    }
-}
-
-pub fn executable(release: bool) -> Option<PathBuf> {
-    let executable = target_dir()?
-        .join(if release { "release" } else { "debug" })
-        .join(EXECUTABLE);
-    Some(executable)
-}
-
-pub fn spawn_mpc(release: bool, node: &str, args: &[String]) -> anyhow::Result<Child> {
-    let executable = executable(release)
-        .with_context(|| format!("could not find target dir while starting {node} node"))?;
-
-    Command::new(&executable)
-        .args(args)
-        .env("RUST_LOG", "mpc_recovery=INFO")
-        .envs(std::env::vars())
-        .stdout(Stdio::inherit())
-        .stderr(Stdio::inherit())
-        .kill_on_drop(true)
-        .spawn()
-        .with_context(|| format!("failed to run {node} node: {}", executable.display()))
-}

mpc-recovery/Cargo.toml

@@ -59,3 +59,7 @@ sha2 = "0.9.9"
 
 [dev-dependencies]
 rsa = "0.8.2"
+
+[features]
+default = []
+disable-open-telemetry = []

mpc-recovery/src/lib.rs

@@ -212,14 +212,13 @@ pub async fn run(cmd: Cli) -> anyhow::Result<()> {
             jwt_signature_pk_url,
             logging_options,
         } => {
-            let _subscriber_guard = logging::default_subscriber_with_opentelemetry(
+            let _subscriber_guard = logging::subscribe_global(
                 EnvFilter::from_default_env(),
                 &logging_options,
                 env.clone(),
                 "leader".to_string(),
             )
-            .await
-            .global();
+            .await;
             let gcp_service =
                 GcpService::new(env.clone(), gcp_project_id, gcp_datastore_url).await?;
             let account_creator_signer =
@@ -258,14 +257,13 @@ pub async fn run(cmd: Cli) -> anyhow::Result<()> {
             jwt_signature_pk_url,
             logging_options,
         } => {
-            let _subscriber_guard = logging::default_subscriber_with_opentelemetry(
+            let _subscriber_guard = logging::subscribe_global(
                 EnvFilter::from_default_env(),
                 &logging_options,
                 env.clone(),
                 node_id.to_string(),
             )
-            .await
-            .global();
+            .await;
             let gcp_service =
                 GcpService::new(env.clone(), gcp_project_id, gcp_datastore_url).await?;
             let oidc_providers = OidcProviderList {
@@ -309,14 +307,13 @@ pub async fn run(cmd: Cli) -> anyhow::Result<()> {
             gcp_datastore_url,
             logging_options,
         } => {
-            let _subscriber_guard = logging::default_subscriber_with_opentelemetry(
+            let _subscriber_guard = logging::subscribe_global(
                 EnvFilter::from_default_env(),
                 &logging_options,
                 env.clone(),
                 node_id.to_string(),
             )
-            .await
-            .global();
+            .await;
             let gcp_service = GcpService::new(
                 env.clone(),
                 gcp_project_id.clone(),