utils: symbolize heap profiles (#10153)

## Problem

Jemalloc heap profiles aren't symbolized. This is inconvenient, and
doesn't work with Grafana Cloud Profiles.

Resolves #9964.

## Summary of changes

Symbolize the heap profiles in-process, and strip unnecessary cruft.

This uses about 100 MB additional memory to cache the DWARF information,
but I believe this is already the case with CPU profiles, which use the
same library for symbolization. With cached DWARF information, the
symbolization CPU overhead is negligible.

Example profiles:

*
[pageserver.pb.gz](https://github.com/user-attachments/files/18141395/pageserver.pb.gz)
*
[safekeeper.pb.gz](https://github.com/user-attachments/files/18141396/safekeeper.pb.gz)

Cargo.lock

@@ -10,9 +10,9 @@ checksum = "8b5ace29ee3216de37c0546865ad08edef58b0f9e76838ed8959a84a990e58c5"
 
 [[package]]
 name = "addr2line"
-version = "0.21.0"
+version = "0.24.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8a30b2e23b9e17a9f90641c7ab1549cd9b44f296d3ccbf309d2863cfe398a0cb"
+checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1"
 dependencies = [
  "gimli",
 ]
@@ -23,6 +23,12 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "adler2"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
+
 [[package]]
 name = "ahash"
 version = "0.8.11"
@@ -871,17 +877,17 @@ dependencies = [
 
 [[package]]
 name = "backtrace"
-version = "0.3.69"
+version = "0.3.74"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2089b7e3f35b9dd2d0ed921ead4f6d318c27680d4a5bd167b3ee120edb105837"
+checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a"
 dependencies = [
  "addr2line",
- "cc",
  "cfg-if",
  "libc",
- "miniz_oxide",
+ "miniz_oxide 0.8.0",
  "object",
  "rustc-demangle",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -1127,7 +1133,7 @@ dependencies = [
  "num-traits",
  "serde",
  "wasm-bindgen",
- "windows-targets 0.52.4",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -2107,7 +2113,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743"
 dependencies = [
  "crc32fast",
- "miniz_oxide",
+ "miniz_oxide 0.7.1",
 ]
 
 [[package]]
@@ -2308,9 +2314,9 @@ dependencies = [
 
 [[package]]
 name = "gimli"
-version = "0.28.1"
+version = "0.31.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253"
+checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
 
 [[package]]
 name = "git-version"
@@ -3404,6 +3410,15 @@ dependencies = [
  "adler",
 ]
 
+[[package]]
+name = "miniz_oxide"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e2d80299ef12ff69b16a84bb182e3b9df68b5a91574d3d4fa6e41b65deec4df1"
+dependencies = [
+ "adler2",
+]
+
 [[package]]
 name = "mio"
 version = "0.8.11"
@@ -3638,9 +3653,9 @@ dependencies = [
 
 [[package]]
 name = "object"
-version = "0.32.2"
+version = "0.36.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441"
+checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e"
 dependencies = [
  "memchr",
 ]
@@ -5323,9 +5338,9 @@ dependencies = [
 
 [[package]]
 name = "rustc-demangle"
-version = "0.1.23"
+version = "0.1.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76"
+checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
 
 [[package]]
 name = "rustc-hash"
@@ -7219,6 +7234,7 @@ dependencies = [
  "anyhow",
  "arc-swap",
  "async-compression",
+ "backtrace",
  "bincode",
  "byteorder",
  "bytes",
@@ -7229,12 +7245,14 @@ dependencies = [
  "criterion",
  "diatomic-waker",
  "fail",
+ "flate2",
  "futures",
  "git-version",
  "hex",
  "hex-literal",
  "humantime",
  "hyper 0.14.30",
+ "itertools 0.10.5",
  "jemalloc_pprof",
  "jsonwebtoken",
  "metrics",
@@ -7591,7 +7609,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e48a53791691ab099e5e2ad123536d0fff50652600abaf43bbf952894110d0be"
 dependencies = [
  "windows-core",
- "windows-targets 0.52.4",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -7600,7 +7618,7 @@ version = "0.52.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9"
 dependencies = [
- "windows-targets 0.52.4",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -7618,7 +7636,7 @@ version = "0.52.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
 dependencies = [
- "windows-targets 0.52.4",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -7638,17 +7656,18 @@ dependencies = [
 
 [[package]]
 name = "windows-targets"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7dd37b7e5ab9018759f893a1952c9420d060016fc19a472b4bb20d1bdd694d1b"
+checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
 dependencies = [
- "windows_aarch64_gnullvm 0.52.4",
+ "windows_aarch64_gnullvm 0.52.6",
- "windows_aarch64_msvc 0.52.4",
+ "windows_aarch64_msvc 0.52.6",
- "windows_i686_gnu 0.52.4",
+ "windows_i686_gnu 0.52.6",
- "windows_i686_msvc 0.52.4",
+ "windows_i686_gnullvm",
- "windows_x86_64_gnu 0.52.4",
+ "windows_i686_msvc 0.52.6",
- "windows_x86_64_gnullvm 0.52.4",
+ "windows_x86_64_gnu 0.52.6",
- "windows_x86_64_msvc 0.52.4",
+ "windows_x86_64_gnullvm 0.52.6",
+ "windows_x86_64_msvc 0.52.6",
 ]
 
 [[package]]
@@ -7659,9 +7678,9 @@ checksum = "91ae572e1b79dba883e0d315474df7305d12f569b400fcf90581b06062f7e1bc"
 
 [[package]]
 name = "windows_aarch64_gnullvm"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bcf46cf4c365c6f2d1cc93ce535f2c8b244591df96ceee75d8e83deb70a9cac9"
+checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
 
 [[package]]
 name = "windows_aarch64_msvc"
@@ -7671,9 +7690,9 @@ checksum = "b2ef27e0d7bdfcfc7b868b317c1d32c641a6fe4629c171b8928c7b08d98d7cf3"
 
 [[package]]
 name = "windows_aarch64_msvc"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "da9f259dd3bcf6990b55bffd094c4f7235817ba4ceebde8e6d11cd0c5633b675"
+checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
 
 [[package]]
 name = "windows_i686_gnu"
@@ -7683,9 +7702,15 @@ checksum = "622a1962a7db830d6fd0a69683c80a18fda201879f0f447f065a3b7467daa241"
 
 [[package]]
 name = "windows_i686_gnu"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b474d8268f99e0995f25b9f095bc7434632601028cf86590aea5c8a5cb7801d3"
+checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
 
 [[package]]
 name = "windows_i686_msvc"
@@ -7695,9 +7720,9 @@ checksum = "4542c6e364ce21bf45d69fdd2a8e455fa38d316158cfd43b3ac1c5b1b19f8e00"
 
 [[package]]
 name = "windows_i686_msvc"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1515e9a29e5bed743cb4415a9ecf5dfca648ce85ee42e15873c3cd8610ff8e02"
+checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
 
 [[package]]
 name = "windows_x86_64_gnu"
@@ -7707,9 +7732,9 @@ checksum = "ca2b8a661f7628cbd23440e50b05d705db3686f894fc9580820623656af974b1"
 
 [[package]]
 name = "windows_x86_64_gnu"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5eee091590e89cc02ad514ffe3ead9eb6b660aedca2183455434b93546371a03"
+checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
 
 [[package]]
 name = "windows_x86_64_gnullvm"
@@ -7719,9 +7744,9 @@ checksum = "7896dbc1f41e08872e9d5e8f8baa8fdd2677f29468c4e156210174edc7f7b953"
 
 [[package]]
 name = "windows_x86_64_gnullvm"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77ca79f2451b49fa9e2af39f0747fe999fcda4f5e241b2898624dca97a1f2177"
+checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
 
 [[package]]
 name = "windows_x86_64_msvc"
@@ -7731,9 +7756,9 @@ checksum = "1a515f5799fe4961cb532f983ce2b23082366b898e52ffbce459c86f67c8378a"
 
 [[package]]
 name = "windows_x86_64_msvc"
-version = "0.52.4"
+version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32b752e52a2da0ddfbdbcc6fceadfeede4c939ed16d13e648833a61dfb611ed8"
+checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
 
 [[package]]
 name = "winnow"

Cargo.toml

@@ -52,6 +52,7 @@ anyhow = { version = "1.0", features = ["backtrace"] }
 arc-swap = "1.6"
 async-compression = { version = "0.4.0", features = ["tokio", "gzip", "zstd"] }
 atomic-take = "1.1.0"
+backtrace = "0.3.74"
 flate2 = "1.0.26"
 async-stream = "0.3"
 async-trait = "0.1"

libs/utils/Cargo.toml

@@ -15,17 +15,20 @@ arc-swap.workspace = true
 sentry.workspace = true
 async-compression.workspace = true
 anyhow.workspace = true
+backtrace.workspace = true
 bincode.workspace = true
 bytes.workspace = true
 camino.workspace = true
 chrono.workspace = true
 diatomic-waker.workspace = true
+flate2.workspace = true
 git-version.workspace = true
 hex = { workspace = true, features = ["serde"] }
 humantime.workspace = true
 hyper0 = { workspace = true, features = ["full"] }
+itertools.workspace = true
 fail.workspace = true
-futures = { workspace = true}
+futures = { workspace = true }
 jemalloc_pprof.workspace = true
 jsonwebtoken.workspace = true
 nix.workspace = true

libs/utils/src/http/endpoint.rs

@@ -1,15 +1,22 @@
 use crate::auth::{AuthError, Claims, SwappableJwtAuth};
 use crate::http::error::{api_error_handler, route_error_handler, ApiError};
 use crate::http::request::{get_query_param, parse_query_param};
+use crate::pprof;
+use ::pprof::protos::Message as _;
+use ::pprof::ProfilerGuardBuilder;
 use anyhow::{anyhow, Context};
+use bytes::{Bytes, BytesMut};
 use hyper::header::{HeaderName, AUTHORIZATION, CONTENT_DISPOSITION};
 use hyper::http::HeaderValue;
 use hyper::Method;
 use hyper::{header::CONTENT_TYPE, Body, Request, Response};
 use metrics::{register_int_counter, Encoder, IntCounter, TextEncoder};
 use once_cell::sync::Lazy;
+use regex::Regex;
 use routerify::ext::RequestExt;
 use routerify::{Middleware, RequestInfo, Router, RouterBuilder};
+use tokio::sync::{mpsc, Mutex};
+use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::io::ReaderStream;
 use tracing::{debug, info, info_span, warn, Instrument};
 
@@ -18,11 +25,6 @@ use std::io::Write as _;
 use std::str::FromStr;
 use std::time::Duration;
 
-use bytes::{Bytes, BytesMut};
-use pprof::protos::Message as _;
-use tokio::sync::{mpsc, Mutex};
-use tokio_stream::wrappers::ReceiverStream;
-
 static SERVE_METRICS_COUNT: Lazy<IntCounter> = Lazy::new(|| {
     register_int_counter!(
         "libmetrics_metric_handler_requests_total",
@@ -365,7 +367,7 @@ pub async fn profile_cpu_handler(req: Request<Body>) -> Result<Response<Body>, A
 
     // Take the profile.
     let report = tokio::task::spawn_blocking(move || {
-        let guard = pprof::ProfilerGuardBuilder::default()
+        let guard = ProfilerGuardBuilder::default()
             .frequency(frequency_hz)
             .blocklist(&["libc", "libgcc", "pthread", "vdso"])
             .build()?;
@@ -457,10 +459,34 @@ pub async fn profile_heap_handler(req: Request<Body>) -> Result<Response<Body>,
         }
 
         Format::Pprof => {
-            let data = tokio::task::spawn_blocking(move || prof_ctl.dump_pprof())
+            let data = tokio::task::spawn_blocking(move || {
-                .await
+                let bytes = prof_ctl.dump_pprof()?;
-                .map_err(|join_err| ApiError::InternalServerError(join_err.into()))?
+                // Symbolize the profile.
-                .map_err(ApiError::InternalServerError)?;
+                // TODO: consider moving this upstream to jemalloc_pprof and avoiding the
+                // serialization roundtrip.
+                static STRIP_FUNCTIONS: Lazy<Vec<(Regex, bool)>> = Lazy::new(|| {
+                    // Functions to strip from profiles. If true, also remove child frames.
+                    vec![
+                        (Regex::new("^__rust").unwrap(), false),
+                        (Regex::new("^_start$").unwrap(), false),
+                        (Regex::new("^irallocx_prof").unwrap(), true),
+                        (Regex::new("^prof_alloc_prep").unwrap(), true),
+                        (Regex::new("^std::rt::lang_start").unwrap(), false),
+                        (Regex::new("^std::sys::backtrace::__rust").unwrap(), false),
+                    ]
+                });
+                let profile = pprof::decode(&bytes)?;
+                let profile = pprof::symbolize(profile)?;
+                let profile = pprof::strip_locations(
+                    profile,
+                    &["libc", "libgcc", "pthread", "vdso"],
+                    &STRIP_FUNCTIONS,
+                );
+                pprof::encode(&profile)
+            })
+            .await
+            .map_err(|join_err| ApiError::InternalServerError(join_err.into()))?
+            .map_err(ApiError::InternalServerError)?;
             Response::builder()
                 .status(200)
                 .header(CONTENT_TYPE, "application/octet-stream")

libs/utils/src/lib.rs

@@ -96,6 +96,8 @@ pub mod circuit_breaker;
 
 pub mod try_rcu;
 
+pub mod pprof;
+
 // Re-export used in macro. Avoids adding git-version as dep in target crates.
 #[doc(hidden)]
 pub use git_version;

libs/utils/src/pprof.rs

@@ -0,0 +1,190 @@
+use flate2::write::{GzDecoder, GzEncoder};
+use flate2::Compression;
+use itertools::Itertools as _;
+use once_cell::sync::Lazy;
+use pprof::protos::{Function, Line, Message as _, Profile};
+use regex::Regex;
+
+use std::borrow::Cow;
+use std::collections::{HashMap, HashSet};
+use std::ffi::c_void;
+use std::io::Write as _;
+
+/// Decodes a gzip-compressed Protobuf-encoded pprof profile.
+pub fn decode(bytes: &[u8]) -> anyhow::Result<Profile> {
+    let mut gz = GzDecoder::new(Vec::new());
+    gz.write_all(bytes)?;
+    Ok(Profile::parse_from_bytes(&gz.finish()?)?)
+}
+
+/// Encodes a pprof profile as gzip-compressed Protobuf.
+pub fn encode(profile: &Profile) -> anyhow::Result<Vec<u8>> {
+    let mut gz = GzEncoder::new(Vec::new(), Compression::default());
+    profile.write_to_writer(&mut gz)?;
+    Ok(gz.finish()?)
+}
+
+/// Symbolizes a pprof profile using the current binary.
+pub fn symbolize(mut profile: Profile) -> anyhow::Result<Profile> {
+    if !profile.function.is_empty() {
+        return Ok(profile); // already symbolized
+    }
+
+    // Collect function names.
+    let mut functions: HashMap<String, Function> = HashMap::new();
+    let mut strings: HashMap<String, i64> = profile
+        .string_table
+        .into_iter()
+        .enumerate()
+        .map(|(i, s)| (s, i as i64))
+        .collect();
+
+    // Helper to look up or register a string.
+    let mut string_id = |s: &str| -> i64 {
+        // Don't use .entry() to avoid unnecessary allocations.
+        if let Some(id) = strings.get(s) {
+            return *id;
+        }
+        let id = strings.len() as i64;
+        strings.insert(s.to_string(), id);
+        id
+    };
+
+    for loc in &mut profile.location {
+        if !loc.line.is_empty() {
+            continue;
+        }
+
+        // Resolve the line and function for each location.
+        backtrace::resolve(loc.address as *mut c_void, |symbol| {
+            let Some(symname) = symbol.name() else {
+                return;
+            };
+            let mut name = symname.to_string();
+
+            // Strip the Rust monomorphization suffix from the symbol name.
+            static SUFFIX_REGEX: Lazy<Regex> =
+                Lazy::new(|| Regex::new("::h[0-9a-f]{16}$").expect("invalid regex"));
+            if let Some(m) = SUFFIX_REGEX.find(&name) {
+                name.truncate(m.start());
+            }
+
+            let function_id = match functions.get(&name) {
+                Some(function) => function.id,
+                None => {
+                    let id = functions.len() as u64 + 1;
+                    let system_name = String::from_utf8_lossy(symname.as_bytes());
+                    let filename = symbol
+                        .filename()
+                        .map(|path| path.to_string_lossy())
+                        .unwrap_or(Cow::Borrowed(""));
+                    let function = Function {
+                        id,
+                        name: string_id(&name),
+                        system_name: string_id(&system_name),
+                        filename: string_id(&filename),
+                        ..Default::default()
+                    };
+                    functions.insert(name, function);
+                    id
+                }
+            };
+            loc.line.push(Line {
+                function_id,
+                line: symbol.lineno().unwrap_or(0) as i64,
+                ..Default::default()
+            });
+        });
+    }
+
+    // Store the resolved functions, and mark the mapping as resolved.
+    profile.function = functions.into_values().sorted_by_key(|f| f.id).collect();
+    profile.string_table = strings
+        .into_iter()
+        .sorted_by_key(|(_, i)| *i)
+        .map(|(s, _)| s)
+        .collect();
+
+    for mapping in &mut profile.mapping {
+        mapping.has_functions = true;
+        mapping.has_filenames = true;
+    }
+
+    Ok(profile)
+}
+
+/// Strips locations (stack frames) matching the given mappings (substring) or function names
+/// (regex). The function bool specifies whether child frames should be stripped as well.
+///
+/// The string definitions are left behind in the profile for simplicity, to avoid rewriting all
+/// string references.
+pub fn strip_locations(
+    mut profile: Profile,
+    mappings: &[&str],
+    functions: &[(Regex, bool)],
+) -> Profile {
+    // Strip mappings.
+    let mut strip_mappings: HashSet<u64> = HashSet::new();
+
+    profile.mapping.retain(|mapping| {
+        let Some(name) = profile.string_table.get(mapping.filename as usize) else {
+            return true;
+        };
+        if mappings.iter().any(|substr| name.contains(substr)) {
+            strip_mappings.insert(mapping.id);
+            return false;
+        }
+        true
+    });
+
+    // Strip functions.
+    let mut strip_functions: HashMap<u64, bool> = HashMap::new();
+
+    profile.function.retain(|function| {
+        let Some(name) = profile.string_table.get(function.name as usize) else {
+            return true;
+        };
+        for (regex, strip_children) in functions {
+            if regex.is_match(name) {
+                strip_functions.insert(function.id, *strip_children);
+                return false;
+            }
+        }
+        true
+    });
+
+    // Strip locations. The bool specifies whether child frames should be stripped too.
+    let mut strip_locations: HashMap<u64, bool> = HashMap::new();
+
+    profile.location.retain(|location| {
+        for line in &location.line {
+            if let Some(strip_children) = strip_functions.get(&line.function_id) {
+                strip_locations.insert(location.id, *strip_children);
+                return false;
+            }
+        }
+        if strip_mappings.contains(&location.mapping_id) {
+            strip_locations.insert(location.id, false);
+            return false;
+        }
+        true
+    });
+
+    // Strip sample locations.
+    for sample in &mut profile.sample {
+        // First, find the uppermost function with child removal and truncate the stack.
+        if let Some(truncate) = sample
+            .location_id
+            .iter()
+            .rposition(|id| strip_locations.get(id) == Some(&true))
+        {
+            sample.location_id.drain(..=truncate);
+        }
+        // Next, strip any individual frames without child removal.
+        sample
+            .location_id
+            .retain(|id| !strip_locations.contains_key(id));
+    }
+
+    profile
+}