profiling.sampling CPU mode may attribute on-CPU samples to detached tstates when sharing native_thread_id #159312

#define _GNU_SOURCE
#include <Python.h>
#include <stdio.h>
#include <sys/syscall.h>
#include <unistd.h>

static long
gettid_linux(void)
{
    return (long)syscall(SYS_gettid);
}

static void
print_tstate_info(const char *where, PyThreadState *tstate)
{
    PyThreadState *current = PyThreadState_GetUnchecked();
    unsigned long native_tid = tstate ? tstate->native_thread_id : 0;
    long tid = gettid_linux();

    fprintf(stderr,
            "TSTATE %-24s tstate=%p current=%p native_thread_id=%lu "
            "gettid=%ld match=%s\n",
            where, (void *)tstate, (void *)current, native_tid, tid,
            tstate && native_tid == (unsigned long)tid ? "yes" : "no");
    fflush(stderr);
}

static PyObject *
print_tstate(PyObject *self, PyObject *args)
{
    const char *where;
    if (!PyArg_ParseTuple(args, "s", &where)) {
        return NULL;
    }
    print_tstate_info(where, PyThreadState_Get());
    Py_RETURN_NONE;
}

static PyObject *
switch_tstate_and_call(PyObject *self, PyObject *callback)
{
    if (!PyCallable_Check(callback)) {
        PyErr_SetString(PyExc_TypeError, "callback must be callable");
        return NULL;
    }

    PyThreadState *tstate0 = PyThreadState_Get();
    PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate0);
    print_tstate_info("tstate0 in funB", tstate0);

    PyThreadState *saved0 = PyEval_SaveThread();
    print_tstate_info("tstate0 detached", saved0);

    PyThreadState *tstate1 = PyThreadState_New(interp);
    if (tstate1 == NULL) {
        PyEval_RestoreThread(saved0);
        return PyErr_NoMemory();
    }

    print_tstate_info("tstate1 created", tstate1);
    PyEval_RestoreThread(tstate1);
    print_tstate_info("tstate1 attached", tstate1);

    /* The callback enters funC -> funD, and funD intentionally does not return. */
    PyObject *result = PyObject_CallNoArgs(callback);

    /* Cleanup remains correct if funC is changed to return. */
    PyObject *raised = result ? NULL : PyErr_GetRaisedException();
    Py_XDECREF(result);
    PyThreadState *saved1 = PyEval_SaveThread();
    PyEval_RestoreThread(saved0);
    PyThreadState_Clear(saved1);
    PyThreadState_Delete(saved1);

    if (raised) {
        PyErr_SetRaisedException(raised);
        return NULL;
    }
    Py_RETURN_NONE;
}

static PyMethodDef methods[] = {
    {"print_tstate", print_tstate, METH_VARARGS, NULL},
    {"switch_tstate_and_call", switch_tstate_and_call, METH_O, NULL},
    {NULL, NULL, 0, NULL}
};

static PyModuleDef_Slot slots[] = {
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
    {0, NULL}
};

static struct PyModuleDef module = {
    PyModuleDef_HEAD_INIT, "_native", NULL, 0, methods, slots
};

PyMODINIT_FUNC
PyInit__native(void)
{
    return PyModuleDef_Init(&module);
}