Fixes

Author: squadgazzz

.github/workflows/deploy-profiling.yaml

@@ -6,6 +6,9 @@ on:
         description: 'Tag suffix for the profiling image (e.g., "profiling" or "profiling-v1.2.3")'
         required: false
         default: 'profiling'
+  push:
+    branches:
+      - jemalloc-profiling-support
 
 jobs:
   deploy-profiling:

Cargo.toml

@@ -46,7 +46,7 @@ serde_json = "1.0.117"
 serde_with = "3.8.1"
 sqlx = { version = "0.7", default-features = false, features = ["runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "postgres", "macros"] }
 strum = { version = "0.27.2", features = ["derive"] }
-tempfile = "3.11"
+tempfile = "3.10.1"
 thiserror = "1.0.61"
 toml = "0.8.14"
 tokio = { version = "1.38.0", features = ["tracing"] }

crates/driver/src/infra/observe/mod.rs

@@ -38,6 +38,8 @@ pub mod metrics;
 pub fn init(obs_config: observe::Config) {
     observe::tracing::initialize_reentrant(&obs_config);
     metrics::init();
+    #[cfg(all(unix, feature = "jemalloc-profiling"))]
+    observe::heap_dump_handler::spawn_heap_dump_handler();
 }
 
 /// Observe a received auction.

crates/driver/src/run.rs

@@ -27,8 +27,6 @@ use {
 /// driver from multiple binaries.
 pub async fn start(args: impl Iterator<Item = String>) {
     observe::panic_hook::install();
-    #[cfg(all(unix, feature = "jemalloc-profiling"))]
-    observe::heap_dump_handler::spawn_heap_dump_handler();
     let args = cli::Args::parse_from(args);
     run_with(args, None).await
 }

crates/observe/src/heap_dump_handler.rs

@@ -1,9 +1,10 @@
 use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
     net::{UnixListener, UnixStream},
+    time::{Duration, timeout},
 };
 
-/// Spawns a new thread that listens for connections to a UNIX socket
+/// Spawns a new async task that listens for connections to a UNIX socket
 /// at "/tmp/heap_dump_<process_name>.sock".
 /// When "dump" command is sent, it generates a heap profile using
 /// jemalloc_pprof and streams the binary protobuf data back through the socket.
@@ -31,7 +32,7 @@ pub fn spawn_heap_dump_handler() {
     }
 
     tokio::spawn(async move {
-        let name = binary_name().unwrap_or_default();
+        let name = binary_name().unwrap_or("unknown".to_string());
         let socket_path = format!("/tmp/heap_dump_{name}.sock");
 
         tracing::info!(socket = socket_path, "heap dump handler started");
@@ -43,16 +44,29 @@ pub fn spawn_heap_dump_handler() {
         };
 
         loop {
-            // Catch any panics in connection handling to prevent the handler from crashing
-            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
-                let listener = &handle.listener;
-                async move { handle_connection(listener).await }
-            }));
-
-            match result {
-                Ok(future) => future.await,
+            // Accept connection in main loop, then spawn task for each connection
+            // Sequential processing prevents multiple simultaneous expensive heap dumps
+            match handle.listener.accept().await {
+                Ok((socket, _addr)) => {
+                    let handle = tokio::spawn(async move {
+                        handle_connection_with_socket(socket).await;
+                    });
+
+                    // 5-minute timeout to prevent stuck dumps from blocking future requests
+                    match timeout(Duration::from_secs(300), handle).await {
+                        Ok(Ok(())) => {
+                            // Task completed successfully
+                        }
+                        Ok(Err(err)) => {
+                            tracing::error!(?err, "panic in heap dump connection handler");
+                        }
+                        Err(_) => {
+                            tracing::error!("heap dump request timed out after 5 minutes");
+                        }
+                    }
+                }
                 Err(err) => {
-                    tracing::error!(?err, "panic in heap dump connection handler");
+                    tracing::debug!(?err, "failed to accept connection");
                 }
             }
         }
@@ -80,12 +94,7 @@ fn binary_name() -> Option<String> {
     )
 }
 
-async fn handle_connection(listener: &UnixListener) {
-    let Ok((mut socket, _addr)) = listener.accept().await else {
-        tracing::debug!("failed to accept connection");
-        return;
-    };
-
+async fn handle_connection_with_socket(mut socket: UnixStream) {
     let message = read_line(&mut socket).await;
     match message.as_deref() {
         Some("dump") => {

crates/solvers/src/run.rs

@@ -12,8 +12,6 @@ use {
 
 pub async fn start(args: impl IntoIterator<Item = String>) {
     observe::panic_hook::install();
-    #[cfg(all(unix, feature = "jemalloc-profiling"))]
-    observe::heap_dump_handler::spawn_heap_dump_handler();
     let args = cli::Args::parse_from(args);
     run_with(args, None).await;
 }
@@ -34,6 +32,8 @@ async fn run_with(args: cli::Args, bind: Option<oneshot::Sender<SocketAddr>>) {
         None,
     );
     observe::tracing::initialize_reentrant(&obs_config);
+    #[cfg(all(unix, feature = "jemalloc-profiling"))]
+    observe::heap_dump_handler::spawn_heap_dump_handler();
 
     let commit_hash = option_env!("VERGEN_GIT_SHA").unwrap_or("COMMIT_INFO_NOT_FOUND");