[AdminService] Minor improvements for admin service. (aptos-labs#10866)

zi0Black · Nov 10, 2023 · 008de22 · 008de22
1 parent d41b8c5
commit 008de22
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/aptos-admin-service/Cargo.toml b/crates/aptos-admin-service/Cargo.toml
@@ -37,4 +37,5 @@ url = { workspace = true }
 [target.'cfg(target_os = "linux")'.dependencies]
 aptos-profiler = { workspace = true }
 pprof = { workspace = true }
+regex = { workspace = true }
 rstack-self = { workspace = true }
diff --git a/crates/aptos-admin-service/src/server/profiling.rs b/crates/aptos-admin-service/src/server/profiling.rs
@@ -9,6 +9,7 @@ use http::header::{HeaderValue, CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYP
 use hyper::{Body, Request, Response, StatusCode};
 use lazy_static::lazy_static;
 use pprof::protos::Message;
+use regex::Regex;
 use std::{collections::HashMap, time::Duration};
 
 lazy_static! {
@@ -100,6 +101,7 @@ async fn start_cpu_profiling(
     let mut body = Vec::new();
     let report = guard
         .report()
+        .frames_post_processor(frames_post_processor())
         .build()
         .map_err(|e| anyhow!("Failed to generate cpu profiling report: {e:?}."))?;
 
@@ -117,3 +119,13 @@ async fn start_cpu_profiling(
 
     Ok(body)
 }
+
+fn frames_post_processor() -> impl Fn(&mut pprof::Frames) {
+    let regex = Regex::new(r"^(.*)-(\d*)$").unwrap();
+
+    move |frames| {
+        if let Some((_, [name, _])) = regex.captures(&frames.thread_name).map(|c| c.extract()) {
+            frames.thread_name = name.to_string();
+        }
+    }
+}
diff --git a/crates/aptos-admin-service/src/server/thread_dump.rs b/crates/aptos-admin-service/src/server/thread_dump.rs
@@ -15,6 +15,8 @@ lazy_static! {
     static ref THREAD_DUMP_MUTEX: Mutex<()> = Mutex::new(());
 }
 
+static MAX_NUM_FRAMES_WITHOUT_VERBOSE: usize = 20;
+
 pub async fn handle_thread_dump_request(req: Request<Body>) -> hyper::Result<Response<Body>> {
     let query = req.uri().query().unwrap_or("");
     let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect();
@@ -43,8 +45,16 @@ pub async fn handle_thread_dump_request(req: Request<Body>) -> hyper::Result<Res
         None => false,
     };
 
+    let verbose: bool = match query_pairs.get("verbose") {
+        Some(val) => match val.parse() {
+            Ok(val) => val,
+            Err(err) => return Ok(reply_with_status(StatusCode::BAD_REQUEST, err.to_string())),
+        },
+        None => false,
+    };
+
     info!("Starting dumping stack trace for all threads.");
-    match start_thread_dump(snapshot, location, frame_ip).await {
+    match do_thread_dump(snapshot, location, frame_ip, verbose).await {
         Ok(body) => {
             info!("Thread dumping is done.");
             let headers: Vec<(_, HeaderValue)> =
@@ -61,10 +71,11 @@ pub async fn handle_thread_dump_request(req: Request<Body>) -> hyper::Result<Res
     }
 }
 
-async fn start_thread_dump(
+async fn do_thread_dump(
     snapshot: bool,
     location: bool,
     frame_ip: bool,
+    verbose: bool,
 ) -> anyhow::Result<String> {
     let lock = THREAD_DUMP_MUTEX.try_lock();
     ensure!(lock.is_some(), "A thread dumping task is already running.");
@@ -75,10 +86,49 @@ async fn start_thread_dump(
         .trace(Command::new(exe).arg("--stacktrace"))
         .map_err(Error::msg)?;
 
+    let mut wait_threads = Vec::new();
+    let mut sleep_threads = Vec::new();
     let mut body = String::new();
     for thread in trace.threads() {
+        let frames = thread.frames();
+        if !frames.is_empty() {
+            let symbols = frames[0].symbols();
+            if !symbols.is_empty() {
+                if let Some(name) = symbols[0].name() {
+                    if name.contains("epoll_wait") {
+                        wait_threads.push(thread.name());
+                        continue;
+                    }
+
+                    if name.contains("clock_nanosleep") {
+                        sleep_threads.push(thread.name());
+                        continue;
+                    }
+                }
+            }
+        }
+
+        if frames.len() > 1 {
+            let symbols = frames[1].symbols();
+            if !symbols.is_empty() {
+                if let Some(name) = symbols[0].name() {
+                    if name.contains("futex_wait")
+                        || name.contains("pthread_cond_wait")
+                        || name.contains("pthread_cond_timedwait")
+                    {
+                        wait_threads.push(thread.name());
+                        continue;
+                    }
+                }
+            }
+        }
+
         body.push_str(&format!("Thread {} ({}):\n", thread.id(), thread.name()));
-        for frame in thread.frames() {
+        for (count, frame) in frames.iter().enumerate() {
+            if !verbose && count >= MAX_NUM_FRAMES_WITHOUT_VERBOSE {
+                break;
+            }
+
             if frame_ip {
                 body.push_str(&format!("Frame ip: {}\n", frame.ip()));
             }
@@ -107,5 +157,17 @@ async fn start_thread_dump(
         body.push_str("\n\n");
     }
 
+    body.push_str("Wait threads:");
+    for wait_thread in wait_threads {
+        body.push_str(&format!(" {wait_thread}"));
+    }
+    body.push_str("\n\n");
+
+    body.push_str("Sleep threads:");
+    for sleep_thread in sleep_threads {
+        body.push_str(&format!(" {sleep_thread}"));
+    }
+    body.push_str("\n\n");
+
     Ok(body)
 }