| #include <Python.h> |
| #include "py3compat.h" |
| #include <sched.h> |
| #include <errno.h> |
| #include <syscall.h> |
| |
| #ifndef __unused |
| #define __unused __attribute__ ((unused)) |
| #endif |
| |
| #ifndef CPU_ALLOC |
| |
| /* From glibc-headers-2.12-1.7.el6.x86_64 */ |
| |
| /* Basic access functions. */ |
| #define __CPUELT(cpu) ((cpu) / __NCPUBITS) |
| #define __CPUMASK(cpu) ((__cpu_mask) 1 << ((cpu) % __NCPUBITS)) |
| |
| #define CPU_ALLOC_SIZE(ncpus) \ |
| ((((ncpus) + __NCPUBITS - 1) / __NCPUBITS) * sizeof (__cpu_mask)) |
| |
| /* Access functions for CPU masks. */ |
| # if __GNUC_PREREQ (2, 91) |
| # define __CPU_ZERO_S(setsize, cpusetp) \ |
| do __builtin_memset (cpusetp, '\0', setsize); while (0) |
| # else |
| # define __CPU_ZERO_S(setsize, cpusetp) \ |
| do { \ |
| size_t __i; \ |
| size_t __imax = (setsize) / sizeof (__cpu_mask); \ |
| __cpu_mask *__bits = (cpusetp)->__bits; \ |
| for (__i = 0; __i < __imax; ++__i) \ |
| __bits[__i] = 0; \ |
| } while (0) |
| # endif |
| |
| #define CPU_FREE(cpus) free(cpus) |
| #define CPU_ALLOC(ncpus) malloc(CPU_ALLOC_SIZE(ncpus)) |
| |
| # define __CPU_SET_S(cpu, setsize, cpusetp) \ |
| (__extension__ \ |
| ({ size_t __cpu = (cpu); \ |
| __cpu < 8 * (setsize) \ |
| ? (((__cpu_mask *) ((cpusetp)->__bits))[__CPUELT (__cpu)] \ |
| |= __CPUMASK (__cpu)) \ |
| : 0; })) |
| |
| # define __CPU_ISSET_S(cpu, setsize, cpusetp) \ |
| (__extension__ \ |
| ({ size_t __cpu = (cpu); \ |
| __cpu < 8 * (setsize) \ |
| ? ((((__const __cpu_mask *) ((cpusetp)->__bits))[__CPUELT (__cpu)] \ |
| & __CPUMASK (__cpu))) != 0 \ |
| : 0; })) |
| |
| #define CPU_SET_S(cpu, setsize, cpusetp) __CPU_SET_S (cpu, setsize, cpusetp) |
| #define CPU_ISSET_S(cpu, setsize, cpusetp) __CPU_ISSET_S (cpu, setsize, cpusetp) |
| #define CPU_ZERO_S(setsize, cpusetp) __CPU_ZERO_S (setsize, cpusetp) |
| |
| #endif /* CPU_ALLOC */ |
| |
| |
| #ifndef SCHED_RESET_ON_FORK |
| #define SCHED_RESET_ON_FORK 0x40000000 |
| #endif |
| |
| /* |
| * The following bitmask declarations, bitmask_*() routines, and associated |
| * _setbit() and _getbit() routines are: |
| * Copyright (c) 2004 Silicon Graphics, Inc. (SGI) All rights reserved. |
| * SGI publishes it under the terms of the GNU General Public License, v2, |
| * as published by the Free Software Foundation. |
| */ |
| #define howmany(x,y) (((x)+((y)-1))/(y)) |
| #define bitsperlong (8 * sizeof(unsigned long)) |
| #define longsperbits(n) howmany(n, bitsperlong) |
| #define bytesperbits(x) ((x+7)/8) |
| |
| /* |
| * Number of bits in a CPU bitmask on current system |
| */ |
| static int __get_max_number_of_cpus(void) |
| { |
| int n; |
| int cpus = 2048; |
| |
| for (;;) { |
| unsigned long buffer[longsperbits(cpus)]; |
| memset(buffer, 0, sizeof(buffer)); |
| /* the library version does not return size of cpumask_t */ |
| n = syscall(SYS_sched_getaffinity, 0, bytesperbits(cpus), &buffer); |
| if (n < 0 && errno == EINVAL && cpus < 1024 * 1024) { |
| cpus *= 2; |
| continue; |
| } |
| return n * 8; |
| } |
| return -1; |
| } |
| |
| static PyObject *get_max_number_of_cpus(PyObject *self __unused, PyObject *args __unused) |
| { |
| int ret = __get_max_number_of_cpus(); |
| |
| if (ret < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| return Py_BuildValue("i", ret); |
| } |
| |
| static PyObject *get_affinity(PyObject *self __unused, PyObject *args) |
| { |
| PyObject *list = NULL; |
| cpu_set_t *cpus; |
| int pid, cpu; |
| size_t cpusetsize; |
| int max_cpus; |
| |
| if (!PyArg_ParseTuple(args, "i", &pid)) |
| goto out_error; |
| |
| max_cpus = __get_max_number_of_cpus(); |
| if (max_cpus < 0) |
| goto out_error; |
| |
| cpus = CPU_ALLOC(max_cpus); |
| if (cpus == NULL) |
| return PyErr_NoMemory(); |
| |
| cpusetsize = CPU_ALLOC_SIZE(max_cpus); |
| CPU_ZERO_S(cpusetsize, cpus); |
| |
| if (sched_getaffinity(pid, cpusetsize, cpus) < 0) |
| goto out_free; |
| |
| list = PyList_New(0); |
| for (cpu = 0; cpu < max_cpus; ++cpu) |
| if (CPU_ISSET_S(cpu, cpusetsize, cpus)) |
| PyList_Append(list, Py_BuildValue("i", cpu)); |
| |
| CPU_FREE(cpus); |
| |
| return list; |
| out_free: |
| CPU_FREE(cpus); |
| out_error: |
| return PyErr_SetFromErrno(PyExc_OSError); |
| } |
| |
| static PyObject *set_affinity(PyObject *self __unused, PyObject *args) |
| { |
| int pid, nr_elements, i, max_cpus; |
| cpu_set_t *cpus; |
| PyObject *list; |
| size_t cpusetsize; |
| |
| if (!PyArg_ParseTuple(args, "iO", &pid, &list)) |
| goto out_error; |
| |
| max_cpus = __get_max_number_of_cpus(); |
| if (max_cpus < 0) |
| goto out_error; |
| |
| cpus = CPU_ALLOC(max_cpus); |
| if (cpus == NULL) |
| return PyErr_NoMemory(); |
| |
| cpusetsize = CPU_ALLOC_SIZE(max_cpus); |
| CPU_ZERO_S(cpusetsize, cpus); |
| |
| nr_elements = PyList_Size(list); |
| for (i = 0; i < nr_elements; ++i) { |
| int cpu = PyInt_AsLong(PyList_GetItem(list, i)); |
| |
| if (cpu >= max_cpus) { |
| PyErr_SetString(PyExc_OSError, "Invalid CPU"); |
| CPU_FREE(cpus); |
| return NULL; |
| } |
| CPU_SET_S(cpu, cpusetsize, cpus); |
| } |
| |
| i = sched_setaffinity(pid, cpusetsize, cpus); |
| CPU_FREE(cpus); |
| |
| if (i < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| Py_INCREF(Py_None); |
| return Py_None; |
| out_error: |
| return PyErr_SetFromErrno(PyExc_OSError); |
| } |
| |
| static PyObject *get_scheduler(PyObject *self __unused, PyObject *args) |
| { |
| int pid, scheduler; |
| |
| if (!PyArg_ParseTuple(args, "i", &pid)) |
| return NULL; |
| |
| scheduler = sched_getscheduler(pid); |
| if (scheduler < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| return Py_BuildValue("i", scheduler); |
| } |
| |
| static PyObject *set_scheduler(PyObject *self __unused, PyObject *args) |
| { |
| int pid, policy, priority; |
| struct sched_param param; |
| |
| if (!PyArg_ParseTuple(args, "iii", &pid, &policy, &priority)) |
| return NULL; |
| |
| memset(¶m, 0, sizeof(param)); |
| param.sched_priority = priority; |
| |
| if (sched_setscheduler(pid, policy, ¶m) < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| Py_RETURN_NONE; |
| } |
| |
| static PyObject *get_priority(PyObject *self __unused, PyObject *args) |
| { |
| int pid; |
| struct sched_param param = { .sched_priority = -1, }; |
| |
| if (!PyArg_ParseTuple(args, "i", &pid)) |
| return NULL; |
| |
| if (sched_getparam(pid, ¶m) != 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| return Py_BuildValue("i", param.sched_priority); |
| } |
| |
| #ifndef SCHED_BATCH |
| #define SCHED_BATCH 3 |
| #endif |
| #ifndef SCHED_IDLE |
| #define SCHED_IDLE 5 |
| #endif |
| #ifndef SCHED_DEADLINE |
| #define SCHED_DEADLINE 6 |
| #endif |
| |
| |
| static PyObject *schedstr(PyObject *self __unused, PyObject *args) |
| { |
| int scheduler; |
| char *s; |
| |
| if (!PyArg_ParseTuple(args, "i", &scheduler)) |
| return NULL; |
| |
| switch (scheduler & ~SCHED_RESET_ON_FORK) { |
| case SCHED_OTHER: s = "SCHED_OTHER"; break; |
| case SCHED_RR: s = "SCHED_RR"; break; |
| case SCHED_FIFO: s = "SCHED_FIFO"; break; |
| case SCHED_BATCH: s = "SCHED_BATCH"; break; |
| case SCHED_IDLE: s = "SCHED_IDLE"; break; |
| case SCHED_DEADLINE: s = "SCHED_DEADLINE"; break; |
| default: s = "UNKNOWN"; break; |
| } |
| |
| return PyStr_FromString(s); |
| } |
| |
| static PyObject *schedfromstr(PyObject *self __unused, PyObject *args) |
| { |
| int scheduler; |
| char *s; |
| |
| if (!PyArg_ParseTuple(args, "s", &s)) |
| return NULL; |
| |
| if (strcmp(s, "SCHED_OTHER") == 0) |
| scheduler = SCHED_OTHER; |
| else if (strcmp(s, "SCHED_RR") == 0) |
| scheduler = SCHED_RR; |
| else if (strcmp(s, "SCHED_FIFO") == 0) |
| scheduler = SCHED_FIFO; |
| else if (strcmp(s, "SCHED_BATCH") == 0) |
| scheduler = SCHED_BATCH; |
| else if (strcmp(s, "SCHED_IDLE") == 0) |
| scheduler = SCHED_IDLE; |
| else if (strcmp(s, "SCHED_DEADLINE") == 0) |
| scheduler = SCHED_DEADLINE; |
| else { |
| PyErr_SetString(PyExc_OSError, "Unknown scheduler"); |
| return NULL; |
| } |
| |
| return Py_BuildValue("i", scheduler); |
| } |
| |
| static PyObject *get_priority_min(PyObject *self __unused, PyObject *args) |
| { |
| int policy, min; |
| |
| if (!PyArg_ParseTuple(args, "i", &policy)) |
| return NULL; |
| |
| min = sched_get_priority_min(policy); |
| if (min < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| return Py_BuildValue("i", min); |
| } |
| |
| static PyObject *get_priority_max(PyObject *self __unused, PyObject *args) |
| { |
| int policy, max; |
| |
| if (!PyArg_ParseTuple(args, "i", &policy)) |
| return NULL; |
| |
| max = sched_get_priority_max(policy); |
| if (max < 0) |
| return PyErr_SetFromErrno(PyExc_OSError); |
| |
| return Py_BuildValue("i", max); |
| } |
| |
| static struct PyMethodDef PySchedutilsModuleMethods[] = { |
| { |
| .ml_name = "get_affinity", |
| .ml_meth = (PyCFunction)get_affinity, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "set_affinity", |
| .ml_meth = (PyCFunction)set_affinity, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "get_scheduler", |
| .ml_meth = (PyCFunction)get_scheduler, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "set_scheduler", |
| .ml_meth = (PyCFunction)set_scheduler, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "get_priority", |
| .ml_meth = (PyCFunction)get_priority, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "schedstr", |
| .ml_meth = (PyCFunction)schedstr, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "schedfromstr", |
| .ml_meth = (PyCFunction)schedfromstr, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "get_priority_min", |
| .ml_meth = (PyCFunction)get_priority_min, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "get_priority_max", |
| .ml_meth = (PyCFunction)get_priority_max, |
| .ml_flags = METH_VARARGS, |
| }, |
| { |
| .ml_name = "get_max_number_of_cpus", |
| .ml_meth = (PyCFunction)get_max_number_of_cpus, |
| .ml_flags = METH_NOARGS, |
| }, |
| { .ml_name = NULL, }, |
| }; |
| |
| static struct PyModuleDef moduledef = { |
| PyModuleDef_HEAD_INIT, |
| .m_name = "schedutils", |
| .m_doc = NULL, |
| .m_size = -1, |
| .m_methods = PySchedutilsModuleMethods, |
| }; |
| |
| MODULE_INIT_FUNC(schedutils) |
| { |
| PyObject *m = PyModule_Create(&moduledef); |
| if (m == NULL) |
| return NULL; |
| |
| PyModule_AddIntConstant(m, "SCHED_OTHER", SCHED_OTHER); |
| PyModule_AddIntConstant(m, "SCHED_FIFO", SCHED_FIFO); |
| PyModule_AddIntConstant(m, "SCHED_RR", SCHED_RR); |
| PyModule_AddIntConstant(m, "SCHED_BATCH", SCHED_BATCH); |
| PyModule_AddIntConstant(m, "SCHED_IDLE", SCHED_IDLE); |
| PyModule_AddIntConstant(m, "SCHED_DEADLINE", SCHED_DEADLINE); |
| PyModule_AddIntConstant(m, "SCHED_RESET_ON_FORK", SCHED_RESET_ON_FORK); |
| |
| return m; |
| } |
| |