Skip to content

Commit

Permalink
pythongh-123022: Fix crash with Py_Initialize in background thread (p…
Browse files Browse the repository at this point in the history
…ython#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.
  • Loading branch information
colesbury authored Aug 17, 2024
1 parent 40632b1 commit d061ffe
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 5 deletions.
9 changes: 8 additions & 1 deletion Lib/test/test_embed.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -1802,6 +1802,13 @@ def test_init_main_interpreter_settings(self):

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):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix crash in free-threaded build when calling :c:func:`Py_Initialize` from
a non-main thread.
16 changes: 12 additions & 4 deletions Objects/mimalloc/os.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
18 changes: 18 additions & 0 deletions Programs/_testembed.c
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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},
Expand Down

0 comments on commit d061ffe

Please sign in to comment.