mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-11 07:36:53 +00:00
fix(runtime): finalize issue 2941 profiling cleanup (#2983)
* perf(runtime): narrow profiling support and upgrade starshard * style(notify): normalize starshard imports * perf(ecstore): reduce list_path_raw coordination overhead * docs(scripts): add issue 2941 perf capture workflow * fix(runtime): finalize issue 2941 profiling cleanup * build(deps): bump quick-xml to 0.40.0 * chore(scripts): untrack local perf capture guide * fix(scripts): honor label in perf capture output
This commit is contained in:
houseme
+1
-7
@@ -176,13 +176,7 @@ libmimalloc-sys = { version = "0.1.47", features = ["extended"] }
|
||||
|
||||
|
||||
|
||||
# Only enable pprof-based profiling on non-Windows targets.
|
||||
[target.'cfg(all(not(target_os = "windows"), not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))))'.dependencies]
|
||||
starshard = { workspace = true }
|
||||
backtrace = { workspace = true }
|
||||
rand = { workspace = true }
|
||||
pprof = { workspace = true }
|
||||
|
||||
# Only enable pprof-based profiling on linux + gnu + x86_64.
|
||||
[target.'cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))'.dependencies]
|
||||
tikv-jemallocator = { workspace = true }
|
||||
tikv-jemalloc-ctl = { workspace = true }
|
||||
|
||||
@@ -15,8 +15,11 @@
|
||||
use crate::admin::{auth::validate_admin_request, router::Operation};
|
||||
use crate::auth::{check_key_valid, get_session_token};
|
||||
use crate::server::RemoteAddr;
|
||||
#[cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))]
|
||||
use http::HeaderMap;
|
||||
use http::StatusCode;
|
||||
#[cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))]
|
||||
use http::header::CONTENT_TYPE;
|
||||
use http::{HeaderMap, StatusCode};
|
||||
use matchit::Params;
|
||||
use rustfs_policy::policy::action::{Action, AdminAction};
|
||||
use s3s::{Body, S3Request, S3Response, S3Result, s3_error};
|
||||
@@ -49,14 +52,29 @@ impl Operation for TriggerProfileCPU {
|
||||
authorize_profile_request(&req).await?;
|
||||
info!("Triggering CPU profile dump via S3 request...");
|
||||
|
||||
let dur = std::time::Duration::from_secs(60);
|
||||
match crate::profiling::dump_cpu_pprof_for(dur).await {
|
||||
Ok(path) => {
|
||||
let mut header = HeaderMap::new();
|
||||
header.insert(CONTENT_TYPE, "text/html".parse().unwrap());
|
||||
Ok(S3Response::with_headers((StatusCode::OK, Body::from(path.display().to_string())), header))
|
||||
#[cfg(not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64")))]
|
||||
{
|
||||
return Ok(S3Response::new((
|
||||
StatusCode::NOT_IMPLEMENTED,
|
||||
Body::from(
|
||||
crate::profiling::dump_cpu_pprof_for(std::time::Duration::from_secs(0))
|
||||
.await
|
||||
.unwrap_err(),
|
||||
),
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))]
|
||||
{
|
||||
let dur = std::time::Duration::from_secs(60);
|
||||
match crate::profiling::dump_cpu_pprof_for(dur).await {
|
||||
Ok(path) => {
|
||||
let mut header = HeaderMap::new();
|
||||
header.insert(CONTENT_TYPE, "text/html".parse().unwrap());
|
||||
Ok(S3Response::with_headers((StatusCode::OK, Body::from(path.display().to_string())), header))
|
||||
}
|
||||
Err(e) => Err(s3s::s3_error!(InternalError, "{}", format!("Failed to dump CPU profile: {e}"))),
|
||||
}
|
||||
Err(e) => Err(s3s::s3_error!(InternalError, "{}", format!("Failed to dump CPU profile: {e}"))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -68,13 +86,24 @@ impl Operation for TriggerProfileMemory {
|
||||
authorize_profile_request(&req).await?;
|
||||
info!("Triggering Memory profile dump via S3 request...");
|
||||
|
||||
match crate::profiling::dump_memory_pprof_now().await {
|
||||
Ok(path) => {
|
||||
let mut header = HeaderMap::new();
|
||||
header.insert(CONTENT_TYPE, "text/html".parse().unwrap());
|
||||
Ok(S3Response::with_headers((StatusCode::OK, Body::from(path.display().to_string())), header))
|
||||
#[cfg(not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64")))]
|
||||
{
|
||||
return Ok(S3Response::new((
|
||||
StatusCode::NOT_IMPLEMENTED,
|
||||
Body::from(crate::profiling::dump_memory_pprof_now().await.unwrap_err()),
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))]
|
||||
{
|
||||
match crate::profiling::dump_memory_pprof_now().await {
|
||||
Ok(path) => {
|
||||
let mut header = HeaderMap::new();
|
||||
header.insert(CONTENT_TYPE, "text/html".parse().unwrap());
|
||||
Ok(S3Response::with_headers((StatusCode::OK, Body::from(path.display().to_string())), header))
|
||||
}
|
||||
Err(e) => Err(s3s::s3_error!(InternalError, "{}", format!("Failed to dump Memory profile: {e}"))),
|
||||
}
|
||||
Err(e) => Err(s3s::s3_error!(InternalError, "{}", format!("Failed to dump Memory profile: {e}"))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -51,6 +51,7 @@ use crate::config::Config;
|
||||
use crate::init::{add_bucket_notification_configuration, init_buffer_profile_system, init_kms_system};
|
||||
use crate::server::{init_event_notifier, shutdown_event_notifier, start_audit_system, start_http_server, stop_audit_system};
|
||||
use rustfs_common::{GlobalReadiness, SystemStage, set_global_addr};
|
||||
use rustfs_config::ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS;
|
||||
use rustfs_credentials::init_global_action_credentials;
|
||||
use rustfs_ecstore::store::init_lock_clients;
|
||||
use rustfs_ecstore::{
|
||||
@@ -317,7 +318,7 @@ impl RustFSServerBuilder {
|
||||
));
|
||||
}
|
||||
|
||||
let allow_insecure_defaults = get_env_bool(rustfs_config::ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS, false);
|
||||
let allow_insecure_defaults = get_env_bool(ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS, false);
|
||||
if !config.default_credentials_allowed_for_addr(server_addr, allow_insecure_defaults) {
|
||||
return Err(ServerError::Init(
|
||||
"default root credentials are not allowed on non-loopback listeners; set access_key and secret_key to non-default values, bind to loopback, or set RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS=true for local development only"
|
||||
|
||||
+4
-11
@@ -34,6 +34,7 @@ use rustfs::server::{
|
||||
start_audit_system, start_http_server, stop_audit_system, wait_for_shutdown,
|
||||
};
|
||||
use rustfs_common::{GlobalReadiness, SystemStage, set_global_addr};
|
||||
use rustfs_config::ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS;
|
||||
use rustfs_credentials::init_global_action_credentials;
|
||||
use rustfs_ecstore::store::init_lock_clients;
|
||||
use rustfs_ecstore::{
|
||||
@@ -76,15 +77,7 @@ const ENV_HEAL_ENABLED_DEPRECATED: &str = "RUSTFS_ENABLE_HEAL";
|
||||
#[global_allocator]
|
||||
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
||||
|
||||
#[cfg(all(
|
||||
not(target_os = "windows"),
|
||||
not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))
|
||||
))]
|
||||
#[global_allocator]
|
||||
static GLOBAL: rustfs::profiling::allocator::TracingAllocator<mimalloc::MiMalloc> =
|
||||
rustfs::profiling::allocator::TracingAllocator::new(mimalloc::MiMalloc);
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
#[cfg(not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64")))]
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
@@ -145,7 +138,7 @@ const DEFAULT_CREDENTIALS_WARNING_MESSAGE: &str = "Detected default root credent
|
||||
const DEFAULT_CREDENTIALS_ERROR_MESSAGE: &str = "Default root credentials are not allowed on non-loopback listeners; set RUSTFS_ACCESS_KEY and RUSTFS_SECRET_KEY to non-default values, bind to loopback, or set RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS=true for local development only";
|
||||
|
||||
fn allow_insecure_default_credentials() -> bool {
|
||||
get_env_bool(rustfs_config::ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS, false)
|
||||
get_env_bool(ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS, false)
|
||||
}
|
||||
|
||||
async fn async_main() -> Result<()> {
|
||||
@@ -845,7 +838,7 @@ mod tests {
|
||||
for message in [DEFAULT_CREDENTIALS_WARNING_MESSAGE, DEFAULT_CREDENTIALS_ERROR_MESSAGE] {
|
||||
assert!(message.contains(rustfs_config::ENV_RUSTFS_ACCESS_KEY));
|
||||
assert!(message.contains(rustfs_config::ENV_RUSTFS_SECRET_KEY));
|
||||
assert!(message.contains(rustfs_config::ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS));
|
||||
assert!(message.contains(ENV_RUSTFS_ALLOW_INSECURE_DEFAULT_CREDENTIALS));
|
||||
assert!(!message.contains(rustfs_credentials::DEFAULT_ACCESS_KEY));
|
||||
assert!(!message.contains(rustfs_credentials::DEFAULT_SECRET_KEY));
|
||||
}
|
||||
|
||||
+26
-127
@@ -12,154 +12,53 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#[cfg(all(
|
||||
not(target_os = "windows"),
|
||||
not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))
|
||||
))]
|
||||
pub mod allocator;
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows_impl {
|
||||
#[cfg(not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64")))]
|
||||
mod unsupported_impl {
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tracing::info;
|
||||
|
||||
pub async fn init_from_env() {
|
||||
info!("Profiling initialization skipped on Windows platform (not supported)");
|
||||
let target_env = option_env!("CARGO_CFG_TARGET_ENV").unwrap_or("unknown");
|
||||
info!(
|
||||
target_os = std::env::consts::OS,
|
||||
target_env,
|
||||
target_arch = std::env::consts::ARCH,
|
||||
"Profiling initialization skipped on unsupported platform"
|
||||
);
|
||||
}
|
||||
|
||||
/// Stop all background profiling tasks
|
||||
pub fn shutdown_profiling() {
|
||||
info!("profiling: shutdown called on Windows platform (no-op)");
|
||||
let target_env = option_env!("CARGO_CFG_TARGET_ENV").unwrap_or("unknown");
|
||||
info!(
|
||||
target_os = std::env::consts::OS,
|
||||
target_env,
|
||||
target_arch = std::env::consts::ARCH,
|
||||
"profiling: shutdown called on unsupported platform (no-op)"
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn dump_cpu_pprof_for(_duration: Duration) -> Result<PathBuf, String> {
|
||||
Err("CPU profiling is not supported on Windows platform".to_string())
|
||||
Err(unsupported_message("CPU profiling"))
|
||||
}
|
||||
|
||||
pub async fn dump_memory_pprof_now() -> Result<PathBuf, String> {
|
||||
Err("Memory profiling is not supported on Windows platform".to_string())
|
||||
Err(unsupported_message("Memory profiling"))
|
||||
}
|
||||
}
|
||||
#[cfg(target_os = "windows")]
|
||||
pub use windows_impl::{dump_cpu_pprof_for, dump_memory_pprof_now, init_from_env, shutdown_profiling};
|
||||
|
||||
#[cfg(all(
|
||||
not(target_os = "windows"),
|
||||
not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))
|
||||
))]
|
||||
mod generic_impl {
|
||||
use super::allocator;
|
||||
use rustfs_config::{
|
||||
DEFAULT_ENABLE_PROFILING, DEFAULT_MEM_INTERVAL_SECS, DEFAULT_MEM_PERIODIC, DEFAULT_OUTPUT_DIR, ENV_ENABLE_PROFILING,
|
||||
ENV_MEM_INTERVAL_SECS, ENV_MEM_PERIODIC, ENV_OUTPUT_DIR,
|
||||
};
|
||||
use rustfs_utils::{get_env_bool, get_env_str, get_env_u64};
|
||||
use std::fs::create_dir_all;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
// Global cancellation token for periodic profiling tasks
|
||||
static PROFILING_CANCEL_TOKEN: OnceLock<CancellationToken> = OnceLock::new();
|
||||
|
||||
fn get_platform_info() -> (String, String, String) {
|
||||
(
|
||||
std::env::consts::OS.to_string(),
|
||||
option_env!("CARGO_CFG_TARGET_ENV").unwrap_or("unknown").to_string(),
|
||||
std::env::consts::ARCH.to_string(),
|
||||
fn unsupported_message(feature: &str) -> String {
|
||||
let target_env = option_env!("CARGO_CFG_TARGET_ENV").unwrap_or("unknown");
|
||||
format!(
|
||||
"{feature} is only supported on linux x86_64 gnu. target_os={}, target_env={target_env}, target_arch={}",
|
||||
std::env::consts::OS,
|
||||
std::env::consts::ARCH
|
||||
)
|
||||
}
|
||||
|
||||
fn output_dir() -> PathBuf {
|
||||
let dir = get_env_str(ENV_OUTPUT_DIR, DEFAULT_OUTPUT_DIR);
|
||||
let p = PathBuf::from(dir);
|
||||
if let Err(e) = create_dir_all(&p) {
|
||||
warn!("profiling: create output dir {} failed: {}, fallback to current dir", p.display(), e);
|
||||
return PathBuf::from(".");
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
fn ts() -> String {
|
||||
jiff::Zoned::now().strftime("%Y%m%dT%H%M%S").to_string()
|
||||
}
|
||||
|
||||
pub async fn init_from_env() {
|
||||
let enabled = get_env_bool(ENV_ENABLE_PROFILING, DEFAULT_ENABLE_PROFILING);
|
||||
if !enabled {
|
||||
debug!("profiling: disabled by env");
|
||||
return;
|
||||
}
|
||||
|
||||
allocator::set_enabled(true);
|
||||
info!("profiling: Memory profiling enabled (mimalloc + tracing)");
|
||||
|
||||
// Initialize cancellation token
|
||||
let token = PROFILING_CANCEL_TOKEN.get_or_init(CancellationToken::new).clone();
|
||||
|
||||
// Memory periodic dump
|
||||
let mem_periodic = get_env_bool(ENV_MEM_PERIODIC, DEFAULT_MEM_PERIODIC);
|
||||
let mem_interval = Duration::from_secs(get_env_u64(ENV_MEM_INTERVAL_SECS, DEFAULT_MEM_INTERVAL_SECS));
|
||||
if mem_periodic {
|
||||
start_memory_periodic(mem_interval, token).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_memory_periodic(interval: Duration, token: CancellationToken) {
|
||||
info!(?interval, "start periodic memory pprof dump");
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
info!("periodic memory profiling task cancelled");
|
||||
break;
|
||||
}
|
||||
_ = sleep(interval) => {
|
||||
let out = output_dir().join(format!("mem_profile_periodic_{}.pb", ts()));
|
||||
match allocator::dump_profile(&out) {
|
||||
Ok(_) => info!("periodic memory profile dumped to {}", out.display()),
|
||||
Err(e) => error!("periodic mem dump failed: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Stop all background profiling tasks
|
||||
pub fn shutdown_profiling() {
|
||||
if let Some(token) = PROFILING_CANCEL_TOKEN.get() {
|
||||
token.cancel();
|
||||
}
|
||||
allocator::set_enabled(false);
|
||||
}
|
||||
|
||||
pub async fn dump_cpu_pprof_for(_duration: Duration) -> Result<PathBuf, String> {
|
||||
let (target_os, target_env, target_arch) = get_platform_info();
|
||||
let msg = format!(
|
||||
"CPU profiling is not supported on this platform. target_os={target_os}, target_env={target_env}, target_arch={target_arch}"
|
||||
);
|
||||
Err(msg)
|
||||
}
|
||||
|
||||
pub async fn dump_memory_pprof_now() -> Result<PathBuf, String> {
|
||||
let out = output_dir().join(format!("mem_profile_{}.pb", ts()));
|
||||
allocator::dump_profile(&out).map(|_| {
|
||||
info!("Memory profile exported: {}", out.display());
|
||||
out
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(
|
||||
not(target_os = "windows"),
|
||||
not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))
|
||||
))]
|
||||
pub use generic_impl::{dump_cpu_pprof_for, dump_memory_pprof_now, init_from_env, shutdown_profiling};
|
||||
#[cfg(not(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64")))]
|
||||
pub use unsupported_impl::{dump_cpu_pprof_for, dump_memory_pprof_now, init_from_env, shutdown_profiling};
|
||||
|
||||
#[cfg(all(target_os = "linux", target_env = "gnu", target_arch = "x86_64"))]
|
||||
mod linux_impl {
|
||||
|
||||
@@ -1,529 +0,0 @@
|
||||
// Copyright 2024 RustFS Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use backtrace::Backtrace;
|
||||
use pprof::protos::Message;
|
||||
use rand::RngExt;
|
||||
use starshard::ShardedHashMap;
|
||||
use std::alloc::{GlobalAlloc, Layout};
|
||||
use std::cell::Cell;
|
||||
use std::collections::HashMap;
|
||||
use std::fs::File;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, LazyLock, Weak};
|
||||
|
||||
type AllocatorShardedMap = ShardedHashMap<usize, (usize, Arc<Vec<usize>>)>;
|
||||
type LazyAllocatorShardedMap = LazyLock<AllocatorShardedMap>;
|
||||
|
||||
type AllocatorSampleHashMap = HashMap<*const Vec<usize>, (i64, i64, Arc<Vec<usize>>)>;
|
||||
/// A wrapper around a GlobalAlloc that samples allocations and records stack traces.
|
||||
pub struct TracingAllocator<A: GlobalAlloc> {
|
||||
inner: A,
|
||||
}
|
||||
|
||||
// Thread-local reentrancy guard to prevent infinite recursion when recording allocations
|
||||
thread_local! {
|
||||
static REENTRANCY_GUARD: Cell<bool> = const { Cell::new(false) };
|
||||
}
|
||||
|
||||
// Global configuration
|
||||
static SAMPLE_RATE: AtomicUsize = AtomicUsize::new(512 * 1024); // Default: sample every 512KB on average
|
||||
static ENABLED: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
// Global storage for profile data
|
||||
// Map: Address (usize) -> (Size (usize), StackTrace (Arc<Vec<usize>>))
|
||||
// We store the Arc to keep the stack trace alive as long as the allocation is live.
|
||||
static LIVE_ALLOCATIONS: LazyAllocatorShardedMap = LazyLock::new(|| ShardedHashMap::new(64));
|
||||
|
||||
// Cache for deduplicating stack traces.
|
||||
// Map: StackHash (u64) -> Weak<Vec<usize>>
|
||||
// We use Weak references so that unused stack traces can be dropped when all referring allocations are freed.
|
||||
static STACK_CACHE: LazyLock<ShardedHashMap<u64, Weak<Vec<usize>>>> = LazyLock::new(|| ShardedHashMap::new(64));
|
||||
|
||||
impl<A: GlobalAlloc> TracingAllocator<A> {
|
||||
pub const fn new(inner: A) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
// Public configuration functions
|
||||
#[allow(dead_code)]
|
||||
pub fn set_sample_rate(rate: usize) {
|
||||
SAMPLE_RATE.store(rate, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn set_enabled(enabled: bool) {
|
||||
// Force initialization of LazyLocks before enabling profiling to avoid recursion during init.
|
||||
// Accessing them is enough to trigger initialization.
|
||||
let _ = &*LIVE_ALLOCATIONS;
|
||||
let _ = &*STACK_CACHE;
|
||||
|
||||
ENABLED.store(enabled, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn should_sample(size: usize) -> bool {
|
||||
if !ENABLED.load(Ordering::Relaxed) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let rate = SAMPLE_RATE.load(Ordering::Relaxed);
|
||||
if rate == 0 {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Use a fresh RNG each time.
|
||||
let mut rng = rand::rng();
|
||||
rng.random_range(0..rate) < size
|
||||
}
|
||||
|
||||
// Internal function, assumes guard is already held
|
||||
fn record_alloc(ptr: *mut u8, size: usize) {
|
||||
// Capture stack trace
|
||||
let bt = Backtrace::new_unresolved();
|
||||
let mut frames = Vec::new();
|
||||
for frame in bt.frames() {
|
||||
frames.push(frame.symbol_address() as usize);
|
||||
}
|
||||
|
||||
// Calculate hash of the stack trace
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
frames.hash(&mut hasher);
|
||||
let stack_hash = hasher.finish();
|
||||
|
||||
// Deduplicate stack trace using STACK_CACHE
|
||||
let stack_arc = if let Some(weak) = STACK_CACHE.get(&stack_hash) {
|
||||
if let Some(arc) = weak.upgrade() {
|
||||
arc
|
||||
} else {
|
||||
// Entry exists but is dead, replace it
|
||||
let arc = Arc::new(frames);
|
||||
STACK_CACHE.insert(stack_hash, Arc::downgrade(&arc));
|
||||
arc
|
||||
}
|
||||
} else {
|
||||
// New entry
|
||||
let arc = Arc::new(frames);
|
||||
STACK_CACHE.insert(stack_hash, Arc::downgrade(&arc));
|
||||
arc
|
||||
};
|
||||
|
||||
// Store the allocation info with the Arc
|
||||
LIVE_ALLOCATIONS.insert(ptr as usize, (size, stack_arc));
|
||||
}
|
||||
|
||||
// Internal function, assumes guard is already held
|
||||
fn record_dealloc(ptr: *mut u8) {
|
||||
// Remove from live allocations.
|
||||
// The Arc<Vec<usize>> will be dropped.
|
||||
// If it was the last reference, the Vec<usize> is freed.
|
||||
// The Weak pointer in STACK_CACHE remains but becomes upgrade-able to None.
|
||||
LIVE_ALLOCATIONS.remove(&(ptr as usize));
|
||||
}
|
||||
|
||||
/// Dump the current profile to a pprof protobuf file
|
||||
pub fn dump_profile(path: &Path) -> Result<(), String> {
|
||||
// Prevent reentrancy during dump
|
||||
if REENTRANCY_GUARD.replace(true) {
|
||||
return Err("Reentrancy detected during dump".to_string());
|
||||
}
|
||||
|
||||
// Perform a lazy cleanup of the cache during dump
|
||||
cleanup_cache();
|
||||
|
||||
let result = dump_profile_inner(path);
|
||||
|
||||
REENTRANCY_GUARD.set(false);
|
||||
result
|
||||
}
|
||||
|
||||
// Clean up dead entries from STACK_CACHE
|
||||
fn cleanup_cache() {
|
||||
// We collect dead keys first to avoid locking issues during iteration if any
|
||||
let mut dead_keys = Vec::new();
|
||||
|
||||
// Note: This iteration might be slow if the cache is huge, but dump_profile is infrequent.
|
||||
for entry in STACK_CACHE.iter() {
|
||||
let (key, weak) = entry;
|
||||
if weak.upgrade().is_none() {
|
||||
dead_keys.push(key);
|
||||
}
|
||||
}
|
||||
|
||||
for key in dead_keys {
|
||||
STACK_CACHE.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
fn dump_profile_inner(path: &Path) -> Result<(), String> {
|
||||
use pprof::protos as pb;
|
||||
|
||||
let mut profile = pb::Profile::default();
|
||||
|
||||
// Basic metadata
|
||||
profile.string_table.push("".to_string()); // 0: empty
|
||||
profile.string_table.push("alloc_objects".to_string()); // 1
|
||||
profile.string_table.push("count".to_string()); // 2
|
||||
profile.string_table.push("alloc_space".to_string()); // 3
|
||||
profile.string_table.push("bytes".to_string()); // 4
|
||||
|
||||
let sample_type_count = pb::ValueType {
|
||||
ty: 1, // "alloc_objects"
|
||||
unit: 2, // "count"
|
||||
..Default::default()
|
||||
};
|
||||
let sample_type_bytes = pb::ValueType {
|
||||
ty: 3, // "alloc_space"
|
||||
unit: 4, // "bytes"
|
||||
..Default::default()
|
||||
};
|
||||
profile.sample_type = vec![sample_type_count, sample_type_bytes];
|
||||
|
||||
// Helper to get string ID
|
||||
let mut string_map: HashMap<String, i64> = HashMap::new();
|
||||
string_map.insert("".to_string(), 0);
|
||||
string_map.insert("alloc_objects".to_string(), 1);
|
||||
string_map.insert("count".to_string(), 2);
|
||||
string_map.insert("alloc_space".to_string(), 3);
|
||||
string_map.insert("bytes".to_string(), 4);
|
||||
|
||||
let mut get_string_id = |s: String| -> i64 {
|
||||
if let Some(&id) = string_map.get(&s) {
|
||||
id
|
||||
} else {
|
||||
let id = profile.string_table.len() as i64;
|
||||
profile.string_table.push(s.clone());
|
||||
string_map.insert(s, id);
|
||||
id
|
||||
}
|
||||
};
|
||||
|
||||
// Helper to get location ID
|
||||
let mut location_map: HashMap<usize, u64> = HashMap::new(); // addr -> loc_id
|
||||
let mut function_map: HashMap<usize, u64> = HashMap::new(); // addr -> func_id
|
||||
|
||||
// Collect samples
|
||||
// Aggregate by Stack Trace Pointer (deduplication via Arc pointer)
|
||||
// Map: Arc pointer -> (Count, Bytes, Arc<Vec<usize>>)
|
||||
let mut aggregated_samples: AllocatorSampleHashMap = HashMap::new();
|
||||
|
||||
// Step 1: Collect data from LIVE_ALLOCATIONS while holding the lock (implicitly via iter)
|
||||
// We do NOT perform symbol resolution here to avoid deadlocks.
|
||||
for entry in LIVE_ALLOCATIONS.iter() {
|
||||
let (_ptr, (size, stack_arc)) = entry;
|
||||
let stack_arc_clone = stack_arc.clone();
|
||||
let key = Arc::as_ptr(&stack_arc_clone);
|
||||
|
||||
let agg = aggregated_samples.entry(key).or_insert_with(|| (0, 0, stack_arc_clone));
|
||||
agg.0 += 1;
|
||||
agg.1 += size as i64;
|
||||
}
|
||||
// LIVE_ALLOCATIONS lock is released here as the iterator is dropped.
|
||||
|
||||
// Step 2: Process samples and resolve symbols (outside of LIVE_ALLOCATIONS lock)
|
||||
for (_key, (count, bytes, frames)) in aggregated_samples {
|
||||
let mut sample = pb::Sample {
|
||||
value: vec![count, bytes],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Process frames
|
||||
for &addr in frames.iter() {
|
||||
let loc_id = if let Some(&id) = location_map.get(&addr) {
|
||||
id
|
||||
} else {
|
||||
// Resolve symbol
|
||||
// This might take time and locks, but we are safe now.
|
||||
let mut func_name = "unknown".to_string();
|
||||
let mut file_name = "unknown".to_string();
|
||||
let mut line_no = 0;
|
||||
|
||||
backtrace::resolve(addr as *mut std::ffi::c_void, |symbol| {
|
||||
if let Some(name) = symbol.name() {
|
||||
func_name = name.to_string();
|
||||
}
|
||||
if let Some(filename) = symbol.filename() {
|
||||
file_name = filename.to_string_lossy().to_string();
|
||||
}
|
||||
if let Some(line) = symbol.lineno() {
|
||||
line_no = line as i64;
|
||||
}
|
||||
});
|
||||
|
||||
// Create Function
|
||||
let func_id = if let Some(&id) = function_map.get(&addr) {
|
||||
id
|
||||
} else {
|
||||
let id = (profile.function.len() + 1) as u64;
|
||||
let name_id = get_string_id(func_name);
|
||||
let file_id = get_string_id(file_name);
|
||||
|
||||
let func = pb::Function {
|
||||
id,
|
||||
name: name_id,
|
||||
system_name: name_id,
|
||||
filename: file_id,
|
||||
start_line: 0,
|
||||
..Default::default()
|
||||
};
|
||||
profile.function.push(func);
|
||||
function_map.insert(addr, id);
|
||||
id
|
||||
};
|
||||
|
||||
// Create Location
|
||||
let id = (profile.location.len() + 1) as u64;
|
||||
let line = pb::Line {
|
||||
function_id: func_id,
|
||||
line: line_no,
|
||||
..Default::default()
|
||||
};
|
||||
let loc = pb::Location {
|
||||
id,
|
||||
mapping_id: 0,
|
||||
address: addr as u64,
|
||||
line: vec![line],
|
||||
is_folded: false,
|
||||
..Default::default()
|
||||
};
|
||||
profile.location.push(loc);
|
||||
location_map.insert(addr, id);
|
||||
id
|
||||
};
|
||||
sample.location_id.push(loc_id);
|
||||
}
|
||||
profile.sample.push(sample);
|
||||
}
|
||||
|
||||
// Write to file
|
||||
let mut buf = Vec::with_capacity(1024 * 1024);
|
||||
profile.write_to_vec(&mut buf).map_err(|e| format!("encode failed: {e}"))?;
|
||||
|
||||
let mut f = File::create(path).map_err(|e| format!("create file failed: {e}"))?;
|
||||
f.write_all(&buf).map_err(|e| format!("write file failed: {e}"))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Helper to handle sampling logic
|
||||
#[inline(always)]
|
||||
fn handle_alloc_sampling(ptr: *mut u8, size: usize) {
|
||||
if !ptr.is_null() {
|
||||
// Check reentrancy guard BEFORE calling should_sample
|
||||
if !REENTRANCY_GUARD.replace(true) {
|
||||
if should_sample(size) {
|
||||
record_alloc(ptr, size);
|
||||
}
|
||||
REENTRANCY_GUARD.set(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper to handle dealloc logic
|
||||
#[inline(always)]
|
||||
fn handle_dealloc_sampling(ptr: *mut u8) {
|
||||
if !REENTRANCY_GUARD.replace(true) {
|
||||
record_dealloc(ptr);
|
||||
REENTRANCY_GUARD.set(false);
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: This allocator wrapper preserves the `GlobalAlloc` contract by
|
||||
// delegating all allocation operations to `inner` with the exact caller-provided
|
||||
// layouts and pointers. Sampling records metadata only and does not take
|
||||
// ownership of allocation memory.
|
||||
#[allow(unsafe_code)]
|
||||
unsafe impl<A: GlobalAlloc> GlobalAlloc for TracingAllocator<A> {
|
||||
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
||||
// SAFETY: Delegating to inner allocator.
|
||||
let ptr = unsafe { self.inner.alloc(layout) };
|
||||
handle_alloc_sampling(ptr, layout.size());
|
||||
ptr
|
||||
}
|
||||
|
||||
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
||||
handle_dealloc_sampling(ptr);
|
||||
// SAFETY: Delegating to inner allocator.
|
||||
unsafe { self.inner.dealloc(ptr, layout) };
|
||||
}
|
||||
|
||||
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
|
||||
// SAFETY: Delegating to inner allocator.
|
||||
let ptr = unsafe { self.inner.alloc_zeroed(layout) };
|
||||
handle_alloc_sampling(ptr, layout.size());
|
||||
ptr
|
||||
}
|
||||
|
||||
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
|
||||
handle_dealloc_sampling(ptr);
|
||||
|
||||
// SAFETY: Delegating to inner allocator.
|
||||
let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) };
|
||||
|
||||
handle_alloc_sampling(new_ptr, new_size);
|
||||
new_ptr
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
// SAFETY: Tests call the unsafe `GlobalAlloc` methods with layouts created by
|
||||
// `Layout::from_size_align` and deallocate each returned pointer with the same
|
||||
// layout.
|
||||
#[allow(unsafe_code)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serial_test::serial;
|
||||
use std::alloc::System;
|
||||
use std::thread;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
// Use System allocator for testing
|
||||
static TEST_ALLOCATOR: TracingAllocator<System> = TracingAllocator::new(System);
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_basic_allocation_tracking() {
|
||||
// Enable profiling and force sampling (rate = 1 means sample everything)
|
||||
set_enabled(true);
|
||||
set_sample_rate(1);
|
||||
|
||||
unsafe {
|
||||
let layout = Layout::from_size_align(1024, 8).unwrap();
|
||||
let ptr = TEST_ALLOCATOR.alloc(layout);
|
||||
assert!(!ptr.is_null());
|
||||
|
||||
// Verify allocation is recorded
|
||||
assert!(LIVE_ALLOCATIONS.get(&(ptr as usize)).is_some());
|
||||
|
||||
TEST_ALLOCATOR.dealloc(ptr, layout);
|
||||
|
||||
// Verify allocation is removed
|
||||
assert!(LIVE_ALLOCATIONS.get(&(ptr as usize)).is_none());
|
||||
}
|
||||
|
||||
// Reset
|
||||
set_enabled(false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_reentrancy_guard() {
|
||||
set_enabled(true);
|
||||
set_sample_rate(1);
|
||||
|
||||
// Manually set guard to simulate reentrancy
|
||||
REENTRANCY_GUARD.set(true);
|
||||
|
||||
unsafe {
|
||||
let layout = Layout::from_size_align(128, 8).unwrap();
|
||||
let ptr = TEST_ALLOCATOR.alloc(layout);
|
||||
|
||||
// Should NOT be recorded because guard was true
|
||||
assert!(LIVE_ALLOCATIONS.get(&(ptr as usize)).is_none());
|
||||
|
||||
TEST_ALLOCATOR.dealloc(ptr, layout);
|
||||
}
|
||||
|
||||
REENTRANCY_GUARD.set(false);
|
||||
set_enabled(false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_sampling_logic() {
|
||||
set_enabled(true);
|
||||
// Set a high rate so small allocations are unlikely to be sampled
|
||||
set_sample_rate(1_000_000);
|
||||
|
||||
let mut sampled_count = 0;
|
||||
let iterations = 100;
|
||||
|
||||
unsafe {
|
||||
let layout = Layout::from_size_align(8, 8).unwrap();
|
||||
for _ in 0..iterations {
|
||||
let ptr = TEST_ALLOCATOR.alloc(layout);
|
||||
if LIVE_ALLOCATIONS.get(&(ptr as usize)).is_some() {
|
||||
sampled_count += 1;
|
||||
}
|
||||
TEST_ALLOCATOR.dealloc(ptr, layout);
|
||||
}
|
||||
}
|
||||
|
||||
// With high sample rate and small size, sampled count should be low (likely 0)
|
||||
// This is probabilistic, but 0 is very likely.
|
||||
assert!(sampled_count < iterations);
|
||||
|
||||
set_enabled(false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_profile_dump() {
|
||||
set_enabled(true);
|
||||
// Use a larger sample rate to avoid capturing too much noise from the test runner
|
||||
// and ensure we only capture our large allocation.
|
||||
set_sample_rate(1024 * 1024);
|
||||
|
||||
unsafe {
|
||||
// Allocate a large enough chunk to likely be sampled (2MB > 1MB rate)
|
||||
let layout = Layout::from_size_align(2 * 1024 * 1024, 8).unwrap();
|
||||
let ptr = TEST_ALLOCATOR.alloc(layout);
|
||||
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let path = file.path();
|
||||
|
||||
let result = dump_profile(path);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let metadata = std::fs::metadata(path).unwrap();
|
||||
assert!(metadata.len() > 0);
|
||||
|
||||
TEST_ALLOCATOR.dealloc(ptr, layout);
|
||||
}
|
||||
set_enabled(false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_concurrent_allocations() {
|
||||
set_enabled(true);
|
||||
set_sample_rate(1);
|
||||
|
||||
let threads: Vec<_> = (0..10)
|
||||
.map(|_| {
|
||||
thread::spawn(|| {
|
||||
unsafe {
|
||||
let layout = Layout::from_size_align(64, 8).unwrap();
|
||||
for _ in 0..100 {
|
||||
let ptr = TEST_ALLOCATOR.alloc(layout);
|
||||
// Just ensure no panic/crash
|
||||
TEST_ALLOCATOR.dealloc(ptr, layout);
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
for t in threads {
|
||||
t.join().unwrap();
|
||||
}
|
||||
|
||||
// After all threads join and dealloc, map should be empty (ignoring other potential allocations in test runner)
|
||||
// Note: In a real test runner, other tests might be running, so we can't assert empty.
|
||||
// But we verified no crashes.
|
||||
set_enabled(false);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user