summaryrefslogtreecommitdiffstats
path: root/Modules
diff options
context:
space:
mode:
authorAntoine Pitrou <solipsis@pitrou.net>2012-08-04 14:16:35 (GMT)
committerAntoine Pitrou <solipsis@pitrou.net>2012-08-04 14:16:35 (GMT)
commit848698727fcbb633246b56ab57080b4d5493c186 (patch)
treeaf5543aa381427085e745fc5c139c94ea15b94ee /Modules
parenta0e3febe5ae4b01f2076d1f0f09b76f366611a16 (diff)
downloadcpython-848698727fcbb633246b56ab57080b4d5493c186.zip
cpython-848698727fcbb633246b56ab57080b4d5493c186.tar.gz
cpython-848698727fcbb633246b56ab57080b4d5493c186.tar.bz2
Issue #12655: Instead of requiring a custom type, os.sched_getaffinity and
os.sched_setaffinity now use regular sets of integers to represent the CPUs a process is restricted to.
Diffstat (limited to 'Modules')
-rw-r--r--Modules/posixmodule.c451
1 files changed, 130 insertions, 321 deletions
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index ce10d1b..1f2b0cc 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -5720,322 +5720,100 @@ posix_sched_yield(PyObject *self, PyObject *noargs)
#ifdef HAVE_SCHED_SETAFFINITY
-typedef struct {
- PyObject_HEAD;
- Py_ssize_t size;
- int ncpus;
- cpu_set_t *set;
-} Py_cpu_set;
-
-static PyTypeObject cpu_set_type;
-
-static void
-cpu_set_dealloc(Py_cpu_set *set)
-{
- assert(set->set);
- CPU_FREE(set->set);
- Py_TYPE(set)->tp_free(set);
-}
-
-static Py_cpu_set *
-make_new_cpu_set(PyTypeObject *type, Py_ssize_t size)
-{
- Py_cpu_set *set;
+/* The minimum number of CPUs allocated in a cpu_set_t */
+static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
- if (size < 0) {
- PyErr_SetString(PyExc_ValueError, "negative size");
- return NULL;
- }
- set = (Py_cpu_set *)type->tp_alloc(type, 0);
- if (!set)
- return NULL;
- set->ncpus = size;
- set->size = CPU_ALLOC_SIZE(size);
- set->set = CPU_ALLOC(size);
- if (!set->set) {
- type->tp_free(set);
- PyErr_NoMemory();
- return NULL;
- }
- CPU_ZERO_S(set->size, set->set);
- return set;
-}
-
-static PyObject *
-cpu_set_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
-{
- int size;
-
- if (!_PyArg_NoKeywords("cpu_set()", kwargs) ||
- !PyArg_ParseTuple(args, "i:cpu_set", &size))
- return NULL;
- return (PyObject *)make_new_cpu_set(type, size);
-}
-
-PyDoc_STRVAR(cpu_set_sizeof_doc,
-"cpu_set.__sizeof__() -> int\n\n\
-Returns size in memory, in bytes.");
-
-static PyObject *
-cpu_set_sizeof(Py_cpu_set *set, PyObject *noargs)
-{
- Py_ssize_t res = 0;
-
- res = sizeof(Py_cpu_set);
- res += set->size;
- return PyLong_FromSsize_t(res);
-}
+PyDoc_STRVAR(posix_sched_setaffinity__doc__,
+"sched_setaffinity(pid, cpu_set)\n\n\
+Set the affinity of the process with PID *pid* to *cpu_set*.");
static PyObject *
-cpu_set_repr(Py_cpu_set *set)
-{
- return PyUnicode_FromFormat("<cpu_set with %li entries>", set->ncpus);
-}
-
-static Py_ssize_t
-cpu_set_len(Py_cpu_set *set)
-{
- return set->ncpus;
-}
-
-static int
-_get_cpu(Py_cpu_set *set, const char *requester, PyObject *args)
+posix_sched_setaffinity(PyObject *self, PyObject *args)
{
- int cpu;
- if (!PyArg_ParseTuple(args, requester, &cpu))
- return -1;
- if (cpu < 0) {
- PyErr_SetString(PyExc_ValueError, "cpu < 0 not valid");
- return -1;
- }
- if (cpu >= set->ncpus) {
- PyErr_SetString(PyExc_ValueError, "cpu too large for set");
- return -1;
- }
- return cpu;
-}
-
-PyDoc_STRVAR(cpu_set_set_doc,
-"cpu_set.set(i)\n\n\
-Add CPU *i* to the set.");
+ pid_t pid;
+ int ncpus;
+ size_t setsize;
+ cpu_set_t *mask = NULL;
+ PyObject *iterable, *iterator = NULL, *item;
-static PyObject *
-cpu_set_set(Py_cpu_set *set, PyObject *args)
-{
- int cpu = _get_cpu(set, "i|set", args);
- if (cpu == -1)
+ if (!PyArg_ParseTuple(args, _Py_PARSE_PID "O:sched_setaffinity",
+ &pid, &iterable))
return NULL;
- CPU_SET_S(cpu, set->size, set->set);
- Py_RETURN_NONE;
-}
-
-PyDoc_STRVAR(cpu_set_count_doc,
-"cpu_set.count() -> int\n\n\
-Return the number of CPUs active in the set.");
-
-static PyObject *
-cpu_set_count(Py_cpu_set *set, PyObject *noargs)
-{
- return PyLong_FromLong(CPU_COUNT_S(set->size, set->set));
-}
-PyDoc_STRVAR(cpu_set_clear_doc,
-"cpu_set.clear(i)\n\n\
-Remove CPU *i* from the set.");
-
-static PyObject *
-cpu_set_clear(Py_cpu_set *set, PyObject *args)
-{
- int cpu = _get_cpu(set, "i|clear", args);
- if (cpu == -1)
+ iterator = PyObject_GetIter(iterable);
+ if (iterator == NULL)
return NULL;
- CPU_CLR_S(cpu, set->size, set->set);
- Py_RETURN_NONE;
-}
-PyDoc_STRVAR(cpu_set_isset_doc,
-"cpu_set.isset(i) -> bool\n\n\
-Test if CPU *i* is in the set.");
-
-static PyObject *
-cpu_set_isset(Py_cpu_set *set, PyObject *args)
-{
- int cpu = _get_cpu(set, "i|isset", args);
- if (cpu == -1)
- return NULL;
- if (CPU_ISSET_S(cpu, set->size, set->set))
- Py_RETURN_TRUE;
- Py_RETURN_FALSE;
-}
+ ncpus = NCPUS_START;
+ setsize = CPU_ALLOC_SIZE(ncpus);
+ mask = CPU_ALLOC(ncpus);
+ if (mask == NULL) {
+ PyErr_NoMemory();
+ goto error;
+ }
+ CPU_ZERO_S(setsize, mask);
-PyDoc_STRVAR(cpu_set_zero_doc,
-"cpu_set.zero()\n\n\
-Clear the cpu_set.");
+ while ((item = PyIter_Next(iterator))) {
+ long cpu;
+ if (!PyLong_Check(item)) {
+ PyErr_Format(PyExc_TypeError,
+ "expected an iterator of ints, "
+ "but iterator yielded %R",
+ Py_TYPE(item));
+ Py_DECREF(item);
+ goto error;
+ }
+ cpu = PyLong_AsLong(item);
+ Py_DECREF(item);
+ if (cpu < 0) {
+ if (!PyErr_Occurred())
+ PyErr_SetString(PyExc_ValueError, "negative CPU number");
+ goto error;
+ }
+ if (cpu > INT_MAX - 1) {
+ PyErr_SetString(PyExc_OverflowError, "CPU number too large");
+ goto error;
+ }
+ if (cpu >= ncpus) {
+ /* Grow CPU mask to fit the CPU number */
+ int newncpus = ncpus;
+ cpu_set_t *newmask;
+ size_t newsetsize;
+ while (newncpus <= cpu) {
+ if (newncpus > INT_MAX / 2)
+ newncpus = cpu + 1;
+ else
+ newncpus = newncpus * 2;
+ }
+ newmask = CPU_ALLOC(newncpus);
+ if (newmask == NULL) {
+ PyErr_NoMemory();
+ goto error;
+ }
+ newsetsize = CPU_ALLOC_SIZE(newncpus);
+ CPU_ZERO_S(newsetsize, newmask);
+ memcpy(newmask, mask, setsize);
+ CPU_FREE(mask);
+ setsize = newsetsize;
+ mask = newmask;
+ ncpus = newncpus;
+ }
+ CPU_SET_S(cpu, setsize, mask);
+ }
+ Py_CLEAR(iterator);
-static PyObject *
-cpu_set_zero(Py_cpu_set *set, PyObject *noargs)
-{
- CPU_ZERO_S(set->size, set->set);
+ if (sched_setaffinity(pid, setsize, mask)) {
+ posix_error();
+ goto error;
+ }
+ CPU_FREE(mask);
Py_RETURN_NONE;
-}
-
-static PyObject *
-cpu_set_richcompare(Py_cpu_set *set, Py_cpu_set *other, int op)
-{
- int eq;
-
- if ((op != Py_EQ && op != Py_NE) || Py_TYPE(other) != &cpu_set_type)
- Py_RETURN_NOTIMPLEMENTED;
- eq = set->ncpus == other->ncpus && CPU_EQUAL_S(set->size, set->set, other->set);
- if ((op == Py_EQ) ? eq : !eq)
- Py_RETURN_TRUE;
- else
- Py_RETURN_FALSE;
-}
-
-#define CPU_SET_BINOP(name, op) \
- static PyObject * \
- do_cpu_set_##name(Py_cpu_set *left, Py_cpu_set *right, Py_cpu_set *res) { \
- if (res) { \
- Py_INCREF(res); \
- } \
- else { \
- res = make_new_cpu_set(&cpu_set_type, left->ncpus); \
- if (!res) \
- return NULL; \
- } \
- if (Py_TYPE(right) != &cpu_set_type || left->ncpus != right->ncpus) { \
- Py_DECREF(res); \
- Py_RETURN_NOTIMPLEMENTED; \
- } \
- assert(left->size == right->size && right->size == res->size); \
- op(res->size, res->set, left->set, right->set); \
- return (PyObject *)res; \
- } \
- static PyObject * \
- cpu_set_##name(Py_cpu_set *left, Py_cpu_set *right) { \
- return do_cpu_set_##name(left, right, NULL); \
- } \
- static PyObject * \
- cpu_set_i##name(Py_cpu_set *left, Py_cpu_set *right) { \
- return do_cpu_set_##name(left, right, left); \
- } \
-
-CPU_SET_BINOP(and, CPU_AND_S)
-CPU_SET_BINOP(or, CPU_OR_S)
-CPU_SET_BINOP(xor, CPU_XOR_S)
-#undef CPU_SET_BINOP
-
-PyDoc_STRVAR(cpu_set_doc,
-"cpu_set(size)\n\n\
-Create an empty mask of CPUs.");
-
-static PyNumberMethods cpu_set_as_number = {
- 0, /*nb_add*/
- 0, /*nb_subtract*/
- 0, /*nb_multiply*/
- 0, /*nb_remainder*/
- 0, /*nb_divmod*/
- 0, /*nb_power*/
- 0, /*nb_negative*/
- 0, /*nb_positive*/
- 0, /*nb_absolute*/
- 0, /*nb_bool*/
- 0, /*nb_invert*/
- 0, /*nb_lshift*/
- 0, /*nb_rshift*/
- (binaryfunc)cpu_set_and, /*nb_and*/
- (binaryfunc)cpu_set_xor, /*nb_xor*/
- (binaryfunc)cpu_set_or, /*nb_or*/
- 0, /*nb_int*/
- 0, /*nb_reserved*/
- 0, /*nb_float*/
- 0, /*nb_inplace_add*/
- 0, /*nb_inplace_subtract*/
- 0, /*nb_inplace_multiply*/
- 0, /*nb_inplace_remainder*/
- 0, /*nb_inplace_power*/
- 0, /*nb_inplace_lshift*/
- 0, /*nb_inplace_rshift*/
- (binaryfunc)cpu_set_iand, /*nb_inplace_and*/
- (binaryfunc)cpu_set_ixor, /*nb_inplace_xor*/
- (binaryfunc)cpu_set_ior, /*nb_inplace_or*/
-};
-
-static PySequenceMethods cpu_set_as_sequence = {
- (lenfunc)cpu_set_len, /* sq_length */
-};
-
-static PyMethodDef cpu_set_methods[] = {
- {"clear", (PyCFunction)cpu_set_clear, METH_VARARGS, cpu_set_clear_doc},
- {"count", (PyCFunction)cpu_set_count, METH_NOARGS, cpu_set_count_doc},
- {"isset", (PyCFunction)cpu_set_isset, METH_VARARGS, cpu_set_isset_doc},
- {"set", (PyCFunction)cpu_set_set, METH_VARARGS, cpu_set_set_doc},
- {"zero", (PyCFunction)cpu_set_zero, METH_NOARGS, cpu_set_zero_doc},
- {"__sizeof__", (PyCFunction)cpu_set_sizeof, METH_NOARGS, cpu_set_sizeof_doc},
- {NULL, NULL} /* sentinel */
-};
-
-static PyTypeObject cpu_set_type = {
- PyVarObject_HEAD_INIT(&PyType_Type, 0)
- "posix.cpu_set", /* tp_name */
- sizeof(Py_cpu_set), /* tp_basicsize */
- 0, /* tp_itemsize */
- /* methods */
- (destructor)cpu_set_dealloc, /* tp_dealloc */
- 0, /* tp_print */
- 0, /* tp_getattr */
- 0, /* tp_setattr */
- 0, /* tp_reserved */
- (reprfunc)cpu_set_repr, /* tp_repr */
- &cpu_set_as_number, /* tp_as_number */
- &cpu_set_as_sequence, /* tp_as_sequence */
- 0, /* tp_as_mapping */
- PyObject_HashNotImplemented, /* tp_hash */
- 0, /* tp_call */
- 0, /* tp_str */
- PyObject_GenericGetAttr, /* tp_getattro */
- 0, /* tp_setattro */
- 0, /* tp_as_buffer */
- Py_TPFLAGS_DEFAULT, /* tp_flags */
- cpu_set_doc, /* tp_doc */
- 0, /* tp_traverse */
- 0, /* tp_clear */
- (richcmpfunc)cpu_set_richcompare, /* tp_richcompare */
- 0, /* tp_weaklistoffset */
- 0, /* tp_iter */
- 0, /* tp_iternext */
- cpu_set_methods, /* tp_methods */
- 0, /* tp_members */
- 0, /* tp_getset */
- 0, /* tp_base */
- 0, /* tp_dict */
- 0, /* tp_descr_get */
- 0, /* tp_descr_set */
- 0, /* tp_dictoffset */
- 0, /* tp_init */
- PyType_GenericAlloc, /* tp_alloc */
- cpu_set_new, /* tp_new */
- PyObject_Del, /* tp_free */
-};
-
-PyDoc_STRVAR(posix_sched_setaffinity__doc__,
-"sched_setaffinity(pid, cpu_set)\n\n\
-Set the affinity of the process with PID *pid* to *cpu_set*.");
-
-static PyObject *
-posix_sched_setaffinity(PyObject *self, PyObject *args)
-{
- pid_t pid;
- Py_cpu_set *cpu_set;
-
- if (!PyArg_ParseTuple(args, _Py_PARSE_PID "O!:sched_setaffinity",
- &pid, &cpu_set_type, &cpu_set))
- return NULL;
- if (sched_setaffinity(pid, cpu_set->size, cpu_set->set))
- return posix_error();
- Py_RETURN_NONE;
+error:
+ if (mask)
+ CPU_FREE(mask);
+ Py_XDECREF(iterator);
+ return NULL;
}
PyDoc_STRVAR(posix_sched_getaffinity__doc__,
@@ -6047,20 +5825,58 @@ static PyObject *
posix_sched_getaffinity(PyObject *self, PyObject *args)
{
pid_t pid;
- int ncpus;
- Py_cpu_set *res;
+ int cpu, ncpus, count;
+ size_t setsize;
+ cpu_set_t *mask = NULL;
+ PyObject *res = NULL;
- if (!PyArg_ParseTuple(args, _Py_PARSE_PID "i:sched_getaffinity",
- &pid, &ncpus))
+ if (!PyArg_ParseTuple(args, _Py_PARSE_PID ":sched_getaffinity",
+ &pid))
return NULL;
- res = make_new_cpu_set(&cpu_set_type, ncpus);
- if (!res)
- return NULL;
- if (sched_getaffinity(pid, res->size, res->set)) {
- Py_DECREF(res);
- return posix_error();
+
+ ncpus = NCPUS_START;
+ while (1) {
+ setsize = CPU_ALLOC_SIZE(ncpus);
+ mask = CPU_ALLOC(ncpus);
+ if (mask == NULL)
+ return PyErr_NoMemory();
+ if (sched_getaffinity(pid, setsize, mask) == 0)
+ break;
+ CPU_FREE(mask);
+ if (errno != EINVAL)
+ return posix_error();
+ if (ncpus > INT_MAX / 2) {
+ PyErr_SetString(PyExc_OverflowError, "could not allocate "
+ "a large enough CPU set");
+ return NULL;
+ }
+ ncpus = ncpus * 2;
}
- return (PyObject *)res;
+
+ res = PySet_New(NULL);
+ if (res == NULL)
+ goto error;
+ for (cpu = 0, count = CPU_COUNT_S(setsize, mask); count; cpu++) {
+ if (CPU_ISSET_S(cpu, setsize, mask)) {
+ PyObject *cpu_num = PyLong_FromLong(cpu);
+ --count;
+ if (cpu_num == NULL)
+ goto error;
+ if (PySet_Add(res, cpu_num)) {
+ Py_DECREF(cpu_num);
+ goto error;
+ }
+ Py_DECREF(cpu_num);
+ }
+ }
+ CPU_FREE(mask);
+ return res;
+
+error:
+ if (mask)
+ CPU_FREE(mask);
+ Py_XDECREF(res);
+ return NULL;
}
#endif /* HAVE_SCHED_SETAFFINITY */
@@ -11997,13 +11813,6 @@ INITFUNC(void)
Py_INCREF(PyExc_OSError);
PyModule_AddObject(m, "error", PyExc_OSError);
-#ifdef HAVE_SCHED_SETAFFINITY
- if (PyType_Ready(&cpu_set_type) < 0)
- return NULL;
- Py_INCREF(&cpu_set_type);
- PyModule_AddObject(m, "cpu_set", (PyObject *)&cpu_set_type);
-#endif
-
#ifdef HAVE_PUTENV
if (posix_putenv_garbage == NULL)
posix_putenv_garbage = PyDict_New();