diff --git a/Doc/library/os.rst b/Doc/library/os.rst --- a/Doc/library/os.rst +++ b/Doc/library/os.rst @@ -3141,47 +3141,17 @@ operating system. Voluntarily relinquish the CPU. -.. class:: cpu_set(ncpus) - - :class:`cpu_set` represents a set of CPUs on which a process is eligible to - run. *ncpus* is the number of CPUs the set should describe. Methods on - :class:`cpu_set` allow CPUs to be add or removed. - - :class:`cpu_set` supports the AND, OR, and XOR bitwise operations. For - example, given two cpu_sets, ``one`` and ``two``, ``one | two`` returns a - :class:`cpu_set` containing the cpus enabled both in ``one`` and ``two``. - - .. method:: set(i) - - Enable CPU *i*. - - .. method:: clear(i) - - Remove CPU *i*. - - .. method:: isset(i) - - Return ``True`` if CPU *i* is enabled in the set. - - .. method:: count() - - Return the number of enabled CPUs in the set. - - .. method:: zero() - - Clear the set completely. - - .. function:: sched_setaffinity(pid, mask) - Restrict the process with PID *pid* to a set of CPUs. *mask* is a - :class:`cpu_set` instance. - - -.. function:: sched_getaffinity(pid, size) - - Return the :class:`cpu_set` the process with PID *pid* is restricted to. The - result will contain *size* CPUs. + Restrict the process with PID *pid* (or the current process if zero) to a + set of CPUs. *mask* is an iterable of integers representing the set of + CPUs to which the process should be restricted. + + +.. function:: sched_getaffinity(pid) + + Return the set of CPUs the process with PID *pid* (or the current process + if zero) is restricted to. .. _os-path: diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py --- a/Lib/test/test_posix.py +++ b/Lib/test/test_posix.py @@ -855,7 +855,7 @@ class PosixTester(unittest.TestCase): requires_sched_h = unittest.skipUnless(hasattr(posix, 'sched_yield'), "don't have scheduling support") - requires_sched_affinity = unittest.skipUnless(hasattr(posix, 'cpu_set'), + requires_sched_affinity = unittest.skipUnless(hasattr(posix, 'sched_setaffinity'), "don't have sched affinity support") @requires_sched_h @@ -936,83 +936,28 @@ class PosixTester(unittest.TestCase): self.assertLess(interval, 1.) @requires_sched_affinity - def test_sched_affinity(self): - mask = posix.sched_getaffinity(0, 1024) - self.assertGreaterEqual(mask.count(), 1) - self.assertIsInstance(mask, posix.cpu_set) - self.assertRaises(OSError, posix.sched_getaffinity, -1, 1024) - empty = posix.cpu_set(10) - posix.sched_setaffinity(0, mask) - self.assertRaises(OSError, posix.sched_setaffinity, 0, empty) - self.assertRaises(OSError, posix.sched_setaffinity, -1, mask) + def test_sched_getaffinity(self): + mask = posix.sched_getaffinity(0) + self.assertIsInstance(mask, set) + self.assertGreaterEqual(len(mask), 1) + self.assertRaises(OSError, posix.sched_getaffinity, -1) + for cpu in mask: + self.assertIsInstance(cpu, int) + self.assertGreaterEqual(cpu, 0) + self.assertLess(cpu, 1 << 32) @requires_sched_affinity - def test_cpu_set_basic(self): - s = posix.cpu_set(10) - self.assertEqual(len(s), 10) - self.assertEqual(s.count(), 0) - s.set(0) - s.set(9) - self.assertTrue(s.isset(0)) - self.assertTrue(s.isset(9)) - self.assertFalse(s.isset(5)) - self.assertEqual(s.count(), 2) - s.clear(0) - self.assertFalse(s.isset(0)) - self.assertEqual(s.count(), 1) - s.zero() - self.assertFalse(s.isset(0)) - self.assertFalse(s.isset(9)) - self.assertEqual(s.count(), 0) - self.assertRaises(ValueError, s.set, -1) - self.assertRaises(ValueError, s.set, 10) - self.assertRaises(ValueError, s.clear, -1) - self.assertRaises(ValueError, s.clear, 10) - self.assertRaises(ValueError, s.isset, -1) - self.assertRaises(ValueError, s.isset, 10) - - @requires_sched_affinity - def test_cpu_set_cmp(self): - self.assertNotEqual(posix.cpu_set(11), posix.cpu_set(12)) - l = posix.cpu_set(10) - r = posix.cpu_set(10) - self.assertEqual(l, r) - l.set(1) - self.assertNotEqual(l, r) - r.set(1) - self.assertEqual(l, r) - - @requires_sched_affinity - def test_cpu_set_bitwise(self): - l = posix.cpu_set(5) - l.set(0) - l.set(1) - r = posix.cpu_set(5) - r.set(1) - r.set(2) - b = l & r - self.assertEqual(b.count(), 1) - self.assertTrue(b.isset(1)) - b = l | r - self.assertEqual(b.count(), 3) - self.assertTrue(b.isset(0)) - self.assertTrue(b.isset(1)) - self.assertTrue(b.isset(2)) - b = l ^ r - self.assertEqual(b.count(), 2) - self.assertTrue(b.isset(0)) - self.assertFalse(b.isset(1)) - self.assertTrue(b.isset(2)) - b = l - b |= r - self.assertIs(b, l) - self.assertEqual(l.count(), 3) - - @requires_sched_affinity - @support.cpython_only - def test_cpu_set_sizeof(self): - self.assertGreater(sys.getsizeof(posix.cpu_set(1000)), - sys.getsizeof(posix.cpu_set(1))) + def test_sched_setaffinity(self): + mask = posix.sched_getaffinity(0) + if len(mask) > 1: + # Empty masks are forbidden + mask.pop() + posix.sched_setaffinity(0, mask) + self.assertEqual(posix.sched_getaffinity(0), mask) + self.assertRaises(OSError, posix.sched_setaffinity, 0, []) + self.assertRaises(ValueError, posix.sched_setaffinity, 0, [-10]) + self.assertRaises(OverflowError, posix.sched_setaffinity, 0, [1<<128]) + self.assertRaises(OSError, posix.sched_setaffinity, -1, mask) def test_rtld_constants(self): # check presence of major RTLD_* constants diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -5720,306 +5720,6 @@ posix_sched_yield(PyObject *self, PyObje #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; - - 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); -} - -static PyObject * -cpu_set_repr(Py_cpu_set *set) -{ - return PyUnicode_FromFormat("", 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) -{ - 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."); - -static PyObject * -cpu_set_set(Py_cpu_set *set, PyObject *args) -{ - int cpu = _get_cpu(set, "i|set", args); - if (cpu == -1) - 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) - 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; -} - -PyDoc_STRVAR(cpu_set_zero_doc, -"cpu_set.zero()\n\n\ -Clear the cpu_set."); - -static PyObject * -cpu_set_zero(Py_cpu_set *set, PyObject *noargs) -{ - CPU_ZERO_S(set->size, set->set); - 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*."); @@ -6028,14 +5728,89 @@ 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(); + int ncpus; + size_t setsize; + cpu_set_t *mask = NULL; + PyObject *iterable, *iterator = NULL, *item; + + if (!PyArg_ParseTuple(args, _Py_PARSE_PID "O:sched_setaffinity", + &pid, &iterable)) + return NULL; + + iterator = PyObject_GetIter(iterable); + if (iterator == NULL) + return NULL; + + ncpus = 16; + setsize = CPU_ALLOC_SIZE(ncpus); + mask = CPU_ALLOC(ncpus); + if (mask == NULL) { + PyErr_NoMemory(); + goto error; + } + CPU_ZERO_S(setsize, mask); + + 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); + + if (sched_setaffinity(pid, setsize, mask)) { + posix_error(); + goto error; + } + CPU_FREE(mask); Py_RETURN_NONE; + +error: + if (mask) + CPU_FREE(mask); + Py_XDECREF(iterator); + return NULL; } PyDoc_STRVAR(posix_sched_getaffinity__doc__, @@ -6047,20 +5822,58 @@ static PyObject * posix_sched_getaffinity(PyObject *self, PyObject *args) { pid_t pid; - int ncpus; - Py_cpu_set *res; - - if (!PyArg_ParseTuple(args, _Py_PARSE_PID "i:sched_getaffinity", - &pid, &ncpus)) - 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(); - } - return (PyObject *)res; + int cpu, ncpus, count; + size_t setsize; + cpu_set_t *mask = NULL; + PyObject *res = NULL; + + if (!PyArg_ParseTuple(args, _Py_PARSE_PID ":sched_getaffinity", + &pid)) + return NULL; + + ncpus = 16; + 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; + } + + 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 +11810,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();