From 39938fa7e641ac52394d5b8b5a4f034d3b1a4552 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Tue, 22 Apr 2025 06:29:30 +0000
Subject: [PATCH] Join jemalloc background threads prior to exit.

Co-authored-by: Jade Ellis <jade@ellis.link>
Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/core/alloc/je.rs |  4 ++++
 src/main/runtime.rs  | 19 ++++++++++++-------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/core/alloc/je.rs b/src/core/alloc/je.rs
index 2424e99c..e138233e 100644
--- a/src/core/alloc/je.rs
+++ b/src/core/alloc/je.rs
@@ -274,6 +274,10 @@ pub fn set_dirty_decay<I: Into<Option<usize>>>(arena: I, decay_ms: isize) -> Res
 	}
 }
 
+pub fn background_thread_enable(enable: bool) -> Result<bool> {
+	set::<u8>(&mallctl!("background_thread"), enable.into()).map(is_nonzero!())
+}
+
 #[inline]
 #[must_use]
 pub fn is_affine_arena() -> bool { is_percpu_arena() || is_phycpu_arena() }
diff --git a/src/main/runtime.rs b/src/main/runtime.rs
index 1c58ea81..e9029012 100644
--- a/src/main/runtime.rs
+++ b/src/main/runtime.rs
@@ -98,12 +98,7 @@ pub(super) fn shutdown(server: &Arc<Server>, runtime: tokio::runtime::Runtime) {
 		Level::INFO
 	};
 
-	debug!(
-		timeout = ?SHUTDOWN_TIMEOUT,
-		"Waiting for runtime..."
-	);
-
-	runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+	wait_shutdown(server, runtime);
 	let runtime_metrics = server.server.metrics.runtime_interval().unwrap_or_default();
 
 	event!(LEVEL, ?runtime_metrics, "Final runtime metrics");
@@ -111,13 +106,23 @@ pub(super) fn shutdown(server: &Arc<Server>, runtime: tokio::runtime::Runtime) {
 
 #[cfg(not(tokio_unstable))]
 #[tracing::instrument(name = "stop", level = "info", skip_all)]
-pub(super) fn shutdown(_server: &Arc<Server>, runtime: tokio::runtime::Runtime) {
+pub(super) fn shutdown(server: &Arc<Server>, runtime: tokio::runtime::Runtime) {
+	wait_shutdown(server, runtime);
+}
+
+fn wait_shutdown(_server: &Arc<Server>, runtime: tokio::runtime::Runtime) {
 	debug!(
 		timeout = ?SHUTDOWN_TIMEOUT,
 		"Waiting for runtime..."
 	);
 
 	runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+	// Join any jemalloc threads so they don't appear in use at exit.
+	#[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))]
+	conduwuit_core::alloc::je::background_thread_enable(false)
+		.log_debug_err()
+		.ok();
 }
 
 #[tracing::instrument(