traverseda
/
cpython
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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.
This commit is contained in:
Antoine Pitrou 2012-08-04 16:16:35 +02:00
parent a0e3febe5a
commit 848698727f
4 changed files with 162 additions and 434 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
Doc/library/os.rst
View File

@ -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.
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, size)
.. function:: sched_getaffinity(pid)
Return the :class:`cpu_set` the process with PID *pid* is restricted to. The
result will contain *size* CPUs.
Return the set of CPUs the process with PID *pid* (or the current process
if zero) is restricted to.
.. _os-path:

97
Lib/test/test_posix.py
View File

@ -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,84 +936,29 @@ 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)
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_sched_setaffinity(self):
mask = posix.sched_getaffinity(0)
if len(mask) > 1:
# Empty masks are forbidden
mask.pop()
posix.sched_setaffinity(0, mask)
self.assertRaises(OSError, posix.sched_setaffinity, 0, empty)
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)
@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_rtld_constants(self):
# check presence of major RTLD_* constants
posix.RTLD_LAZY

4
Misc/NEWS
View File

@ -77,6 +77,10 @@ Core and Builtins
Library
-------
- 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.
- Issue #15538: Fix compilation of the getnameinfo() / getaddrinfo()
emulation code. Patch by Philipp Hagemeister.

451
Modules/posixmodule.c
View File

@ -5720,305 +5720,8 @@ 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;
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("<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)
{
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 */
};
/* The minimum number of CPUs allocated in a cpu_set_t */
static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
PyDoc_STRVAR(posix_sched_setaffinity__doc__,
"sched_setaffinity(pid, cpu_set)\n\n\
@ -6028,14 +5731,89 @@ static PyObject *
posix_sched_setaffinity(PyObject *self, PyObject *args)
{
pid_t pid;
Py_cpu_set *cpu_set;
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, &cpu_set_type, &cpu_set))
if (!PyArg_ParseTuple(args, _Py_PARSE_PID "O:sched_setaffinity",
&pid, &iterable))
return NULL;
if (sched_setaffinity(pid, cpu_set->size, cpu_set->set))
return posix_error();
iterator = PyObject_GetIter(iterable);
if (iterator == NULL)
return NULL;
ncpus = NCPUS_START;
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 +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);
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;
}
return (PyObject *)res;
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 +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();