diff --git a/Doc/library/resource.rst b/Doc/library/resource.rst --- a/Doc/library/resource.rst +++ b/Doc/library/resource.rst @@ -61,6 +61,27 @@ (unless the process has an effective UID of super-user). Can also raise :exc:`error` if the underlying system call fails. +.. function:: prlimit(pid, resource[, limits]) + + Combines :func:`setrlimit` and :func:`getrlimit` in one function and + supports to get and set the resources limits of an arbitrary process. If + *pid* is 0, then the call applies to the current process. *resource* and + *limits* have the same meaning as in :func:`setrlimit`, except that + *limits* is optional. + + When *limits* is not given the function returns the *resource* limit of the + process *pid*. When *limits* is given the *resource* limit of the process is + set and the former resource limit is returned. + + Raises :exc:`ProcessLookupError` when *pid* can't be found and + :exc:`ValueError` when the user doesn't have ``CAP_SYS_RESOURCE`` for the + process. + + Availability: Linux + + .. versionadded:: 3.4 + + These symbols define resources whose consumption can be controlled using the :func:`setrlimit` and :func:`getrlimit` functions described below. The values of these symbols are exactly the constants used by C programs. diff --git a/Lib/test/test_resource.py b/Lib/test/test_resource.py --- a/Lib/test/test_resource.py +++ b/Lib/test/test_resource.py @@ -107,6 +107,18 @@ except (ValueError, AttributeError): pass + @unittest.skipUnless(hasattr(resource, 'prlimit'), 'no prlimit') + def test_prlimit(self): + self.assertRaises(TypeError, resource.prlimit) + self.assertRaises(ValueError, resource.prlimit, + 1, resource.RLIMIT_AS) + self.assertRaises(ProcessLookupError, resource.prlimit, + - 1, resource.RLIMIT_AS) + self.assertEqual(resource.prlimit(0, resource.RLIMIT_AS), (-1, -1)) + self.assertEqual(resource.prlimit(0, resource.RLIMIT_AS, (-1, -1)), + (-1, -1)) + + def test_main(verbose=None): support.run_unittest(ResourceTest) diff --git a/Modules/resource.c b/Modules/resource.c --- a/Modules/resource.c +++ b/Modules/resource.c @@ -106,6 +106,44 @@ return result; } +static int +py2rlimit(PyObject *curobj, PyObject *maxobj, struct rlimit *rl_out) +{ +#if !defined(HAVE_LARGEFILE_SUPPORT) + rl_out->rlim_cur = PyLong_AsLong(curobj); + if (rl_out->rlim_cur == (rlim_t)-1 && PyErr_Occurred()) + return -1; + rl_out->rlim_max = PyLong_AsLong(maxobj); + if (rl_out->rlim_max == (rlim_t)-1 && PyErr_Occurred()) + return -1; +#else + /* The limits are probably bigger than a long */ + rl_out->rlim_cur = PyLong_AsLongLong(curobj); + if (rl_out->rlim_cur == (rlim_t)-1 && PyErr_Occurred()) + return -1; + rl_out->rlim_max = PyLong_AsLongLong(maxobj); + if (rl_out->rlim_max == (rlim_t)-1 && PyErr_Occurred()) + return -1; +#endif + + rl_out->rlim_cur = rl_out->rlim_cur & RLIM_INFINITY; + rl_out->rlim_max = rl_out->rlim_max & RLIM_INFINITY; + return 0; + +} + +static PyObject* +rlimit2py(struct rlimit rl) +{ +#if defined(HAVE_LONG_LONG) + if (sizeof(rl.rlim_cur) > sizeof(long)) { + return Py_BuildValue("LL", + (PY_LONG_LONG) rl.rlim_cur, + (PY_LONG_LONG) rl.rlim_max); + } +#endif + return Py_BuildValue("ll", (long) rl.rlim_cur, (long) rl.rlim_max); +} static PyObject * resource_getrlimit(PyObject *self, PyObject *args) @@ -126,15 +164,7 @@ PyErr_SetFromErrno(PyExc_OSError); return NULL; } - -#if defined(HAVE_LONG_LONG) - if (sizeof(rl.rlim_cur) > sizeof(long)) { - return Py_BuildValue("LL", - (PY_LONG_LONG) rl.rlim_cur, - (PY_LONG_LONG) rl.rlim_max); - } -#endif - return Py_BuildValue("ll", (long) rl.rlim_cur, (long) rl.rlim_max); + return rlimit2py(rl); } static PyObject * @@ -154,25 +184,10 @@ return NULL; } -#if !defined(HAVE_LARGEFILE_SUPPORT) - rl.rlim_cur = PyLong_AsLong(curobj); - if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred()) + if (py2rlimit(curobj, maxobj, &rl) < 0) { return NULL; - rl.rlim_max = PyLong_AsLong(maxobj); - if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred()) - return NULL; -#else - /* The limits are probably bigger than a long */ - rl.rlim_cur = PyLong_AsLongLong(curobj); - if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred()) - return NULL; - rl.rlim_max = PyLong_AsLongLong(maxobj); - if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred()) - return NULL; -#endif + } - rl.rlim_cur = rl.rlim_cur & RLIM_INFINITY; - rl.rlim_max = rl.rlim_max & RLIM_INFINITY; if (setrlimit(resource, &rl) == -1) { if (errno == EINVAL) PyErr_SetString(PyExc_ValueError, @@ -188,6 +203,50 @@ return Py_None; } +#ifdef HAVE_PRLIMIT +static PyObject * +resource_prlimit(PyObject *self, PyObject *args) +{ + struct rlimit old_limit, new_limit; + int resource, retval; + pid_t pid; + PyObject *curobj=NULL, *maxobj=NULL; + + if (!PyArg_ParseTuple(args, _Py_PARSE_PID "i|(OO):prlimit", + &pid, &resource, &curobj, &maxobj)) + return NULL; + + if (resource < 0 || resource >= RLIM_NLIMITS) { + PyErr_SetString(PyExc_ValueError, + "invalid resource specified"); + return NULL; + } + + if (curobj != NULL) { + if (py2rlimit(curobj, maxobj, &new_limit) < 0) { + return NULL; + } + retval = prlimit(pid, resource, &new_limit, &old_limit); + } + else { + retval = prlimit(pid, resource, NULL, &old_limit); + } + + if (retval == -1) { + if (errno == EINVAL) + PyErr_SetString(PyExc_ValueError, + "current limit exceeds maximum limit"); + else if (errno == EPERM) + PyErr_SetString(PyExc_ValueError, + "not allowed to raise maximum limit"); + else + PyErr_SetFromErrno(PyExc_OSError); + return NULL; + } + return rlimit2py(old_limit); +} +#endif /* HAVE_PRLIMIT */ + static PyObject * resource_getpagesize(PyObject *self, PyObject *unused) { @@ -212,6 +271,9 @@ resource_methods[] = { {"getrusage", resource_getrusage, METH_VARARGS}, {"getrlimit", resource_getrlimit, METH_VARARGS}, +#ifdef HAVE_PRLIMIT + {"prlimit", resource_prlimit, METH_VARARGS}, +#endif {"setrlimit", resource_setrlimit, METH_VARARGS}, {"getpagesize", resource_getpagesize, METH_NOARGS}, {NULL, NULL} /* sentinel */ diff --git a/configure.ac b/configure.ac --- a/configure.ac +++ b/configure.ac @@ -2848,6 +2848,16 @@ AC_MSG_RESULT(yes)], [AC_MSG_RESULT(no) ]) +AC_MSG_CHECKING(for prlimit) +AC_COMPILE_IFELSE([AC_LANG_PROGRAM([[ +#include +#include + ]], [[void *x=prlimit]])], + [AC_DEFINE(HAVE_PRLIMIT, 1, Define if you have the 'prlimit' functions.) + AC_MSG_RESULT(yes)], + [AC_MSG_RESULT(no) +]) + # On some systems (eg. FreeBSD 5), we would find a definition of the # functions ctermid_r, setgroups in the library, but no prototype # (e.g. because we use _XOPEN_SOURCE). See whether we can take their diff --git a/pyconfig.h.in b/pyconfig.h.in --- a/pyconfig.h.in +++ b/pyconfig.h.in @@ -618,6 +618,9 @@ /* Define to 1 if you have the `pread' function. */ #undef HAVE_PREAD +/* Define if you have the 'prlimit' functions. */ +#undef HAVE_PRLIMIT + /* Define to 1 if you have the header file. */ #undef HAVE_PROCESS_H