Modules/_testcapi/mem.c

#include "parts.h"

#include <stddef.h>

#if defined(__APPLE__)
#  include <TargetConditionals.h>
   // Older macOS SDKs do not define TARGET_OS_OSX
#  if !defined(TARGET_OS_OSX)
#    define TARGET_OS_OSX 1
#  endif
#  if TARGET_OS_OSX
#    include <errno.h>              // errno, ESRCH
#    include <libproc.h>            // proc_pidinfo(), PROC_PIDTASKINFO
#    include <sys/proc_info.h>      // struct proc_taskinfo
#  endif
#endif

#ifdef __FreeBSD__
#  include <fcntl.h>              // O_RDONLY
#  include <kvm.h>                // kvm_openfiles()
#  include <limits.h>             // _POSIX2_LINE_MAX
#  include <sys/sysctl.h>         // KERN_PROC_PID
#  include <sys/user.h>           // kinfo_proc definition
#  include <unistd.h>             // sysconf()
#endif


typedef struct {
    PyMemAllocatorEx alloc;

    size_t malloc_size;
    size_t calloc_nelem;
    size_t calloc_elsize;
    void *realloc_ptr;
    size_t realloc_new_size;
    void *free_ptr;
    void *ctx;
} alloc_hook_t;

static void *
hook_malloc(void *ctx, size_t size)
{
    alloc_hook_t *hook = (alloc_hook_t *)ctx;
    hook->ctx = ctx;
    hook->malloc_size = size;
    return hook->alloc.malloc(hook->alloc.ctx, size);
}

static void *
hook_calloc(void *ctx, size_t nelem, size_t elsize)
{
    alloc_hook_t *hook = (alloc_hook_t *)ctx;
    hook->ctx = ctx;
    hook->calloc_nelem = nelem;
    hook->calloc_elsize = elsize;
    return hook->alloc.calloc(hook->alloc.ctx, nelem, elsize);
}

static void *
hook_realloc(void *ctx, void *ptr, size_t new_size)
{
    alloc_hook_t *hook = (alloc_hook_t *)ctx;
    hook->ctx = ctx;
    hook->realloc_ptr = ptr;
    hook->realloc_new_size = new_size;
    return hook->alloc.realloc(hook->alloc.ctx, ptr, new_size);
}

static void
hook_free(void *ctx, void *ptr)
{
    alloc_hook_t *hook = (alloc_hook_t *)ctx;
    hook->ctx = ctx;
    hook->free_ptr = ptr;
    hook->alloc.free(hook->alloc.ctx, ptr);
}

/* Most part of the following code is inherited from the pyfailmalloc project
 * written by Victor Stinner. */
static struct {
    int installed;
    PyMemAllocatorEx raw;
    PyMemAllocatorEx mem;
    PyMemAllocatorEx obj;
} FmHook;

static struct {
    int start;
    int stop;
    Py_ssize_t count;
} FmData;

static int
fm_nomemory(void)
{
    FmData.count++;
    if (FmData.count > FmData.start &&
        (FmData.stop <= 0 || FmData.count <= FmData.stop))
    {
        return 1;
    }
    return 0;
}

static void *
hook_fmalloc(void *ctx, size_t size)
{
    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
    if (fm_nomemory()) {
        return NULL;
    }
    return alloc->malloc(alloc->ctx, size);
}

static void *
hook_fcalloc(void *ctx, size_t nelem, size_t elsize)
{
    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
    if (fm_nomemory()) {
        return NULL;
    }
    return alloc->calloc(alloc->ctx, nelem, elsize);
}

static void *
hook_frealloc(void *ctx, void *ptr, size_t new_size)
{
    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
    if (fm_nomemory()) {
        return NULL;
    }
    return alloc->realloc(alloc->ctx, ptr, new_size);
}

static void
hook_ffree(void *ctx, void *ptr)
{
    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
    alloc->free(alloc->ctx, ptr);
}

static void
fm_setup_hooks(void)
{
    if (FmHook.installed) {
        return;
    }
    FmHook.installed = 1;

    PyMemAllocatorEx alloc;
    alloc.malloc = hook_fmalloc;
    alloc.calloc = hook_fcalloc;
    alloc.realloc = hook_frealloc;
    alloc.free = hook_ffree;
    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
    PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
    PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);

    alloc.ctx = &FmHook.raw;
    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);

    alloc.ctx = &FmHook.mem;
    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);

    alloc.ctx = &FmHook.obj;
    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
}

static void
fm_remove_hooks(void)
{
    if (FmHook.installed) {
        FmHook.installed = 0;
        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
        PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
        PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);
    }
}

static void
fm_set_nomemory(int start, int stop)
{
    /* Memory allocation fails after 'start' allocation requests, and until
     * 'stop' allocation requests except when 'stop' is negative or equal
     * to 0 (default) in which case allocation failures never stop. */
    FmData.start = start;
    FmData.stop = stop;
    FmData.count = 0;
    fm_setup_hooks();
}

static PyObject *
set_nomemory(PyObject *self, PyObject *args)
{
    /* Memory allocation fails after 'start' allocation requests, and until
     * 'stop' allocation requests except when 'stop' is negative or equal
     * to 0 (default) in which case allocation failures never stop. */
    int start, stop = 0;
    if (!PyArg_ParseTuple(args, "i|i", &start, &stop)) {
        return NULL;
    }
    fm_set_nomemory(start, stop);
    Py_RETURN_NONE;
}

static PyObject *
remove_mem_hooks(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    fm_remove_hooks();
    Py_RETURN_NONE;
}

static PyObject *
call_with_nomemory(PyObject *self, PyObject *args)
{
    /* Call func(*args) with memory allocation failing as in set_nomemory().
     * No bytecode is executed between arming the failure and the call. */
    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
    if (nargs < 3) {
        PyErr_SetString(PyExc_TypeError,
                        "call_with_nomemory() requires at least 3 arguments");
        return NULL;
    }
    int start = PyLong_AsInt(PyTuple_GET_ITEM(args, 0));
    if (start == -1 && PyErr_Occurred()) {
        return NULL;
    }
    int stop = PyLong_AsInt(PyTuple_GET_ITEM(args, 1));
    if (stop == -1 && PyErr_Occurred()) {
        return NULL;
    }
    PyObject *func = PyTuple_GET_ITEM(args, 2);
    /* PyObject_Call() with a prebuilt tuple does not allocate,
     * unlike PyObject_Vectorcall() for a callee without vectorcall. */
    PyObject *callargs = PyTuple_GetSlice(args, 3, nargs);
    if (callargs == NULL) {
        return NULL;
    }
    fm_set_nomemory(start, stop);
    PyObject *res = PyObject_Call(func, callargs, NULL);
    fm_remove_hooks();
    Py_DECREF(callargs);
    return res;
}

static PyObject *
test_setallocators(PyMemAllocatorDomain domain)
{
    PyObject *res = NULL;
    const char *error_msg;
    alloc_hook_t hook;

    memset(&hook, 0, sizeof(hook));

    PyMemAllocatorEx alloc;
    alloc.ctx = &hook;
    alloc.malloc = &hook_malloc;
    alloc.calloc = &hook_calloc;
    alloc.realloc = &hook_realloc;
    alloc.free = &hook_free;
    PyMem_GetAllocator(domain, &hook.alloc);
    PyMem_SetAllocator(domain, &alloc);

    /* malloc, realloc, free */
    size_t size = 42;
    hook.ctx = NULL;
    void *ptr;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            ptr = PyMem_RawMalloc(size);
            break;
        case PYMEM_DOMAIN_MEM:
            ptr = PyMem_Malloc(size);
            break;
        case PYMEM_DOMAIN_OBJ:
            ptr = PyObject_Malloc(size);
            break;
        default:
            ptr = NULL;
            break;
    }

#define CHECK_CTX(FUNC)                     \
    if (hook.ctx != &hook) {                \
        error_msg = FUNC " wrong context";  \
        goto fail;                          \
    }                                       \
    hook.ctx = NULL;  /* reset for next check */

    if (ptr == NULL) {
        error_msg = "malloc failed";
        goto fail;
    }
    CHECK_CTX("malloc");
    if (hook.malloc_size != size) {
        error_msg = "malloc invalid size";
        goto fail;
    }

    size_t size2 = 200;
    void *ptr2;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            ptr2 = PyMem_RawRealloc(ptr, size2);
            break;
        case PYMEM_DOMAIN_MEM:
            ptr2 = PyMem_Realloc(ptr, size2);
            break;
        case PYMEM_DOMAIN_OBJ:
            ptr2 = PyObject_Realloc(ptr, size2);
            break;
        default:
            ptr2 = NULL;
            break;
    }

    if (ptr2 == NULL) {
        error_msg = "realloc failed";
        goto fail;
    }
    CHECK_CTX("realloc");
    if (hook.realloc_ptr != ptr || hook.realloc_new_size != size2) {
        error_msg = "realloc invalid parameters";
        goto fail;
    }

    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            PyMem_RawFree(ptr2);
            break;
        case PYMEM_DOMAIN_MEM:
            PyMem_Free(ptr2);
            break;
        case PYMEM_DOMAIN_OBJ:
            PyObject_Free(ptr2);
            break;
    }

    CHECK_CTX("free");
    if (hook.free_ptr != ptr2) {
        error_msg = "free invalid pointer";
        goto fail;
    }

    /* calloc, free */
    size_t nelem = 2;
    size_t elsize = 5;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            ptr = PyMem_RawCalloc(nelem, elsize);
            break;
        case PYMEM_DOMAIN_MEM:
            ptr = PyMem_Calloc(nelem, elsize);
            break;
        case PYMEM_DOMAIN_OBJ:
            ptr = PyObject_Calloc(nelem, elsize);
            break;
        default:
            ptr = NULL;
            break;
    }

    if (ptr == NULL) {
        error_msg = "calloc failed";
        goto fail;
    }
    CHECK_CTX("calloc");
    if (hook.calloc_nelem != nelem || hook.calloc_elsize != elsize) {
        error_msg = "calloc invalid nelem or elsize";
        goto fail;
    }

    hook.free_ptr = NULL;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            PyMem_RawFree(ptr);
            break;
        case PYMEM_DOMAIN_MEM:
            PyMem_Free(ptr);
            break;
        case PYMEM_DOMAIN_OBJ:
            PyObject_Free(ptr);
            break;
    }

    CHECK_CTX("calloc free");
    if (hook.free_ptr != ptr) {
        error_msg = "calloc free invalid pointer";
        goto fail;
    }

    /* realloc(NULL, size) should behave like malloc(size) */
    size_t size3 = 100;
    void *ptr3;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            ptr3 = PyMem_RawRealloc(NULL, size3);
            break;
        case PYMEM_DOMAIN_MEM:
            ptr3 = PyMem_Realloc(NULL, size3);
            break;
        case PYMEM_DOMAIN_OBJ:
            ptr3 = PyObject_Realloc(NULL, size3);
            break;
        default:
            ptr3 = NULL;
            break;
    }

    CHECK_CTX("realloc(NULL, size)");
    if (ptr3 == NULL) {
        error_msg = "realloc(NULL, size) failed";
        goto fail;
    }
    if (hook.realloc_ptr != NULL || hook.realloc_new_size != size3) {
        error_msg = "realloc(NULL, size) invalid parameters";
        goto fail;
    }

    hook.free_ptr = NULL;
    switch(domain) {
        case PYMEM_DOMAIN_RAW:
            PyMem_RawFree(ptr3);
            break;
        case PYMEM_DOMAIN_MEM:
            PyMem_Free(ptr3);
            break;
        case PYMEM_DOMAIN_OBJ:
            PyObject_Free(ptr3);
            break;
    }

    CHECK_CTX("realloc(NULL, size) free");
    if (hook.free_ptr != ptr3) {
        error_msg = "unexpected pointer passed to free";
        goto fail;
    }

    res = Py_NewRef(Py_None);
    goto finally;

fail:
    PyErr_SetString(PyExc_RuntimeError, error_msg);

finally:
    PyMem_SetAllocator(domain, &hook.alloc);
    return res;

#undef CHECK_CTX
}

static PyObject *
test_pyobject_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    return test_setallocators(PYMEM_DOMAIN_OBJ);
}

static PyObject *
test_pyobject_new(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    PyObject *obj;
    PyTypeObject *type = &PyBaseObject_Type;
    PyTypeObject *var_type = &PyBytes_Type;

    // PyObject_New()
    obj = PyObject_New(PyObject, type);
    if (obj == NULL) {
        goto alloc_failed;
    }
    Py_DECREF(obj);

    // PyObject_NEW()
    obj = PyObject_NEW(PyObject, type);
    if (obj == NULL) {
        goto alloc_failed;
    }
    Py_DECREF(obj);

    // PyObject_NewVar()
    obj = PyObject_NewVar(PyObject, var_type, 3);
    if (obj == NULL) {
        goto alloc_failed;
    }
    memset(PyBytes_AS_STRING(obj), 0, 3 + 1);  // +1 for the null byte
    Py_DECREF(obj);

    // PyObject_NEW_VAR()
    obj = PyObject_NEW_VAR(PyObject, var_type, 3);
    if (obj == NULL) {
        goto alloc_failed;
    }
    memset(PyBytes_AS_STRING(obj), 0, 3 + 1);  // +1 for the null byte
    Py_DECREF(obj);

    Py_RETURN_NONE;

alloc_failed:
    PyErr_NoMemory();
    return NULL;
}

static PyObject *
test_pymem_alloc0(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    void *ptr;

    ptr = PyMem_RawMalloc(0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyMem_RawMalloc(0) returns NULL");
        return NULL;
    }
    PyMem_RawFree(ptr);

    ptr = PyMem_RawCalloc(0, 0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyMem_RawCalloc(0, 0) returns NULL");
        return NULL;
    }
    PyMem_RawFree(ptr);

    ptr = PyMem_Malloc(0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyMem_Malloc(0) returns NULL");
        return NULL;
    }
    PyMem_Free(ptr);

    ptr = PyMem_Calloc(0, 0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyMem_Calloc(0, 0) returns NULL");
        return NULL;
    }
    PyMem_Free(ptr);

    ptr = PyObject_Malloc(0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyObject_Malloc(0) returns NULL");
        return NULL;
    }
    PyObject_Free(ptr);

    ptr = PyObject_Calloc(0, 0);
    if (ptr == NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                        "PyObject_Calloc(0, 0) returns NULL");
        return NULL;
    }
    PyObject_Free(ptr);

    Py_RETURN_NONE;
}

static PyObject *
test_pymem_setrawallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    return test_setallocators(PYMEM_DOMAIN_RAW);
}

static PyObject *
test_pymem_setallocators(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    return test_setallocators(PYMEM_DOMAIN_MEM);
}

static PyObject *
pyobject_malloc_without_gil(PyObject *self, PyObject *args)
{
    char *buffer;

    /* Deliberate bug to test debug hooks on Python memory allocators:
       call PyObject_Malloc() without holding the GIL */
    Py_BEGIN_ALLOW_THREADS
    buffer = PyObject_Malloc(10);
    Py_END_ALLOW_THREADS

    PyObject_Free(buffer);

    Py_RETURN_NONE;
}

static PyObject *
pymem_buffer_overflow(PyObject *self, PyObject *args)
{
    char *buffer;

    /* Deliberate buffer overflow to check that PyMem_Free() detects
       the overflow when debug hooks are installed. */
    buffer = PyMem_Malloc(16);
    if (buffer == NULL) {
        PyErr_NoMemory();
        return NULL;
    }
    buffer[16] = 'x';
    PyMem_Free(buffer);

    Py_RETURN_NONE;
}

static PyObject *
pymem_api_misuse(PyObject *self, PyObject *args)
{
    char *buffer;

    /* Deliberate misusage of Python allocators:
       allococate with PyMem but release with PyMem_Raw. */
    buffer = PyMem_Malloc(16);
    PyMem_RawFree(buffer);

    Py_RETURN_NONE;
}

static PyObject *
pymem_malloc_without_gil(PyObject *self, PyObject *args)
{
    char *buffer;

    /* Deliberate bug to test debug hooks on Python memory allocators:
       call PyMem_Malloc() without holding the GIL */
    Py_BEGIN_ALLOW_THREADS
    buffer = PyMem_Malloc(10);
    Py_END_ALLOW_THREADS

    PyMem_Free(buffer);

    Py_RETURN_NONE;
}


// Tracemalloc tests
static PyObject *
tracemalloc_track(PyObject *self, PyObject *args)
{
    unsigned int domain;
    PyObject *ptr_obj;
    Py_ssize_t size;
    int release_gil = 0;

    if (!PyArg_ParseTuple(args, "IOn|i",
                          &domain, &ptr_obj, &size, &release_gil))
    {
        return NULL;
    }
    void *ptr = PyLong_AsVoidPtr(ptr_obj);
    if (PyErr_Occurred()) {
        return NULL;
    }

    int res;
    if (release_gil) {
        Py_BEGIN_ALLOW_THREADS
        res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
        Py_END_ALLOW_THREADS
    }
    else {
        res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
    }
    if (res < 0) {
        PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Track error");
        return NULL;
    }

    Py_RETURN_NONE;
}

static PyObject *
tracemalloc_untrack(PyObject *self, PyObject *args)
{
    unsigned int domain;
    PyObject *ptr_obj;
    int release_gil = 0;

    if (!PyArg_ParseTuple(args, "IO|i", &domain, &ptr_obj, &release_gil)) {
        return NULL;
    }
    void *ptr = PyLong_AsVoidPtr(ptr_obj);
    if (PyErr_Occurred()) {
        return NULL;
    }

    int res;
    if (release_gil) {
        Py_BEGIN_ALLOW_THREADS
        res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
        Py_END_ALLOW_THREADS
    }
    else {
        res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
    }
    if (res < 0) {
        PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Untrack error");
        return NULL;
    }

    Py_RETURN_NONE;
}


static void
tracemalloc_track_race_thread(void *data)
{
    PyTraceMalloc_Track(123, 10, 1);
    PyTraceMalloc_Untrack(123, 10);

    PyThread_type_lock lock = (PyThread_type_lock)data;
    PyThread_release_lock(lock);
}

// gh-128679: Test fix for tracemalloc.stop() race condition
static PyObject *
tracemalloc_track_race(PyObject *self, PyObject *args)
{
#define NTHREAD 50
    PyObject *tracemalloc = NULL;
    PyObject *stop = NULL;
    PyThread_type_lock locks[NTHREAD];
    memset(locks, 0, sizeof(locks));

    // Call tracemalloc.start()
    tracemalloc = PyImport_ImportModule("tracemalloc");
    if (tracemalloc == NULL) {
        goto error;
    }
    PyObject *start = PyObject_GetAttrString(tracemalloc, "start");
    if (start == NULL) {
        goto error;
    }
    PyObject *res = PyObject_CallNoArgs(start);
    Py_DECREF(start);
    if (res == NULL) {
        goto error;
    }
    Py_DECREF(res);

    stop = PyObject_GetAttrString(tracemalloc, "stop");
    Py_CLEAR(tracemalloc);
    if (stop == NULL) {
        goto error;
    }

    // Start threads
    for (size_t i = 0; i < NTHREAD; i++) {
        PyThread_type_lock lock = PyThread_allocate_lock();
        if (!lock) {
            PyErr_NoMemory();
            goto error;
        }
        locks[i] = lock;
        PyThread_acquire_lock(lock, 1);

        unsigned long thread;
        thread = PyThread_start_new_thread(tracemalloc_track_race_thread,
                                           (void*)lock);
        if (thread == (unsigned long)-1) {
            PyErr_SetString(PyExc_RuntimeError, "can't start new thread");
            goto error;
        }
    }

    // Call tracemalloc.stop() while threads are running
    res = PyObject_CallNoArgs(stop);
    Py_CLEAR(stop);
    if (res == NULL) {
        goto error;
    }
    Py_DECREF(res);

    // Wait until threads complete with the GIL released
    Py_BEGIN_ALLOW_THREADS
    for (size_t i = 0; i < NTHREAD; i++) {
        PyThread_type_lock lock = locks[i];
        PyThread_acquire_lock(lock, 1);
        PyThread_release_lock(lock);
    }
    Py_END_ALLOW_THREADS

    // Free threads locks
    for (size_t i=0; i < NTHREAD; i++) {
        PyThread_type_lock lock = locks[i];
        PyThread_free_lock(lock);
    }
    Py_RETURN_NONE;

error:
    Py_CLEAR(tracemalloc);
    Py_CLEAR(stop);
    for (size_t i=0; i < NTHREAD; i++) {
        PyThread_type_lock lock = locks[i];
        if (lock) {
            PyThread_free_lock(lock);
        }
    }
    return NULL;
#undef NTHREAD
}


#if TARGET_OS_OSX || defined(__FreeBSD__)
// Return RSS only. Per-process swap usage isn't readily available
static PyObject*
get_process_memory_usage(PyObject *self, PyObject *args)
{
    int pid;
    if (!PyArg_ParseTuple(args, "i", &pid)) {
        return NULL;
    }

#if TARGET_OS_OSX
    // macOS: proc_pidinfo(PROC_PIDTASKINFO).pti_resident_size
    struct proc_taskinfo pti;
    int ret = proc_pidinfo(pid, PROC_PIDTASKINFO, 0, &pti, sizeof(pti));
    if (ret <= 0) {
        if (errno == 0) {
            // proc_pidinfo() can return 0 without setting errno when the
            // process does not exist.
            errno = ESRCH;
        }
        return PyErr_SetFromErrno(PyExc_OSError);
    }

    return PyLong_FromUnsignedLongLong(pti.pti_resident_size);
#else
    // FreeBSD: kvm_getprocs(KERN_PROC_PID) and ki_rssize * page_size
    long page_size = sysconf(_SC_PAGESIZE);
    if (page_size <= 0) {
        return PyErr_SetFromErrno(PyExc_OSError);
    }

    // Using /dev/null for vmcore avoids needing dump file.
    // NULL for kernel file uses running kernel.
    char errbuf[_POSIX2_LINE_MAX];
    kvm_t *kd = kvm_openfiles(NULL, "/dev/null", NULL, O_RDONLY, errbuf);
    if (kd == NULL) {
        return PyErr_SetFromErrno(PyExc_OSError);
    }

    // KERN_PROC_PID filters for the specific process ID.
    int n_procs;
    struct kinfo_proc *kp = kvm_getprocs(kd, KERN_PROC_PID, pid, &n_procs);
    if (kp == NULL) {
        PyErr_SetFromErrno(PyExc_OSError);
        goto error;
    }
    if (n_procs <= 0) {
        // Process with PID not found
        errno = ESRCH;
        PyErr_SetFromErrno(PyExc_OSError);
        goto error;
    }
    assert(n_procs == 1);

    // ki_rssize is in pages. Convert to bytes.
    size_t rss = (size_t)kp[0].ki_rssize * page_size;
    kvm_close(kd);

    return PyLong_FromSize_t(rss);

error:
    kvm_close(kd);
    return NULL;
#endif
}
#endif


struct bytes_resize_tracer {
    PyObject *create;
    PyObject *destroy;
};


static int
bytes_resize_tracer(PyObject *obj, PyRefTracerEvent event, void* data)
{
    if (event != PyRefTracer_CREATE && event != PyRefTracer_DESTROY) {
        return 0;
    }

    struct bytes_resize_tracer *tracer = (struct bytes_resize_tracer*)data;
    if (!PyBytes_Check(obj)) {
        return 0;
    }

    switch (event) {
        case PyRefTracer_CREATE:
            tracer->create = obj;
            break;
        case PyRefTracer_DESTROY:
            tracer->destroy = obj;
            break;
        default:
            break;
    }
    return 0;
}


// When _PyBytes_Resize() resizes a bytes object in-place, check that
// PyRefTracer_DESTROY and PyRefTracer_CREATE events are emitted.
// If no_memory is non-zero, inject MemoryError.
static int
check_bytes_resize_tracer(int no_memory)
{
    PyObject *bytes = NULL;
    PyRefTracer old_tracer = NULL;
    void *old_tracer_data = NULL;
    int restore_tracer = 0;

    bytes = PyBytes_FromString("hello");
    if (bytes == NULL) {
        goto error;
    }
    assert(PyUnstable_Object_IsUniquelyReferenced(bytes));

    old_tracer = PyRefTracer_GetTracer(&old_tracer_data);
    restore_tracer = 1;

    struct bytes_resize_tracer tracer = {0};
    if (PyRefTracer_SetTracer(bytes_resize_tracer, &tracer) != 0) {
        goto error;
    }

    PyObject *old_bytes = bytes;  // borrowed reference
    if (no_memory) {
        fm_set_nomemory(0, 0);
        int res = _PyBytes_Resize(&bytes, 100);
        assert(res < 0);
        assert(bytes == NULL);
        fm_remove_hooks();

        assert(PyErr_ExceptionMatches(PyExc_MemoryError));
        PyErr_Clear();
    }
    else {
        if (_PyBytes_Resize(&bytes, 100) < 0) {
            assert(bytes == NULL);
            goto error;
        }
    }

    if (tracer.destroy != old_bytes) {
        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_DESTROY not seen");
        goto error;
    }

    int seen_create;
    if (no_memory) {
        seen_create = (tracer.create == old_bytes);
    }
    else {
        seen_create = (tracer.create == bytes);
    }
    if (!seen_create) {
        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_CREATE not seen");
        goto error;
    }

    Py_CLEAR(bytes);
    if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
        restore_tracer = 0;
        goto error;
    }
    return 0;

error:
    Py_XDECREF(bytes);
    if (restore_tracer) {
        if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
            return -1;
        }
    }
    return -1;
}


static PyObject*
test_bytes_resize_tracer(PyObject *self, PyObject *Py_UNUSED(ignored))
{
    if (check_bytes_resize_tracer(0) < 0) {
        return NULL;
    }
    if (check_bytes_resize_tracer(1) < 0) {
        return NULL;
    }
    Py_RETURN_NONE;
}


static PyMethodDef test_methods[] = {
    {"pymem_api_misuse",              pymem_api_misuse,              METH_NOARGS},
    {"pymem_buffer_overflow",         pymem_buffer_overflow,         METH_NOARGS},
    {"pymem_malloc_without_gil",      pymem_malloc_without_gil,      METH_NOARGS},
    {"pyobject_malloc_without_gil",   pyobject_malloc_without_gil,   METH_NOARGS},
    {"remove_mem_hooks",              remove_mem_hooks,              METH_NOARGS,
        PyDoc_STR("Remove memory hooks.")},
    {"set_nomemory",                  set_nomemory,                  METH_VARARGS,
        PyDoc_STR("set_nomemory(start:int, stop:int = 0)")},
    {"call_with_nomemory",            call_with_nomemory,            METH_VARARGS,
        PyDoc_STR("call_with_nomemory(start:int, stop:int, func, /, *args)")},
    {"test_pymem_alloc0",             test_pymem_alloc0,             METH_NOARGS},
    {"test_pymem_setallocators",      test_pymem_setallocators,      METH_NOARGS},
    {"test_pymem_setrawallocators",   test_pymem_setrawallocators,   METH_NOARGS},
    {"test_pyobject_new",             test_pyobject_new,             METH_NOARGS},
    {"test_pyobject_setallocators",   test_pyobject_setallocators,   METH_NOARGS},
#if TARGET_OS_OSX || defined(__FreeBSD__)
    {"get_process_memory_usage",      get_process_memory_usage,      METH_VARARGS},
#endif
    {"test_bytes_resize_tracer",      test_bytes_resize_tracer,      METH_NOARGS},

    // Tracemalloc tests
    {"tracemalloc_track",             tracemalloc_track,             METH_VARARGS},
    {"tracemalloc_untrack",           tracemalloc_untrack,           METH_VARARGS},
    {"tracemalloc_track_race", tracemalloc_track_race, METH_NOARGS},
    {NULL},
};

int
_PyTestCapi_Init_Mem(PyObject *mod)
{
    if (PyModule_AddFunctions(mod, test_methods) < 0) {
        return -1;
    }

    PyObject *v;
#ifdef WITH_PYMALLOC
    v = Py_True;
#else
    v = Py_False;
#endif
    if (PyModule_AddObjectRef(mod, "WITH_PYMALLOC", v) < 0) {
        return -1;
    }

#ifdef WITH_MIMALLOC
    v = Py_True;
#else
    v = Py_False;
#endif
    if (PyModule_AddObjectRef(mod, "WITH_MIMALLOC", v) < 0) {
        return -1;
    }

    return 0;
}