gh-123022: Fix crash with `Py_Initialize` in background thread (#123052)

Check that the current default heap is initialized in
`_mi_os_get_aligned_hint` and `mi_os_claim_huge_pages`.

The mimalloc function `_mi_os_get_aligned_hint` assumes that there is an
initialized default heap. This is true for our main thread, but not for
background threads. The problematic code path is usually called during
initialization (i.e., `Py_Initialize`), but it may also be called if the
program allocates large amounts of memory in total.

The crash only affected the free-threaded build.

4 files changed, 40 insertions, 5 deletions

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index aab4333..7860c67 100644
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -1,6 +1,6 @@
 # Run the tests in Programs/_testembed.c (tests for the CPython embedding APIs)
 from test import support
-from test.support import import_helper, os_helper, MS_WINDOWS
+from test.support import import_helper, os_helper, threading_helper, MS_WINDOWS
 import unittest
 
 from collections import namedtuple
@@ -1802,6 +1802,13 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
         self.assertEqual(out, expected)
 
+    @threading_helper.requires_working_threading()
+    def test_init_in_background_thread(self):
+        # gh-123022: Check that running Py_Initialize() in a background
+        # thread doesn't crash.
+        out, err = self.run_embedded_interpreter("test_init_in_background_thread")
+        self.assertEqual(err, "")
+
 
 class SetConfigTests(unittest.TestCase):
     def test_set_config(self):
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-08-15-19-28-43.gh-issue-123022.m3EF9E.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-08-15-19-28-43.gh-issue-123022.m3EF9E.rst
new file mode 100644
index 0000000..47107de
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-08-15-19-28-43.gh-issue-123022.m3EF9E.rst
@@ -0,0 +1,2 @@
+Fix crash in free-threaded build when calling :c:func:`Py_Initialize` from
+a non-main thread.
diff --git a/Objects/mimalloc/os.c b/Objects/mimalloc/os.c
index f3bc718..c910316 100644
--- a/Objects/mimalloc/os.c
+++ b/Objects/mimalloc/os.c
@@ -115,8 +115,12 @@ void* _mi_os_get_aligned_hint(size_t try_alignment, size_t size)
   if (hint == 0 || hint > MI_HINT_MAX) {   // wrap or initialize
     uintptr_t init = MI_HINT_BASE;
     #if (MI_SECURE>0 || MI_DEBUG==0)       // security: randomize start of aligned allocations unless in debug mode
-    uintptr_t r = _mi_heap_random_next(mi_prim_get_default_heap());
-    init = init + ((MI_SEGMENT_SIZE * ((r>>17) & 0xFFFFF)) % MI_HINT_AREA);  // (randomly 20 bits)*4MiB == 0 to 4TiB
+    mi_heap_t* heap = mi_prim_get_default_heap();
+    // gh-123022: default heap may not be initialized in CPython in background threads
+    if (mi_heap_is_initialized(heap)) {
+      uintptr_t r = _mi_heap_random_next(heap);
+      init = init + ((MI_SEGMENT_SIZE * ((r>>17) & 0xFFFFF)) % MI_HINT_AREA);  // (randomly 20 bits)*4MiB == 0 to 4TiB
+    }
     #endif
     uintptr_t expected = hint + size;
     mi_atomic_cas_strong_acq_rel(&aligned_base, &expected, init);
@@ -553,8 +557,12 @@ static uint8_t* mi_os_claim_huge_pages(size_t pages, size_t* total_size) {
       // Initialize the start address after the 32TiB area
       start = ((uintptr_t)32 << 40);  // 32TiB virtual start address
     #if (MI_SECURE>0 || MI_DEBUG==0)      // security: randomize start of huge pages unless in debug mode
-      uintptr_t r = _mi_heap_random_next(mi_prim_get_default_heap());
-      start = start + ((uintptr_t)MI_HUGE_OS_PAGE_SIZE * ((r>>17) & 0x0FFF));  // (randomly 12bits)*1GiB == between 0 to 4TiB
+      mi_heap_t* heap = mi_prim_get_default_heap();
+      // gh-123022: default heap may not be initialized in CPython in background threads
+      if (mi_heap_is_initialized(heap)) {
+        uintptr_t r = _mi_heap_random_next(heap);
+        start = start + ((uintptr_t)MI_HUGE_OS_PAGE_SIZE * ((r>>17) & 0x0FFF));  // (randomly 12bits)*1GiB == between 0 to 4TiB
+      }
     #endif
     }
     end = start + size;
diff --git a/Programs/_testembed.c b/Programs/_testembed.c
index 2c726c9..e341f0c 100644
--- a/Programs/_testembed.c
+++ b/Programs/_testembed.c
@@ -8,6 +8,7 @@
 #include <Python.h>
 #include "pycore_initconfig.h"    // _PyConfig_InitCompatConfig()
 #include "pycore_runtime.h"       // _PyRuntime
+#include "pycore_pythread.h"      // PyThread_start_joinable_thread()
 #include "pycore_import.h"        // _PyImport_FrozenBootstrap
 #include <inttypes.h>
 #include <stdio.h>
@@ -2022,6 +2023,22 @@ static int test_init_main_interpreter_settings(void)
     return 0;
 }
 
+static void do_init(void *unused)
+{
+    _testembed_Py_Initialize();
+    Py_Finalize();
+}
+
+static int test_init_in_background_thread(void)
+{
+    PyThread_handle_t handle;
+    PyThread_ident_t ident;
+    if (PyThread_start_joinable_thread(&do_init, NULL, &ident, &handle) < 0) {
+        return -1;
+    }
+    return PyThread_join_thread(handle);
+}
+
 
 #ifndef MS_WINDOWS
 #include "test_frozenmain.h"      // M_test_frozenmain
@@ -2211,6 +2228,7 @@ static struct TestCase TestCases[] = {
     {"test_get_argc_argv", test_get_argc_argv},
     {"test_init_use_frozen_modules", test_init_use_frozen_modules},
     {"test_init_main_interpreter_settings", test_init_main_interpreter_settings},
+    {"test_init_in_background_thread", test_init_in_background_thread},
 
     // Audit
     {"test_open_code_hook", test_open_code_hook},