Merge branch 'main' of github.com:SonicField/cpython

Modules/spinner.c

@@ -0,0 +1,117 @@
+#include <Python.h>
+#include <time.h>
+#include <sched.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+// 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, &param) != 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, &param) != 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;
+}