feat: add support for wrapping the allocator with tracy (#10874)

crates/cli/util/Cargo.toml

@@ -18,12 +18,21 @@ reth-fs-util.workspace = true
 alloy-primitives.workspace = true
 alloy-eips.workspace = true
 
-secp256k1 = { workspace = true, features = ["rand"] }
-rand.workspace = true
-
 # misc
-thiserror.workspace = true
+cfg-if.workspace = true
 eyre.workspace = true
+rand.workspace = true
+secp256k1 = { workspace = true, features = ["rand"] }
+thiserror.workspace = true
+
+tracy-client = { workspace = true, optional = true, features = ["demangle"] }
 
 [target.'cfg(unix)'.dependencies]
+tikv-jemallocator = { workspace = true, optional = true }
 libc = "0.2"
+
+[features]
+jemalloc = ["dep:tikv-jemallocator"]
+jemalloc-prof = ["jemalloc", "tikv-jemallocator?/profiling"]
+
+tracy-allocator = ["dep:tracy-client"]

crates/cli/util/src/allocator.rs

@@ -0,0 +1,35 @@
+//! Custom allocator implementation.
+
+// We use jemalloc for performance reasons.
+cfg_if::cfg_if! {
+    if #[cfg(all(feature = "jemalloc", unix))] {
+        type AllocatorInner = tikv_jemallocator::Jemalloc;
+    } else {
+        type AllocatorInner = std::alloc::System;
+    }
+}
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "tracy-allocator")] {
+        type AllocatorWrapper = tracy_client::ProfiledAllocator<AllocatorInner>;
+        const fn new_allocator_wrapper() -> AllocatorWrapper {
+            AllocatorWrapper::new(AllocatorInner {}, 100)
+        }
+    } else {
+        type AllocatorWrapper = AllocatorInner;
+        const fn new_allocator_wrapper() -> AllocatorWrapper {
+            AllocatorInner {}
+        }
+    }
+}
+
+#[cfg(feature = "tracy-allocator")]
+tracy_client::register_demangler!();
+
+/// Custom allocator.
+pub type Allocator = AllocatorWrapper;
+
+/// Creates a new [custom allocator][Allocator].
+pub const fn new_allocator() -> Allocator {
+    new_allocator_wrapper()
+}

crates/cli/util/src/lib.rs

@@ -8,6 +8,8 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
+pub mod allocator;
+
 /// Helper function to load a secret key from a file.
 pub mod load_secret_key;
 pub use load_secret_key::get_secret_key;