diff --git a/Modules/spinner.c b/Modules/spinner.c new file mode 100644 index 000000000000000..4631f2677bbbba7 --- /dev/null +++ b/Modules/spinner.c @@ -0,0 +1,117 @@ +#include +#include +#include +#include +#include + +// The set_priority function +static PyObject* set_priority(PyObject* self, PyObject* args) { + int policy, priority; + + // Parse the arguments + if (!PyArg_ParseTuple(args, "ii", &policy, &priority)) { + return NULL; + } + + // Get the current thread id + pid_t tid = syscall(SYS_gettid); + + // Set the scheduling policy and priority + struct sched_param param; + param.sched_priority = priority; + if (sched_setscheduler(tid, policy, ¶m) != 0) { + return PyErr_SetFromErrno(PyExc_OSError); + } + + Py_RETURN_NONE; +} + +// The get_priority function +static PyObject* get_priority(PyObject* self, PyObject* args) { + // Get the current thread id + pid_t tid = syscall(SYS_gettid); + + // Get the scheduling policy + int policy = sched_getscheduler(tid); + if (policy == -1) { + return PyErr_SetFromErrno(PyExc_OSError); + } + + // Get the priority + struct sched_param param; + if (sched_getparam(tid, ¶m) != 0) { + return PyErr_SetFromErrno(PyExc_OSError); + } + + return Py_BuildValue("(ii)", policy, param.sched_priority); +} + +// The spin function +static PyObject* spin(PyObject* self, PyObject* args) { + int milliseconds; + double spins; + + // Parse the arguments + if (!PyArg_ParseTuple(args, "i", &milliseconds)) { + return NULL; + } + + // Save the current state and release the GIL + Py_BEGIN_ALLOW_THREADS + + // Get the current time + struct timespec start, now; + clock_gettime(CLOCK_MONOTONIC, &start); + + // Calculate the end time + time_t end_sec = start.tv_sec + milliseconds / 1000; + long end_nsec = start.tv_nsec + (milliseconds % 1000) * 1000000; + if (end_nsec >= 1000000000) { + end_sec++; + end_nsec -= 1000000000; + } + + // Spitn until the end time + spins = 0; + do { + clock_gettime(CLOCK_MONOTONIC, &now); + ++spins; + } while (now.tv_sec < end_sec || (now.tv_sec == end_sec && now.tv_nsec < end_nsec)); + + // Restore the GIL and the saved state + Py_END_ALLOW_THREADS + + return Py_BuildValue("d", spins); +} + +// The module's function table +static PyMethodDef SpinnerMethods[] = { + {"set_priority", set_priority, METH_VARARGS, "Set the current thread's priority."}, + {"get_priority", get_priority, METH_VARARGS, "Get the current thread's priority."}, + {"spin", spin, METH_VARARGS, "Spin the CPU for a specified time."}, + // Add other functions here... + {NULL, NULL, 0, NULL} +}; + +// The module's definition +static struct PyModuleDef spinnermodule = { + PyModuleDef_HEAD_INIT, + "spinner", + "A module that spins the CPU and manages thread priorities.", + -1, + SpinnerMethods +}; + +// The module's initialization function +PyMODINIT_FUNC PyInit_spinner(void) { + PyObject* m; + m = PyModule_Create(&spinnermodule); + if (m == NULL) + return NULL; + PyModule_AddIntConstant(m, "SCHED_IDLE", SCHED_IDLE); + PyModule_AddIntConstant(m, "SCHED_BATCH", SCHED_BATCH); + PyModule_AddIntConstant(m, "SCHED_OTHER", SCHED_OTHER); + PyModule_AddIntConstant(m, "SCHED_FIFO", SCHED_FIFO); + PyModule_AddIntConstant(m, "SCHED_RR", SCHED_FIFO); + return m; +}