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)