Issue #16595: Add prlimit() to resource module

prlimit() is a Linux specific command that combines setrlimit, getrlimit and can set the limit
of other processes.

8 files changed, 169 insertions, 26 deletions

diff --git a/Doc/library/resource.rst b/Doc/library/resource.rst
index 1c0fa9f..d49e13c 100644
--- a/Doc/library/resource.rst
+++ b/Doc/library/resource.rst
@@ -74,6 +74,27 @@ this module for those platforms.
    ``setrlimit`` may 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:`PermissionError` when the user doesn't have ``CAP_SYS_RESOURCE`` for
+   the process.
+
+   Availability: Linux (glibc 2.13+)
+
+   .. 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/Doc/whatsnew/3.4.rst b/Doc/whatsnew/3.4.rst
index e978fcb..8e4f8e6 100644
--- a/Doc/whatsnew/3.4.rst
+++ b/Doc/whatsnew/3.4.rst
@@ -438,6 +438,12 @@ The :mod:`pprint` module now supports *compact* mode for formatting long
 sequences (:issue:`19132`).
 
 
+resource
+--------
+
+New :func:`resource.prlimit` function and Linux specific constants.
+(Contributed by Christian Heimes in :issue:`16595` and :issue:`19324`.)
+
 smtplib
 -------
 
diff --git a/Lib/test/test_resource.py b/Lib/test/test_resource.py
index 3c9c9a8..950f477 100644
--- a/Lib/test/test_resource.py
+++ b/Lib/test/test_resource.py
@@ -139,6 +139,18 @@ class ResourceTest(unittest.TestCase):
         self.assertIsInstance(resource.RLIMIT_SIGPENDING, int)
 
 
+    @unittest.skipUnless(hasattr(resource, 'prlimit'), 'no prlimit')
+    def test_prlimit(self):
+        self.assertRaises(TypeError, resource.prlimit)
+        self.assertRaises(PermissionError, 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/Misc/NEWS b/Misc/NEWS
index 3d512d7..fc88fbf 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -19,6 +19,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #16595: Add prlimit() to resource module.
+
 - Issue #19324: Expose Linux-specific constants in resource module.
 
 - Issue #17400: ipaddress should make it easy to identify rfc6598 addresses.
diff --git a/Modules/resource.c b/Modules/resource.c
index 0ae24c0..c12ce34 100644
--- a/Modules/resource.c
+++ b/Modules/resource.c
@@ -106,6 +106,44 @@ resource_getrusage(PyObject *self, PyObject *args)
     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 @@ resource_getrlimit(PyObject *self, PyObject *args)
         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 *
@@ -166,25 +196,10 @@ resource_setrlimit(PyObject *self, PyObject *args)
     curobj = PyTuple_GET_ITEM(limits, 0);
     maxobj = PyTuple_GET_ITEM(limits, 1);
 
-#if !defined(HAVE_LARGEFILE_SUPPORT)
-    rl.rlim_cur = PyLong_AsLong(curobj);
-    if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred())
-        goto error;
-    rl.rlim_max = PyLong_AsLong(maxobj);
-    if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred())
-        goto error;
-#else
-    /* The limits are probably bigger than a long */
-    rl.rlim_cur = PyLong_AsLongLong(curobj);
-    if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred())
+    if (py2rlimit(curobj, maxobj, &rl) < 0) {
         goto error;
-    rl.rlim_max = PyLong_AsLongLong(maxobj);
-    if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred())
-        goto error;
-#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,
@@ -205,6 +220,48 @@ resource_setrlimit(PyObject *self, PyObject *args)
     return NULL;
 }
 
+#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 {
+            PyErr_SetFromErrno(PyExc_OSError);
+        }
+        return NULL;
+    }
+    return rlimit2py(old_limit);
+}
+#endif /* HAVE_PRLIMIT */
+
 static PyObject *
 resource_getpagesize(PyObject *self, PyObject *unused)
 {
@@ -229,6 +286,9 @@ static struct PyMethodDef
 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 b/configure
index 83d0ee2..fb1404f 100755
--- a/configure
+++ b/configure
@@ -10601,6 +10601,35 @@ $as_echo "no" >&6; }
 
 fi
 rm -f core conftest.err conftest.$ac_objext conftest.$ac_ext
+{ $as_echo "$as_me:${as_lineno-$LINENO}: checking for prlimit" >&5
+$as_echo_n "checking for prlimit... " >&6; }
+cat confdefs.h - <<_ACEOF >conftest.$ac_ext
+/* end confdefs.h.  */
+
+#include <sys/time.h>
+#include <sys/resource.h>
+
+int
+main ()
+{
+void *x=prlimit
+  ;
+  return 0;
+}
+_ACEOF
+if ac_fn_c_try_compile "$LINENO"; then :
+
+$as_echo "#define HAVE_PRLIMIT 1" >>confdefs.h
+
+   { $as_echo "$as_me:${as_lineno-$LINENO}: result: yes" >&5
+$as_echo "yes" >&6; }
+else
+  { $as_echo "$as_me:${as_lineno-$LINENO}: result: no" >&5
+$as_echo "no" >&6; }
+
+fi
+rm -f core conftest.err conftest.$ac_objext conftest.$ac_ext
+
 # 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/configure.ac b/configure.ac
index e1d3e77..d814581 100644
--- a/configure.ac
+++ b/configure.ac
@@ -2927,6 +2927,16 @@ AC_COMPILE_IFELSE([AC_LANG_PROGRAM([[
    AC_MSG_RESULT(yes)],
   [AC_MSG_RESULT(no)
 ])
+AC_MSG_CHECKING(for prlimit)
+AC_COMPILE_IFELSE([AC_LANG_PROGRAM([[
+#include <sys/time.h>
+#include <sys/resource.h>
+    ]], [[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
index 65205e4..13979fc 100644
--- a/pyconfig.h.in
+++ b/pyconfig.h.in
@@ -627,6 +627,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 <process.h> header file. */
 #undef HAVE_PROCESS_H