Merge "processutils: add support for missing process limits"3.8.0

3 files changed, 85 insertions, 11 deletions

diff --git a/oslo_concurrency/prlimit.py b/oslo_concurrency/prlimit.py
index f7718de..e0fc4e3 100644
--- a/oslo_concurrency/prlimit.py
+++ b/oslo_concurrency/prlimit.py
@@ -26,8 +26,15 @@ USAGE_PROGRAM = ('%s -m oslo_concurrency.prlimit'
 RESOURCES = (
     # argparse argument => resource
     ('as', resource.RLIMIT_AS),
+    ('core', resource.RLIMIT_CORE),
+    ('cpu', resource.RLIMIT_CPU),
+    ('data', resource.RLIMIT_DATA),
+    ('fsize', resource.RLIMIT_FSIZE),
+    ('memlock', resource.RLIMIT_MEMLOCK),
     ('nofile', resource.RLIMIT_NOFILE),
+    ('nproc', resource.RLIMIT_NPROC),
     ('rss', resource.RLIMIT_RSS),
+    ('stack', resource.RLIMIT_STACK),
 )
 
 
@@ -35,10 +42,24 @@ def parse_args():
     parser = argparse.ArgumentParser(description='prlimit', prog=USAGE_PROGRAM)
     parser.add_argument('--as', type=int,
                         help='Address space limit in bytes')
+    parser.add_argument('--core', type=int,
+                        help='Core file size limit in bytes')
+    parser.add_argument('--cpu', type=int,
+                        help='CPU time limit in seconds')
+    parser.add_argument('--data', type=int,
+                        help='Data size limit in bytes')
+    parser.add_argument('--fsize', type=int,
+                        help='File size limit in bytes')
+    parser.add_argument('--memlock', type=int,
+                        help='Locked memory limit in bytes')
     parser.add_argument('--nofile', type=int,
                         help='Maximum number of open files')
+    parser.add_argument('--nproc', type=int,
+                        help='Maximum number of processes')
     parser.add_argument('--rss', type=int,
                         help='Maximum Resident Set Size (RSS) in bytes')
+    parser.add_argument('--stack', type=int,
+                        help='Stack size limit in bytes')
     parser.add_argument('program',
                         help='Program (absolute path)')
     parser.add_argument('program_args', metavar="arg", nargs='...',
diff --git a/oslo_concurrency/processutils.py b/oslo_concurrency/processutils.py
index b74e5f2..1f46221 100644
--- a/oslo_concurrency/processutils.py
+++ b/oslo_concurrency/processutils.py
@@ -134,16 +134,36 @@ class ProcessLimits(object):
     Attributes:
 
     * address_space: Address space limit in bytes
-    * number_files: Maximum number of open files.
+    * core_file_size: Core file size limit in bytes
+    * cpu_time: CPU time limit in seconds
+    * data_size: Data size limit in bytes
+    * file_size: File size limit in bytes
+    * memory_locked: Locked memory limit in bytes
+    * number_files: Maximum number of open files
+    * number_processes: Maximum number of processes
     * resident_set_size: Maximum Resident Set Size (RSS) in bytes
+    * stack_size: Stack size limit in bytes
 
     This object can be used for the *prlimit* parameter of :func:`execute`.
     """
 
+    _LIMITS = {
+        "address_space": "--as",
+        "core_file_size": "--core",
+        "cpu_time": "--cpu",
+        "data_size": "--data",
+        "file_size": "--fsize",
+        "memory_locked": "--memlock",
+        "number_files": "--nofile",
+        "number_processes": "--nproc",
+        "resident_set_size": "--rss",
+        "stack_size": "--stack",
+    }
+
     def __init__(self, **kw):
-        self.address_space = kw.pop('address_space', None)
-        self.number_files = kw.pop('number_files', None)
-        self.resident_set_size = kw.pop('resident_set_size', None)
+        for limit in self._LIMITS.keys():
+            setattr(self, limit, kw.pop(limit, None))
+
         if kw:
             raise ValueError("invalid limits: %s"
                              % ', '.join(sorted(kw.keys())))
@@ -151,12 +171,10 @@ class ProcessLimits(object):
     def prlimit_args(self):
         """Create a list of arguments for the prlimit command line."""
         args = []
-        if self.address_space:
-            args.append('--as=%s' % self.address_space)
-        if self.number_files:
-            args.append('--nofile=%s' % self.number_files)
-        if self.resident_set_size:
-            args.append('--rss=%s' % self.resident_set_size)
+        for limit in self._LIMITS.keys():
+            val = getattr(self, limit)
+            if val is not None:
+                args.append("%s=%s" % (self._LIMITS[limit], val))
         return args
 
 
diff --git a/oslo_concurrency/tests/unit/test_processutils.py b/oslo_concurrency/tests/unit/test_processutils.py
index 0fd8045..5c744f5 100644
--- a/oslo_concurrency/tests/unit/test_processutils.py
+++ b/oslo_concurrency/tests/unit/test_processutils.py
@@ -752,7 +752,7 @@ class PrlimitTestCase(test_base.BaseTestCase):
         # Create a new soft limit for a resource, lower than the current
         # soft limit.
         soft_limit, hard_limit = resource.getrlimit(res)
-        if soft_limit < 0:
+        if soft_limit <= 0:
             soft_limit = default_limit
         else:
             soft_limit -= substract
@@ -790,6 +790,31 @@ class PrlimitTestCase(test_base.BaseTestCase):
         prlimit = self.limit_address_space()
         self.check_limit(prlimit, 'RLIMIT_AS', prlimit.address_space)
 
+    def test_core_size(self):
+        size = self.soft_limit(resource.RLIMIT_CORE, 1, 1024)
+        prlimit = processutils.ProcessLimits(core_file_size=size)
+        self.check_limit(prlimit, 'RLIMIT_CORE', prlimit.core_file_size)
+
+    def test_cpu_time(self):
+        time = self.soft_limit(resource.RLIMIT_CPU, 1, 1024)
+        prlimit = processutils.ProcessLimits(cpu_time=time)
+        self.check_limit(prlimit, 'RLIMIT_CPU', prlimit.cpu_time)
+
+    def test_data_size(self):
+        max_memory = self.memory_limit(resource.RLIMIT_DATA)
+        prlimit = processutils.ProcessLimits(data_size=max_memory)
+        self.check_limit(prlimit, 'RLIMIT_DATA', max_memory)
+
+    def test_file_size(self):
+        size = self.soft_limit(resource.RLIMIT_FSIZE, 1, 1024)
+        prlimit = processutils.ProcessLimits(file_size=size)
+        self.check_limit(prlimit, 'RLIMIT_FSIZE', prlimit.file_size)
+
+    def test_memory_locked(self):
+        max_memory = self.memory_limit(resource.RLIMIT_MEMLOCK)
+        prlimit = processutils.ProcessLimits(memory_locked=max_memory)
+        self.check_limit(prlimit, 'RLIMIT_MEMLOCK', max_memory)
+
     def test_resident_set_size(self):
         max_memory = self.memory_limit(resource.RLIMIT_RSS)
         prlimit = processutils.ProcessLimits(resident_set_size=max_memory)
@@ -800,6 +825,16 @@ class PrlimitTestCase(test_base.BaseTestCase):
         prlimit = processutils.ProcessLimits(number_files=nfiles)
         self.check_limit(prlimit, 'RLIMIT_NOFILE', nfiles)
 
+    def test_number_processes(self):
+        nprocs = self.soft_limit(resource.RLIMIT_NPROC, 1, 65535)
+        prlimit = processutils.ProcessLimits(number_processes=nprocs)
+        self.check_limit(prlimit, 'RLIMIT_NPROC', nprocs)
+
+    def test_stack_size(self):
+        max_memory = self.memory_limit(resource.RLIMIT_STACK)
+        prlimit = processutils.ProcessLimits(stack_size=max_memory)
+        self.check_limit(prlimit, 'RLIMIT_STACK', max_memory)
+
     def test_unsupported_prlimit(self):
         self.assertRaises(ValueError, processutils.ProcessLimits, xxx=33)