From 98e3592d378709dad6c599a6ec5d7e9328ab41c0 Mon Sep 17 00:00:00 2001
From: Andrei Eres <eresav@me.com>
Date: Thu, 16 Nov 2023 17:58:52 +0100
Subject: [PATCH] Add basic pyroscope agent

---
 Cargo.lock                                           |  2 ++
 polkadot/node/subsystem-bench/Cargo.toml             |  2 ++
 polkadot/node/subsystem-bench/src/subsystem-bench.rs | 12 ++++++++++++
 3 files changed, 16 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 73fc3cbdeccc..f8b08e10d0ab 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -13039,6 +13039,8 @@ dependencies = [
  "polkadot-primitives",
  "polkadot-primitives-test-helpers",
  "prometheus",
+ "pyroscope",
+ "pyroscope_pprofrs",
  "rand 0.8.5",
  "sc-keystore",
  "sc-network",
diff --git a/polkadot/node/subsystem-bench/Cargo.toml b/polkadot/node/subsystem-bench/Cargo.toml
index 3308b6fe1052..ae44a60adf43 100644
--- a/polkadot/node/subsystem-bench/Cargo.toml
+++ b/polkadot/node/subsystem-bench/Cargo.toml
@@ -53,6 +53,8 @@ prometheus_endpoint = { package = "substrate-prometheus-endpoint", path = "../..
 prometheus = { version = "0.13.0", default-features = false }
 toml = "0.8.8"
 serde = "1.0.192"
+pyroscope = "0.5.7"
+pyroscope_pprofrs = "0.2.7"
 
 [features]
 default = []
diff --git a/polkadot/node/subsystem-bench/src/subsystem-bench.rs b/polkadot/node/subsystem-bench/src/subsystem-bench.rs
index 7dcc8a15074a..c13078c17dc0 100644
--- a/polkadot/node/subsystem-bench/src/subsystem-bench.rs
+++ b/polkadot/node/subsystem-bench/src/subsystem-bench.rs
@@ -143,6 +143,15 @@ fn new_runtime() -> tokio::runtime::Runtime {
 impl BenchCli {
 	fn launch(self) -> eyre::Result<()> {
 		use prometheus::Registry;
+		use pyroscope::PyroscopeAgent;
+		use pyroscope_pprofrs::{pprof_backend, PprofConfig};
+
+		// Pyroscope must be running on port 4040
+		// See https://grafana.com/docs/pyroscope/latest/get-started/#download-and-configure-pyroscope
+		let agent = PyroscopeAgent::builder("http://localhost:4040", "subsystem-bench")
+			.backend(pprof_backend(PprofConfig::new().sample_rate(100)))
+			.build()?;
+		let agent_running = agent.start()?;
 
 		let runtime = new_runtime();
 
@@ -253,6 +262,9 @@ impl BenchCli {
 
 		runtime.block_on(availability::bench_chunk_recovery(&mut env));
 
+		let agent_ready = agent_running.stop()?;
+		agent_ready.shutdown();
+
 		Ok(())
 	}
 }