/* Python's malloc wrappers (see pymem.h) */

#include "Python.h"

#include "pycore_interp.h" // _PyInterpreterState_HasFeature

#include "pycore_mmap.h" // _PyAnnotateMemoryMap()

#include "pycore_object.h" // _PyDebugAllocatorStats() definition

#include "pycore_obmalloc.h"

#include "pycore_obmalloc_init.h"

#include "pycore_pyerrors.h" // _Py_FatalErrorFormat()

#include "pycore_pymem.h"

#include "pycore_pystate.h" // _PyInterpreterState_GET

#include "pycore_stats.h" // OBJECT_STAT_INC_COND()

#include <stdlib.h> // malloc()

#include <stdbool.h>

#ifdef WITH_MIMALLOC

// Forward declarations of functions used in our mimalloc modifications

static void _PyMem_mi_page_clear_qsbr(mi_page_t *page);

static bool _PyMem_mi_page_is_safe_to_free(mi_page_t *page);

static bool _PyMem_mi_page_maybe_free(mi_page_t *page, mi_page_queue_t *pq, bool force);

static void _PyMem_mi_page_reclaimed(mi_page_t *page);

static void _PyMem_mi_heap_collect_qsbr(mi_heap_t *heap);

# include "pycore_mimalloc.h"

# include "mimalloc/static.c"

# include "mimalloc/internal.h" // for stats

#endif

#if defined(Py_GIL_DISABLED) && !defined(WITH_MIMALLOC)

# error "Py_GIL_DISABLED requires WITH_MIMALLOC"

#endif

#undef uint

#define uint pymem_uint

/* Defined in tracemalloc.c */

extern void _PyMem_DumpTraceback(int fd, const void *ptr);

static void _PyObject_DebugDumpAddress(const void *p);

static void _PyMem_DebugCheckAddress(const char *func, char api_id, const void *p);

static void set_up_debug_hooks_domain_unlocked(PyMemAllocatorDomain domain);

static void set_up_debug_hooks_unlocked(void);

static void get_allocator_unlocked(PyMemAllocatorDomain, PyMemAllocatorEx *);

static void set_allocator_unlocked(PyMemAllocatorDomain, PyMemAllocatorEx *);

/***************************************/

/* low-level allocator implementations */

/***************************************/

/* the default raw allocator (wraps malloc) */

void *

_PyMem_RawMalloc(void *Py_UNUSED(ctx), size_t size)

{

/* PyMem_RawMalloc(0) means malloc(1). Some systems would return NULL

for malloc(0), which would be treated as an error. Some platforms would

return a pointer with no memory behind it, which would break pymalloc.

To solve these problems, allocate an extra byte. */

if (size == 0)

size = 1;

return malloc(size);

}

void *

_PyMem_RawCalloc(void *Py_UNUSED(ctx), size_t nelem, size_t elsize)

{

/* PyMem_RawCalloc(0, 0) means calloc(1, 1). Some systems would return NULL

for calloc(0, 0), which would be treated as an error. Some platforms

would return a pointer with no memory behind it, which would break

pymalloc. To solve these problems, allocate an extra byte. */

if (nelem == 0 || elsize == 0) {

nelem = 1;

elsize = 1;

}

return calloc(nelem, elsize);

}

void *

_PyMem_RawRealloc(void *Py_UNUSED(ctx), void *ptr, size_t size)

{

if (size == 0)

size = 1;

return realloc(ptr, size);

}

void

_PyMem_RawFree(void *Py_UNUSED(ctx), void *ptr)

{

free(ptr);

}

#ifdef WITH_MIMALLOC

static void

_PyMem_mi_page_clear_qsbr(mi_page_t *page)

{

#ifdef Py_GIL_DISABLED

// Clear the QSBR goal and remove the page from the QSBR linked list.

page->qsbr_goal = 0;

if (page->qsbr_node.next != NULL) {

llist_remove(&page->qsbr_node);

}

#endif

}

// Check if an empty, newly reclaimed page is safe to free now.

static bool

_PyMem_mi_page_is_safe_to_free(mi_page_t *page)

{

assert(mi_page_all_free(page));

#ifdef Py_GIL_DISABLED

assert(page->qsbr_node.next == NULL);

if (page->use_qsbr && page->qsbr_goal != 0) {

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

if (tstate == NULL) {

return false;

}

return _Py_qbsr_goal_reached(tstate->qsbr, page->qsbr_goal);

}

#endif

return true;

}

#ifdef Py_GIL_DISABLED

// If we are deferring collection of more than this amount of memory for

// mimalloc pages, advance the write sequence. Advancing allows these

// pages to be re-used in a different thread or for a different size class.

#define QSBR_PAGE_MEM_LIMIT 4096*20

// Return true if the global write sequence should be advanced for a mimalloc

// page that is deferred from collection.

static bool

should_advance_qsbr_for_page(struct _qsbr_thread_state *qsbr, mi_page_t *page)

{

size_t bsize = mi_page_block_size(page);

size_t page_size = page->capacity*bsize;

if (page_size > QSBR_PAGE_MEM_LIMIT) {

qsbr->deferred_page_memory = 0;

return true;

}

qsbr->deferred_page_memory += page_size;

if (qsbr->deferred_page_memory > QSBR_PAGE_MEM_LIMIT) {

qsbr->deferred_page_memory = 0;

return true;

}

return false;

}

#endif

static bool

_PyMem_mi_page_maybe_free(mi_page_t *page, mi_page_queue_t *pq, bool force)

{

#ifdef Py_GIL_DISABLED

assert(mi_page_all_free(page));

if (page->use_qsbr) {

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)PyThreadState_GET();

if (page->qsbr_goal != 0 && _Py_qbsr_goal_reached(tstate->qsbr, page->qsbr_goal)) {

_PyMem_mi_page_clear_qsbr(page);

_mi_page_free(page, pq, force);

return true;

}

_PyMem_mi_page_clear_qsbr(page);

page->retire_expire = 0;

if (should_advance_qsbr_for_page(tstate->qsbr, page)) {

page->qsbr_goal = _Py_qsbr_advance(tstate->qsbr->shared);

}

else {

page->qsbr_goal = _Py_qsbr_shared_next(tstate->qsbr->shared);

}

llist_insert_tail(&tstate->mimalloc.page_list, &page->qsbr_node);

return false;

}

#endif

_mi_page_free(page, pq, force);

return true;

}

static void

_PyMem_mi_page_reclaimed(mi_page_t *page)

{

#ifdef Py_GIL_DISABLED

assert(page->qsbr_node.next == NULL);

if (page->qsbr_goal != 0) {

if (mi_page_all_free(page)) {

assert(page->qsbr_node.next == NULL);

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)PyThreadState_GET();

page->retire_expire = 0;

llist_insert_tail(&tstate->mimalloc.page_list, &page->qsbr_node);

}

else {

page->qsbr_goal = 0;

}

}

#endif

}

static void

_PyMem_mi_heap_collect_qsbr(mi_heap_t *heap)

{

#ifdef Py_GIL_DISABLED

if (!heap->page_use_qsbr) {

return;

}

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

struct llist_node *head = &tstate->mimalloc.page_list;

if (llist_empty(head)) {

return;

}

struct llist_node *node;

llist_for_each_safe(node, head) {

mi_page_t *page = llist_data(node, mi_page_t, qsbr_node);

if (!mi_page_all_free(page)) {

// We allocated from this page some point after the delayed free

_PyMem_mi_page_clear_qsbr(page);

continue;

}

if (!_Py_qsbr_poll(tstate->qsbr, page->qsbr_goal)) {

return;

}

_PyMem_mi_page_clear_qsbr(page);

_mi_page_free(page, mi_page_queue_of(page), false);

}

#endif

}

void *

_PyMem_MiMalloc(void *ctx, size_t size)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM];

return mi_heap_malloc(heap, size);

#else

return mi_malloc(size);

#endif

}

void *

_PyMem_MiCalloc(void *ctx, size_t nelem, size_t elsize)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM];

return mi_heap_calloc(heap, nelem, elsize);

#else

return mi_calloc(nelem, elsize);

#endif

}

void *

_PyMem_MiRealloc(void *ctx, void *ptr, size_t size)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

mi_heap_t *heap = &tstate->mimalloc.heaps[_Py_MIMALLOC_HEAP_MEM];

return mi_heap_realloc(heap, ptr, size);

#else

return mi_realloc(ptr, size);

#endif

}

void

_PyMem_MiFree(void *ctx, void *ptr)

{

mi_free(ptr);

}

void *

_PyObject_MiMalloc(void *ctx, size_t nbytes)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

mi_heap_t *heap = tstate->mimalloc.current_object_heap;

return mi_heap_malloc(heap, nbytes);

#else

return mi_malloc(nbytes);

#endif

}

void *

_PyObject_MiCalloc(void *ctx, size_t nelem, size_t elsize)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

mi_heap_t *heap = tstate->mimalloc.current_object_heap;

return mi_heap_calloc(heap, nelem, elsize);

#else

return mi_calloc(nelem, elsize);

#endif

}

void *

_PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes)

{

#ifdef Py_GIL_DISABLED

_PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();

// Implement our own realloc logic so that we can copy PyObject header

// in a thread-safe way.

size_t size = mi_usable_size(ptr);

if (nbytes <= size && nbytes >= (size / 2) && nbytes > 0) {

return ptr;

}

mi_heap_t *heap = tstate->mimalloc.current_object_heap;

void* newp = mi_heap_malloc(heap, nbytes);

if (newp == NULL) {

return NULL;

}

// Free threaded Python allows access from other threads to the PyObject reference count

// fields for a period of time after the object is freed (see InternalDocs/qsbr.md).

// These fields are typically initialized by PyObject_Init() using relaxed

// atomic stores. We need to copy these fields in a thread-safe way here.

// We use the "debug_offset" to determine how many bytes to copy -- it

// includes the PyObject header and plus any extra pre-headers.

size_t offset = heap->debug_offset;

assert(offset % sizeof(void*) == 0);

size_t copy_size = (size < nbytes ? size : nbytes);

if (copy_size >= offset) {

for (size_t i = 0; i != offset; i += sizeof(void*)) {

// Use memcpy to avoid strict-aliasing issues. However, we probably

// still have unavoidable strict-aliasing issues with

// _Py_atomic_store_ptr_relaxed here.

void *word;

memcpy(&word, (char*)ptr + i, sizeof(void*));

_Py_atomic_store_ptr_relaxed((void**)((char*)newp + i), word);

}

_mi_memcpy((char*)newp + offset, (char*)ptr + offset, copy_size - offset);

}

else {

_mi_memcpy(newp, ptr, copy_size);

}

mi_free(ptr);

return newp;

#else

return mi_realloc(ptr, nbytes);

#endif

}

void

_PyObject_MiFree(void *ctx, void *ptr)

{

mi_free(ptr);

}

#endif // WITH_MIMALLOC

#define MALLOC_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}

#ifdef WITH_MIMALLOC

# define MIMALLOC_ALLOC {NULL, _PyMem_MiMalloc, _PyMem_MiCalloc, _PyMem_MiRealloc, _PyMem_MiFree}

# define MIMALLOC_OBJALLOC {NULL, _PyObject_MiMalloc, _PyObject_MiCalloc, _PyObject_MiRealloc, _PyObject_MiFree}

#endif

/* the pymalloc allocator */

// The actual implementation is further down.

#if defined(WITH_PYMALLOC)

void* _PyObject_Malloc(void *ctx, size_t size);

void* _PyObject_Calloc(void *ctx, size_t nelem, size_t elsize);

void _PyObject_Free(void *ctx, void *p);

void* _PyObject_Realloc(void *ctx, void *ptr, size_t size);

# define PYMALLOC_ALLOC {NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free}

#endif // WITH_PYMALLOC

#if defined(Py_GIL_DISABLED)

// Py_GIL_DISABLED requires using mimalloc for "mem" and "obj" domains.

# define PYRAW_ALLOC MALLOC_ALLOC

# define PYMEM_ALLOC MIMALLOC_ALLOC

# define PYOBJ_ALLOC MIMALLOC_OBJALLOC

#elif defined(WITH_PYMALLOC)

# define PYRAW_ALLOC MALLOC_ALLOC

# define PYMEM_ALLOC PYMALLOC_ALLOC

# define PYOBJ_ALLOC PYMALLOC_ALLOC

#else

# define PYRAW_ALLOC MALLOC_ALLOC

# define PYMEM_ALLOC MALLOC_ALLOC

# define PYOBJ_ALLOC MALLOC_ALLOC

#endif

/* the default debug allocators */

// The actual implementation is further down.

void* _PyMem_DebugRawMalloc(void *ctx, size_t size);

void* _PyMem_DebugRawCalloc(void *ctx, size_t nelem, size_t elsize);

void* _PyMem_DebugRawRealloc(void *ctx, void *ptr, size_t size);

void _PyMem_DebugRawFree(void *ctx, void *ptr);

void* _PyMem_DebugMalloc(void *ctx, size_t size);

void* _PyMem_DebugCalloc(void *ctx, size_t nelem, size_t elsize);

void* _PyMem_DebugRealloc(void *ctx, void *ptr, size_t size);

void _PyMem_DebugFree(void *ctx, void *p);

#define PYDBGRAW_ALLOC \

{&_PyRuntime.allocators.debug.raw, _PyMem_DebugRawMalloc, _PyMem_DebugRawCalloc, _PyMem_DebugRawRealloc, _PyMem_DebugRawFree}

#define PYDBGMEM_ALLOC \

{&_PyRuntime.allocators.debug.mem, _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree}

#define PYDBGOBJ_ALLOC \

{&_PyRuntime.allocators.debug.obj, _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree}

/* default raw allocator (not swappable) */

void *

_PyMem_DefaultRawMalloc(size_t size)

{

#ifdef Py_DEBUG

return _PyMem_DebugRawMalloc(&_PyRuntime.allocators.debug.raw, size);

#else

return _PyMem_RawMalloc(NULL, size);

#endif

}

void *

_PyMem_DefaultRawCalloc(size_t nelem, size_t elsize)

{

#ifdef Py_DEBUG

return _PyMem_DebugRawCalloc(&_PyRuntime.allocators.debug.raw, nelem, elsize);

#else

return _PyMem_RawCalloc(NULL, nelem, elsize);

#endif

}

void *

_PyMem_DefaultRawRealloc(void *ptr, size_t size)

{

#ifdef Py_DEBUG

return _PyMem_DebugRawRealloc(&_PyRuntime.allocators.debug.raw, ptr, size);

#else

return _PyMem_RawRealloc(NULL, ptr, size);

#endif

}

void

_PyMem_DefaultRawFree(void *ptr)

{

#ifdef Py_DEBUG

_PyMem_DebugRawFree(&_PyRuntime.allocators.debug.raw, ptr);

#else

_PyMem_RawFree(NULL, ptr);

#endif

}

wchar_t*

_PyMem_DefaultRawWcsdup(const wchar_t *str)

{

assert(str != NULL);

size_t len = wcslen(str);

if (len > (size_t)PY_SSIZE_T_MAX / sizeof(wchar_t) - 1) {

return NULL;

}

size_t size = (len + 1) * sizeof(wchar_t);

wchar_t *str2 = _PyMem_DefaultRawMalloc(size);

if (str2 == NULL) {

return NULL;

}

memcpy(str2, str, size);

return str2;

}

/* the low-level virtual memory allocator */

#ifdef WITH_PYMALLOC

# ifdef MS_WINDOWS

# include <windows.h>

# elif defined(HAVE_MMAP)

# include <sys/mman.h>

# ifdef MAP_ANONYMOUS

# define ARENAS_USE_MMAP

# endif

# endif

#endif

void *

_PyMem_ArenaAlloc(void *Py_UNUSED(ctx), size_t size)

{

#ifdef MS_WINDOWS

return VirtualAlloc(NULL, size,

MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

#elif defined(ARENAS_USE_MMAP)

void *ptr;

ptr = mmap(NULL, size, PROT_READ|PROT_WRITE,

MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);

if (ptr == MAP_FAILED)

return NULL;

assert(ptr != NULL);

(void)_PyAnnotateMemoryMap(ptr, size, "cpython:pymalloc");

return ptr;

#else

return malloc(size);

#endif

}

void

_PyMem_ArenaFree(void *Py_UNUSED(ctx), void *ptr,

#if defined(ARENAS_USE_MMAP)

size_t size

#else

size_t Py_UNUSED(size)

#endif

)

{

#ifdef MS_WINDOWS

/* Unlike free(), VirtualFree() does not special-case NULL to noop. */

if (ptr == NULL) {

return;

}

VirtualFree(ptr, 0, MEM_RELEASE);

#elif defined(ARENAS_USE_MMAP)

/* Unlike free(), munmap() does not special-case NULL to noop. */

if (ptr == NULL) {

return;

}

munmap(ptr, size);

#else

free(ptr);

#endif

}

/*******************************************/

/* end low-level allocator implementations */

/*******************************************/

#define ALLOCATORS_MUTEX (_PyRuntime.allocators.mutex)

#define _PyMem_Raw (_PyRuntime.allocators.standard.raw)

#define _PyMem (_PyRuntime.allocators.standard.mem)

#define _PyObject (_PyRuntime.allocators.standard.obj)

#define _PyMem_Debug (_PyRuntime.allocators.debug)

#define _PyObject_Arena (_PyRuntime.allocators.obj_arena)

/***************************/

/* managing the allocators */

/***************************/

static int

set_default_allocator_unlocked(PyMemAllocatorDomain domain, int debug,

PyMemAllocatorEx *old_alloc)

{

if (old_alloc != NULL) {

get_allocator_unlocked(domain, old_alloc);

}

PyMemAllocatorEx new_alloc;

switch(domain)

{

case PYMEM_DOMAIN_RAW:

new_alloc = (PyMemAllocatorEx)PYRAW_ALLOC;

break;

case PYMEM_DOMAIN_MEM:

new_alloc = (PyMemAllocatorEx)PYMEM_ALLOC;

break;

case PYMEM_DOMAIN_OBJ:

new_alloc = (PyMemAllocatorEx)PYOBJ_ALLOC;

break;

default:

/* unknown domain */

return -1;

}

set_allocator_unlocked(domain, &new_alloc);

if (debug) {

set_up_debug_hooks_domain_unlocked(domain);

}

return 0;

}

#ifdef Py_DEBUG

static const int pydebug = 1;

#else

static const int pydebug = 0;

#endif

int

_PyMem_GetAllocatorName(const char *name, PyMemAllocatorName *allocator)

{

if (name == NULL || *name == '\0') {

/* PYTHONMALLOC is empty or is not set or ignored (-E/-I command line

nameions): use default memory allocators */

*allocator = PYMEM_ALLOCATOR_DEFAULT;

}

else if (strcmp(name, "default") == 0) {

*allocator = PYMEM_ALLOCATOR_DEFAULT;

}

else if (strcmp(name, "debug") == 0) {

*allocator = PYMEM_ALLOCATOR_DEBUG;

}

#if defined(WITH_PYMALLOC) && !defined(Py_GIL_DISABLED)

else if (strcmp(name, "pymalloc") == 0) {

*allocator = PYMEM_ALLOCATOR_PYMALLOC;

}

else if (strcmp(name, "pymalloc_debug") == 0) {

*allocator = PYMEM_ALLOCATOR_PYMALLOC_DEBUG;

}

#endif

#ifdef WITH_MIMALLOC

else if (strcmp(name, "mimalloc") == 0) {

*allocator = PYMEM_ALLOCATOR_MIMALLOC;

}

else if (strcmp(name, "mimalloc_debug") == 0) {

*allocator = PYMEM_ALLOCATOR_MIMALLOC_DEBUG;

}

#endif

#ifndef Py_GIL_DISABLED

else if (strcmp(name, "malloc") == 0) {

*allocator = PYMEM_ALLOCATOR_MALLOC;

}

else if (strcmp(name, "malloc_debug") == 0) {

*allocator = PYMEM_ALLOCATOR_MALLOC_DEBUG;

}

#endif

else {

/* unknown allocator */

return -1;

}

return 0;

}

static int

set_up_allocators_unlocked(PyMemAllocatorName allocator)

{

switch (allocator) {

case PYMEM_ALLOCATOR_NOT_SET:

/* do nothing */

break;

case PYMEM_ALLOCATOR_DEFAULT:

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_RAW, pydebug, NULL);

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_MEM, pydebug, NULL);

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_OBJ, pydebug, NULL);

_PyRuntime.allocators.is_debug_enabled = pydebug;

break;

case PYMEM_ALLOCATOR_DEBUG:

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_RAW, 1, NULL);

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_MEM, 1, NULL);

(void)set_default_allocator_unlocked(PYMEM_DOMAIN_OBJ, 1, NULL);

_PyRuntime.allocators.is_debug_enabled = 1;

break;

#ifdef WITH_PYMALLOC

case PYMEM_ALLOCATOR_PYMALLOC:

case PYMEM_ALLOCATOR_PYMALLOC_DEBUG:

{

PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_RAW, &malloc_alloc);

PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_MEM, &pymalloc);

set_allocator_unlocked(PYMEM_DOMAIN_OBJ, &pymalloc);

int is_debug = (allocator == PYMEM_ALLOCATOR_PYMALLOC_DEBUG);

_PyRuntime.allocators.is_debug_enabled = is_debug;

if (is_debug) {

set_up_debug_hooks_unlocked();

}

break;

}

#endif

#ifdef WITH_MIMALLOC

case PYMEM_ALLOCATOR_MIMALLOC:

case PYMEM_ALLOCATOR_MIMALLOC_DEBUG:

{

PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_RAW, &malloc_alloc);

PyMemAllocatorEx pymalloc = MIMALLOC_ALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_MEM, &pymalloc);

PyMemAllocatorEx objmalloc = MIMALLOC_OBJALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_OBJ, &objmalloc);

int is_debug = (allocator == PYMEM_ALLOCATOR_MIMALLOC_DEBUG);

_PyRuntime.allocators.is_debug_enabled = is_debug;

if (is_debug) {

set_up_debug_hooks_unlocked();

}

break;

}

#endif

case PYMEM_ALLOCATOR_MALLOC:

case PYMEM_ALLOCATOR_MALLOC_DEBUG:

{

PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;

set_allocator_unlocked(PYMEM_DOMAIN_RAW, &malloc_alloc);

set_allocator_unlocked(PYMEM_DOMAIN_MEM, &malloc_alloc);

set_allocator_unlocked(PYMEM_DOMAIN_OBJ, &malloc_alloc);

int is_debug = (allocator == PYMEM_ALLOCATOR_MALLOC_DEBUG);

_PyRuntime.allocators.is_debug_enabled = is_debug;

if (is_debug) {

set_up_debug_hooks_unlocked();

}

break;

}

default:

/* unknown allocator */

return -1;

}

return 0;

}

int

_PyMem_SetupAllocators(PyMemAllocatorName allocator)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

int res = set_up_allocators_unlocked(allocator);

PyMutex_Unlock(&ALLOCATORS_MUTEX);

return res;

}

static int

pymemallocator_eq(PyMemAllocatorEx *a, PyMemAllocatorEx *b)

{

return (memcmp(a, b, sizeof(PyMemAllocatorEx)) == 0);

}

static const char*

get_current_allocator_name_unlocked(void)

{

PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;

#ifdef WITH_PYMALLOC

PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;

#endif

#ifdef WITH_MIMALLOC

PyMemAllocatorEx mimalloc = MIMALLOC_ALLOC;

PyMemAllocatorEx mimalloc_obj = MIMALLOC_OBJALLOC;

#endif

if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&

pymemallocator_eq(&_PyMem, &malloc_alloc) &&

pymemallocator_eq(&_PyObject, &malloc_alloc))

{

return "malloc";

}

#ifdef WITH_PYMALLOC

if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&

pymemallocator_eq(&_PyMem, &pymalloc) &&

pymemallocator_eq(&_PyObject, &pymalloc))

{

return "pymalloc";

}

#endif

#ifdef WITH_MIMALLOC

if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&

pymemallocator_eq(&_PyMem, &mimalloc) &&

pymemallocator_eq(&_PyObject, &mimalloc_obj))

{

return "mimalloc";

}

#endif

PyMemAllocatorEx dbg_raw = PYDBGRAW_ALLOC;

PyMemAllocatorEx dbg_mem = PYDBGMEM_ALLOC;

PyMemAllocatorEx dbg_obj = PYDBGOBJ_ALLOC;

if (pymemallocator_eq(&_PyMem_Raw, &dbg_raw) &&

pymemallocator_eq(&_PyMem, &dbg_mem) &&

pymemallocator_eq(&_PyObject, &dbg_obj))

{

/* Debug hooks installed */

if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&

pymemallocator_eq(&_PyMem_Debug.mem.alloc, &malloc_alloc) &&

pymemallocator_eq(&_PyMem_Debug.obj.alloc, &malloc_alloc))

{

return "malloc_debug";

}

#ifdef WITH_PYMALLOC

if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&

pymemallocator_eq(&_PyMem_Debug.mem.alloc, &pymalloc) &&

pymemallocator_eq(&_PyMem_Debug.obj.alloc, &pymalloc))

{

return "pymalloc_debug";

}

#endif

#ifdef WITH_MIMALLOC

if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&

pymemallocator_eq(&_PyMem_Debug.mem.alloc, &mimalloc) &&

pymemallocator_eq(&_PyMem_Debug.obj.alloc, &mimalloc_obj))

{

return "mimalloc_debug";

}

#endif

}

return NULL;

}

const char*

_PyMem_GetCurrentAllocatorName(void)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

const char *name = get_current_allocator_name_unlocked();

PyMutex_Unlock(&ALLOCATORS_MUTEX);

return name;

}

int

_PyMem_DebugEnabled(void)

{

return _PyRuntime.allocators.is_debug_enabled;

}

#ifdef WITH_PYMALLOC

static int

_PyMem_PymallocEnabled(void)

{

if (_PyMem_DebugEnabled()) {

return (_PyMem_Debug.obj.alloc.malloc == _PyObject_Malloc);

}

else {

return (_PyObject.malloc == _PyObject_Malloc);

}

}

#ifdef WITH_MIMALLOC

static int

_PyMem_MimallocEnabled(void)

{

#ifdef Py_GIL_DISABLED

return 1;

#else

if (_PyMem_DebugEnabled()) {

return (_PyMem_Debug.obj.alloc.malloc == _PyObject_MiMalloc);

}

else {

return (_PyObject.malloc == _PyObject_MiMalloc);

}

#endif

}

#endif // WITH_MIMALLOC

#endif // WITH_PYMALLOC

static void

set_up_debug_hooks_domain_unlocked(PyMemAllocatorDomain domain)

{

PyMemAllocatorEx alloc;

if (domain == PYMEM_DOMAIN_RAW) {

if (_PyMem_Raw.malloc == _PyMem_DebugRawMalloc) {

return;

}

get_allocator_unlocked(domain, &_PyMem_Debug.raw.alloc);

alloc.ctx = &_PyMem_Debug.raw;

alloc.malloc = _PyMem_DebugRawMalloc;

alloc.calloc = _PyMem_DebugRawCalloc;

alloc.realloc = _PyMem_DebugRawRealloc;

alloc.free = _PyMem_DebugRawFree;

set_allocator_unlocked(domain, &alloc);

}

else if (domain == PYMEM_DOMAIN_MEM) {

if (_PyMem.malloc == _PyMem_DebugMalloc) {

return;

}

get_allocator_unlocked(domain, &_PyMem_Debug.mem.alloc);

alloc.ctx = &_PyMem_Debug.mem;

alloc.malloc = _PyMem_DebugMalloc;

alloc.calloc = _PyMem_DebugCalloc;

alloc.realloc = _PyMem_DebugRealloc;

alloc.free = _PyMem_DebugFree;

set_allocator_unlocked(domain, &alloc);

}

else if (domain == PYMEM_DOMAIN_OBJ) {

if (_PyObject.malloc == _PyMem_DebugMalloc) {

return;

}

get_allocator_unlocked(domain, &_PyMem_Debug.obj.alloc);

alloc.ctx = &_PyMem_Debug.obj;

alloc.malloc = _PyMem_DebugMalloc;

alloc.calloc = _PyMem_DebugCalloc;

alloc.realloc = _PyMem_DebugRealloc;

alloc.free = _PyMem_DebugFree;

set_allocator_unlocked(domain, &alloc);

}

}

static void

set_up_debug_hooks_unlocked(void)

{

set_up_debug_hooks_domain_unlocked(PYMEM_DOMAIN_RAW);

set_up_debug_hooks_domain_unlocked(PYMEM_DOMAIN_MEM);

set_up_debug_hooks_domain_unlocked(PYMEM_DOMAIN_OBJ);

_PyRuntime.allocators.is_debug_enabled = 1;

}

void

PyMem_SetupDebugHooks(void)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

set_up_debug_hooks_unlocked();

PyMutex_Unlock(&ALLOCATORS_MUTEX);

}

static void

get_allocator_unlocked(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)

{

switch(domain)

{

case PYMEM_DOMAIN_RAW: *allocator = _PyMem_Raw; break;

case PYMEM_DOMAIN_MEM: *allocator = _PyMem; break;

case PYMEM_DOMAIN_OBJ: *allocator = _PyObject; break;

default:

/* unknown domain: set all attributes to NULL */

allocator->ctx = NULL;

allocator->malloc = NULL;

allocator->calloc = NULL;

allocator->realloc = NULL;

allocator->free = NULL;

}

}

static void

set_allocator_unlocked(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)

{

switch(domain)

{

case PYMEM_DOMAIN_RAW: _PyMem_Raw = *allocator; break;

case PYMEM_DOMAIN_MEM: _PyMem = *allocator; break;

case PYMEM_DOMAIN_OBJ: _PyObject = *allocator; break;

/* ignore unknown domain */

}

}

void

PyMem_GetAllocator(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

get_allocator_unlocked(domain, allocator);

PyMutex_Unlock(&ALLOCATORS_MUTEX);

}

void

PyMem_SetAllocator(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

set_allocator_unlocked(domain, allocator);

PyMutex_Unlock(&ALLOCATORS_MUTEX);

}

void

PyObject_GetArenaAllocator(PyObjectArenaAllocator *allocator)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

*allocator = _PyObject_Arena;

PyMutex_Unlock(&ALLOCATORS_MUTEX);

}

void

PyObject_SetArenaAllocator(PyObjectArenaAllocator *allocator)

{

PyMutex_Lock(&ALLOCATORS_MUTEX);

_PyObject_Arena = *allocator;

PyMutex_Unlock(&ALLOCATORS_MUTEX);

}

/* Note that there is a possible, but very unlikely, race in any place

* below where we call one of the allocator functions. We access two

* fields in each case: "malloc", etc. and "ctx".

*

* It is unlikely that the allocator will be changed while one of those

* calls is happening, much less in that very narrow window.

* Furthermore, the likelihood of a race is drastically reduced by the

* fact that the allocator may not be changed after runtime init