From eb65d1b07830b285bf7ac7678e964f080cd3916a Mon Sep 17 00:00:00 2001
From: Alex Lapenkov <lapenkov.a@yandex.ru>
Date: Sat, 22 Jan 2022 10:14:16 -0800
Subject: [PATCH] Fix FreeBSD system jemalloc TSD cleanup

Before this commit, in case FreeBSD libc jemalloc was overridden by another
jemalloc, proper thread shutdown callback was involved only for the overriding
jemalloc. A call to _malloc_thread_cleanup from libthr would be redirected to
user jemalloc, leaving data about dead threads hanging in system jemalloc. This
change tackles the issue in two ways. First, for current and old system
jemallocs, which we can not modify, the overriding jemalloc would locate and
invoke system cleanup routine. For upcoming jemalloc integrations, the cleanup
registering function will also be redirected to user jemalloc, which means that
system jemalloc's cleanup routine will be registered in user's jemalloc and a
single call to _malloc_thread_cleanup will be sufficient to invoke both
callbacks.
---
 configure.ac                                     |  2 +-
 include/jemalloc/internal/tsd.h                  |  5 ++++-
 .../internal/tsd_malloc_thread_cleanup.h         |  2 +-
 include/jemalloc/internal/tsd_types.h            |  2 +-
 include/jemalloc/internal/tsd_win.h              |  2 +-
 src/tsd.c                                        | 16 +++++++++++-----
 6 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/configure.ac b/configure.ac
index abcd91d2..69b8162f 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1999,7 +1999,7 @@ AC_CHECK_FUNC([_malloc_thread_cleanup],
              )
 if test "x$have__malloc_thread_cleanup" = "x1" ; then
   AC_DEFINE([JEMALLOC_MALLOC_THREAD_CLEANUP], [ ], [ ])
-  wrap_syms="${wrap_syms} _malloc_thread_cleanup"
+  wrap_syms="${wrap_syms} _malloc_thread_cleanup _malloc_tsd_cleanup_register"
   force_tls="1"
 fi
 
diff --git a/include/jemalloc/internal/tsd.h b/include/jemalloc/internal/tsd.h
index 0a46d448..66d68822 100644
--- a/include/jemalloc/internal/tsd.h
+++ b/include/jemalloc/internal/tsd.h
@@ -153,9 +153,12 @@ typedef ql_elm(tsd_t) tsd_link_t;
     				TSD_DATA_SLOWER_INITIALIZER		\
 }
 
+#if defined(JEMALLOC_MALLOC_THREAD_CLEANUP) || defined(_WIN32)
+void _malloc_tsd_cleanup_register(bool (*f)(void));
+#endif
+
 void *malloc_tsd_malloc(size_t size);
 void malloc_tsd_dalloc(void *wrapper);
-void malloc_tsd_cleanup_register(bool (*f)(void));
 tsd_t *malloc_tsd_boot0(void);
 void malloc_tsd_boot1(void);
 void tsd_cleanup(void *arg);
diff --git a/include/jemalloc/internal/tsd_malloc_thread_cleanup.h b/include/jemalloc/internal/tsd_malloc_thread_cleanup.h
index 65852d5c..d8f3ef13 100644
--- a/include/jemalloc/internal/tsd_malloc_thread_cleanup.h
+++ b/include/jemalloc/internal/tsd_malloc_thread_cleanup.h
@@ -21,7 +21,7 @@ tsd_cleanup_wrapper(void) {
 
 JEMALLOC_ALWAYS_INLINE bool
 tsd_boot0(void) {
-	malloc_tsd_cleanup_register(&tsd_cleanup_wrapper);
+	_malloc_tsd_cleanup_register(&tsd_cleanup_wrapper);
 	tsd_booted = true;
 	return false;
 }
diff --git a/include/jemalloc/internal/tsd_types.h b/include/jemalloc/internal/tsd_types.h
index 6200af61..a6ae37da 100644
--- a/include/jemalloc/internal/tsd_types.h
+++ b/include/jemalloc/internal/tsd_types.h
@@ -1,7 +1,7 @@
 #ifndef JEMALLOC_INTERNAL_TSD_TYPES_H
 #define JEMALLOC_INTERNAL_TSD_TYPES_H
 
-#define MALLOC_TSD_CLEANUPS_MAX	2
+#define MALLOC_TSD_CLEANUPS_MAX	4
 
 typedef struct tsd_s tsd_t;
 typedef struct tsdn_s tsdn_t;
diff --git a/include/jemalloc/internal/tsd_win.h b/include/jemalloc/internal/tsd_win.h
index cf30d18e..a91dac88 100644
--- a/include/jemalloc/internal/tsd_win.h
+++ b/include/jemalloc/internal/tsd_win.h
@@ -72,7 +72,7 @@ tsd_boot0(void) {
 	if (tsd_tsd == TLS_OUT_OF_INDEXES) {
 		return true;
 	}
-	malloc_tsd_cleanup_register(&tsd_cleanup_wrapper);
+	_malloc_tsd_cleanup_register(&tsd_cleanup_wrapper);
 	tsd_wrapper_set(&tsd_boot_wrapper);
 	tsd_booted = true;
 	return false;
diff --git a/src/tsd.c b/src/tsd.c
index 4859048e..b98c34bf 100644
--- a/src/tsd.c
+++ b/src/tsd.c
@@ -9,9 +9,6 @@
 /******************************************************************************/
 /* Data. */
 
-static unsigned ncleanups;
-static malloc_tsd_cleanup_t cleanups[MALLOC_TSD_CLEANUPS_MAX];
-
 /* TSD_INITIALIZER triggers "-Wmissing-field-initializer" */
 JEMALLOC_DIAGNOSTIC_PUSH
 JEMALLOC_DIAGNOSTIC_IGNORE_MISSING_STRUCT_FIELD_INITIALIZERS
@@ -337,6 +334,9 @@ malloc_tsd_dalloc(void *wrapper) {
 }
 
 #if defined(JEMALLOC_MALLOC_THREAD_CLEANUP) || defined(_WIN32)
+static unsigned ncleanups;
+static malloc_tsd_cleanup_t cleanups[MALLOC_TSD_CLEANUPS_MAX];
+
 #ifndef _WIN32
 JEMALLOC_EXPORT
 #endif
@@ -361,15 +361,19 @@ _malloc_thread_cleanup(void) {
 		}
 	} while (again);
 }
-#endif
 
+#ifndef _WIN32
+JEMALLOC_EXPORT
+#endif
 void
-malloc_tsd_cleanup_register(bool (*f)(void)) {
+_malloc_tsd_cleanup_register(bool (*f)(void)) {
 	assert(ncleanups < MALLOC_TSD_CLEANUPS_MAX);
 	cleanups[ncleanups] = f;
 	ncleanups++;
 }
 
+#endif
+
 static void
 tsd_do_data_cleanup(tsd_t *tsd) {
 	prof_tdata_cleanup(tsd);
@@ -429,7 +433,9 @@ tsd_t *
 malloc_tsd_boot0(void) {
 	tsd_t *tsd;
 
+#if defined(JEMALLOC_MALLOC_THREAD_CLEANUP) || defined(_WIN32)
 	ncleanups = 0;
+#endif
 	if (malloc_mutex_init(&tsd_nominal_tsds_lock, "tsd_nominal_tsds_lock",
 	    WITNESS_RANK_OMIT, malloc_mutex_rank_exclusive)) {
 		return NULL;