From 3a0d9cdadb8a0dbfd180367459721d13eab0e116 Mon Sep 17 00:00:00 2001
From: Ben Niu <beniu@microsoft.com>
Date: Fri, 30 Aug 2024 17:32:42 -0700
Subject: [PATCH] Use MSVC __declspec(thread) for TSD on Windows

---
 include/jemalloc/internal/tsd_win.h | 74 ++++++++++++++++++++++++++++-
 src/tsd.c                           |  6 ++-
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/include/jemalloc/internal/tsd_win.h b/include/jemalloc/internal/tsd_win.h
index 4b40a8ab..559ee78f 100644
--- a/include/jemalloc/internal/tsd_win.h
+++ b/include/jemalloc/internal/tsd_win.h
@@ -7,11 +7,17 @@
 #include "jemalloc/internal/tsd_internals.h"
 #include "jemalloc/internal/tsd_types.h"
 
+/* val should always be the first field of tsd_wrapper_t since accessing
+   val is the common path and having val as the first field makes it possible
+   that converting a pointer to tsd_wrapper_t to a pointer to val is no more
+   than a type cast. */
 typedef struct {
-	bool initialized;
 	tsd_t val;
+	bool initialized;
 } tsd_wrapper_t;
 
+#if defined(JEMALLOC_LEGACY_WINDOWS_SUPPORT) || !defined(_MSC_VER)
+
 extern DWORD tsd_tsd;
 extern tsd_wrapper_t tsd_boot_wrapper;
 extern bool tsd_booted;
@@ -165,3 +171,69 @@ tsd_set(tsd_t *val) {
 	}
 	wrapper->initialized = true;
 }
+
+#else // defined(JEMALLOC_LEGACY_WINDOWS_SUPPORT) || !defined(_MSC_VER)
+
+#define JEMALLOC_TSD_TYPE_ATTR(type) __declspec(thread) type
+
+extern JEMALLOC_TSD_TYPE_ATTR(tsd_wrapper_t) tsd_wrapper_tls;
+extern bool tsd_booted;
+
+/* Initialization/cleanup. */
+JEMALLOC_ALWAYS_INLINE bool
+tsd_cleanup_wrapper(void) {
+	if (tsd_wrapper_tls.initialized) {
+		tsd_wrapper_tls.initialized = false;
+		tsd_cleanup(&tsd_wrapper_tls.val);
+		if (tsd_wrapper_tls.initialized) {
+			/* Trigger another cleanup round. */
+			return true;
+		}
+	}
+	return false;
+}
+
+/* Initialization/cleanup. */
+JEMALLOC_ALWAYS_INLINE bool
+tsd_boot0(void) {
+	_malloc_tsd_cleanup_register(tsd_cleanup_wrapper);
+	tsd_booted = true;
+	return false;
+}
+
+JEMALLOC_ALWAYS_INLINE void
+tsd_boot1(void) {
+	/* Do nothing. */
+}
+
+JEMALLOC_ALWAYS_INLINE bool
+tsd_boot(void) {
+	return tsd_boot0();
+}
+
+JEMALLOC_ALWAYS_INLINE bool
+tsd_booted_get(void) {
+	return tsd_booted;
+}
+
+JEMALLOC_ALWAYS_INLINE bool
+tsd_get_allocates(void) {
+	return false;
+}
+
+/* Get/set. */
+JEMALLOC_ALWAYS_INLINE tsd_t *
+tsd_get(bool init) {
+	return &(tsd_wrapper_tls.val);
+}
+
+JEMALLOC_ALWAYS_INLINE void
+tsd_set(tsd_t *val) {
+	assert(tsd_booted);
+	if (likely(&(tsd_wrapper_tls.val) != val)) {
+		tsd_wrapper_tls.val = (*val);
+	}
+	tsd_wrapper_tls.initialized = true;
+}
+
+#endif // defined(JEMALLOC_LEGACY_WINDOWS_SUPPORT) || !defined(_MSC_VER)
diff --git a/src/tsd.c b/src/tsd.c
index c9ae2d64..0a2ccc59 100644
--- a/src/tsd.c
+++ b/src/tsd.c
@@ -22,8 +22,12 @@ JEMALLOC_TSD_TYPE_ATTR(tsd_t) tsd_tls = TSD_INITIALIZER;
 pthread_key_t tsd_tsd;
 bool tsd_booted = false;
 #elif (defined(_WIN32))
+#if defined(JEMALLOC_LEGACY_WINDOWS_SUPPORT) || !defined(_MSC_VER)
 DWORD tsd_tsd;
-tsd_wrapper_t tsd_boot_wrapper = {false, TSD_INITIALIZER};
+tsd_wrapper_t tsd_boot_wrapper = {TSD_INITIALIZER, false};
+#else
+JEMALLOC_TSD_TYPE_ATTR(tsd_wrapper_t) tsd_wrapper_tls = { TSD_INITIALIZER, false };
+#endif
 bool tsd_booted = false;
 #if JEMALLOC_WIN32_TLSGETVALUE2
 TGV2 tls_get_value2 = NULL;