From 639e70fcfba18fe67ce5154f2fe52cd753012fd8 Mon Sep 17 00:00:00 2001
From: Slobodan Predolac <spredolac@fb.com>
Date: Tue, 12 May 2026 16:16:37 -0700
Subject: [PATCH] Make TSD state non-atomic

---
 include/jemalloc/internal/tsd_internals.h | 27 ++++-------------------
 src/tsd.c                                 | 10 ++++-----
 2 files changed, 8 insertions(+), 29 deletions(-)

diff --git a/include/jemalloc/internal/tsd_internals.h b/include/jemalloc/internal/tsd_internals.h
index e80b2764..46b4930f 100644
--- a/include/jemalloc/internal/tsd_internals.h
+++ b/include/jemalloc/internal/tsd_internals.h
@@ -126,7 +126,7 @@ typedef void (*test_callback_t)(int *);
 #define TSD_INITIALIZER                                                        \
 	{                                                                      \
 		TSD_DATA_SLOW_INITIALIZER                                      \
-		/* state */ ATOMIC_INIT(tsd_state_uninitialized),              \
+		/* state */ tsd_state_uninitialized,                           \
 		    TSD_DATA_FAST_INITIALIZER TSD_DATA_SLOWER_INITIALIZER      \
 	}
 
@@ -182,16 +182,6 @@ enum {
  */
 #define TSD_MANGLE(n) cant_access_tsd_items_directly_use_a_getter_or_setter_##n
 
-#ifdef JEMALLOC_U8_ATOMICS
-#	define tsd_state_t atomic_u8_t
-#	define tsd_atomic_load atomic_load_u8
-#	define tsd_atomic_store atomic_store_u8
-#else
-#	define tsd_state_t atomic_u32_t
-#	define tsd_atomic_load atomic_load_u32
-#	define tsd_atomic_store atomic_store_u32
-#endif
-
 /* The actual tsd. */
 struct tsd_s {
 	/*
@@ -203,11 +193,8 @@ struct tsd_s {
 #define O(n, t, nt) t TSD_MANGLE(n);
 
 	TSD_DATA_SLOW
-	/*
-	 * We manually limit the state to just a single byte.  Unless the 8-bit
-	 * atomics are unavailable (which is rare).
-	 */
-	tsd_state_t state;
+	/* Encodes one of tsd_state_*; mutated only by the owning thread. */
+	uint8_t state;
 	TSD_DATA_FAST
 	TSD_DATA_SLOWER
 #undef O
@@ -215,13 +202,7 @@ struct tsd_s {
 
 JEMALLOC_ALWAYS_INLINE uint8_t
 tsd_state_get(tsd_t *tsd) {
-	/*
-	 * This should be atomic.  Unfortunately, compilers right now can't tell
-	 * that this can be done as a memory comparison, and forces a load into
-	 * a register that hurts fast-path performance.
-	 */
-	/* return atomic_load_u8(&tsd->state, ATOMIC_RELAXED); */
-	return *(uint8_t *)&tsd->state;
+	return tsd->state;
 }
 
 /*
diff --git a/src/tsd.c b/src/tsd.c
index ff52c9cd..67200f50 100644
--- a/src/tsd.c
+++ b/src/tsd.c
@@ -81,7 +81,7 @@ void
 tsd_slow_update(tsd_t *tsd) {
 	assert(!tsd_booted_get() || tsd_get(false) == tsd
 	    || (tsd_get_allocates() && tsd_get(false) == NULL));
-	tsd_atomic_store(&tsd->state, tsd_state_compute(tsd), ATOMIC_RELAXED);
+	tsd->state = tsd_state_compute(tsd);
 
 	te_recompute_fast_threshold(tsd);
 }
@@ -90,8 +90,7 @@ void
 tsd_state_set(tsd_t *tsd, uint8_t new_state) {
 	assert(!tsd_booted_get() || tsd_get(false) == tsd
 	    || (tsd_get_allocates() && tsd_get(false) == NULL));
-	uint8_t old_state = tsd_atomic_load(&tsd->state, ATOMIC_RELAXED);
-	if (old_state <= tsd_state_nominal_max
+	if (tsd->state <= tsd_state_nominal_max
 	    && new_state <= tsd_state_nominal_max) {
 		/*
 		 * We're transitioning from one nominal state to another.
@@ -99,11 +98,10 @@ tsd_state_set(tsd_t *tsd, uint8_t new_state) {
 		 * than trusting the caller's requested nominal state.
 		 */
 		tsd_slow_update(tsd);
-		return;
 	} else {
-		tsd_atomic_store(&tsd->state, new_state, ATOMIC_RELAXED);
+		tsd->state = new_state;
+		te_recompute_fast_threshold(tsd);
 	}
-	te_recompute_fast_threshold(tsd);
 }
 
 static void