Merge "Add prlimit parameter to execute()"3.4.0

3 files changed, 250 insertions, 0 deletions

diff --git a/oslo_concurrency/prlimit.py b/oslo_concurrency/prlimit.py
new file mode 100644
index 0000000..f7718de
--- /dev/null
+++ b/oslo_concurrency/prlimit.py
@@ -0,0 +1,89 @@
+# Copyright 2016 Red Hat.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+from __future__ import print_function
+
+import argparse
+import os
+import resource
+import sys
+
+USAGE_PROGRAM = ('%s -m oslo_concurrency.prlimit'
+                 % os.path.basename(sys.executable))
+
+RESOURCES = (
+    # argparse argument => resource
+    ('as', resource.RLIMIT_AS),
+    ('nofile', resource.RLIMIT_NOFILE),
+    ('rss', resource.RLIMIT_RSS),
+)
+
+
+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('--nofile', type=int,
+                        help='Maximum number of open files')
+    parser.add_argument('--rss', type=int,
+                        help='Maximum Resident Set Size (RSS) in bytes')
+    parser.add_argument('program',
+                        help='Program (absolute path)')
+    parser.add_argument('program_args', metavar="arg", nargs='...',
+                        help='Program parameters')
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    program = args.program
+    if not os.path.isabs(program):
+        # program uses a relative path: try to find the absolute path
+        # to the executable
+        if sys.version_info >= (3, 3):
+            import shutil
+            program_abs = shutil.which(program)
+        else:
+            import distutils.spawn
+            program_abs = distutils.spawn.find_executable(program)
+        if program_abs:
+            program = program_abs
+
+    for arg_name, rlimit in RESOURCES:
+        value = getattr(args, arg_name)
+        if value is None:
+            continue
+        try:
+            resource.setrlimit(rlimit, (value, value))
+        except ValueError as exc:
+            print("%s: failed to set the %s resource limit: %s"
+                  % (USAGE_PROGRAM, arg_name.upper(), exc),
+                  file=sys.stderr)
+            sys.exit(1)
+
+    try:
+        os.execv(program, [program] + args.program_args)
+    except Exception as exc:
+        print("%s: failed to execute %s: %s"
+              % (USAGE_PROGRAM, program, exc),
+              file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/oslo_concurrency/processutils.py b/oslo_concurrency/processutils.py
index a2f8896..436922a 100644
--- a/oslo_concurrency/processutils.py
+++ b/oslo_concurrency/processutils.py
@@ -24,6 +24,7 @@ import os
 import random
 import shlex
 import signal
+import sys
 import time
 
 import enum
@@ -120,6 +121,38 @@ LOG_FINAL_ERROR = LogErrors.FINAL
 LOG_DEFAULT_ERROR = LogErrors.DEFAULT
 
 
+class ProcessLimits(object):
+    """Resource limits on a process.
+
+    Attributes:
+
+    * address_space: Address space limit in bytes
+    * number_files: Maximum number of open files.
+    * resident_set_size: Maximum Resident Set Size (RSS) in bytes
+
+    This object can be used for the *prlimit* parameter of :func:`execute`.
+    """
+
+    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)
+        if kw:
+            raise ValueError("invalid limits: %s"
+                             % ', '.join(sorted(kw.keys())))
+
+    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)
+        return args
+
+
 def execute(*cmd, **kwargs):
     """Helper method to shell out and execute a command through subprocess.
 
@@ -188,12 +221,22 @@ def execute(*cmd, **kwargs):
                             subprocess.Popen on windows (throws a
                             ValueError)
     :type preexec_fn:       function()
+    :param prlimit:         Set resource limits on the child process. See
+                            below for a detailed description.
+    :type prlimit:          :class:`ProcessLimits`
     :returns:               (stdout, stderr) from process execution
     :raises:                :class:`UnknownArgumentError` on
                             receiving unknown arguments
     :raises:                :class:`ProcessExecutionError`
     :raises:                :class:`OSError`
 
+    The *prlimit* parameter can be used to set resource limits on the child
+    process.  If this parameter is used, the child process will be spawned by a
+    wrapper process which will set limits before spawning the command.
+
+    .. versionchanged:: 3.4
+       Added *prlimit* optional parameter.
+
     .. versionchanged:: 1.5
        Added *cwd* optional parameter.
 
@@ -227,6 +270,7 @@ def execute(*cmd, **kwargs):
     on_execute = kwargs.pop('on_execute', None)
     on_completion = kwargs.pop('on_completion', None)
     preexec_fn = kwargs.pop('preexec_fn', None)
+    prlimit = kwargs.pop('prlimit', None)
 
     if isinstance(check_exit_code, bool):
         ignore_exit_code = not check_exit_code
@@ -256,6 +300,14 @@ def execute(*cmd, **kwargs):
             cmd = shlex.split(root_helper) + list(cmd)
 
     cmd = [str(c) for c in cmd]
+
+    if prlimit:
+        args = [sys.executable, '-m', 'oslo_concurrency.prlimit']
+        args.extend(prlimit.prlimit_args())
+        args.append('--')
+        args.extend(cmd)
+        cmd = args
+
     sanitized_cmd = strutils.mask_password(' '.join(cmd))
 
     watch = timeutils.StopWatch()
diff --git a/oslo_concurrency/tests/unit/test_processutils.py b/oslo_concurrency/tests/unit/test_processutils.py
index ae8d4b3..b868cb5 100644
--- a/oslo_concurrency/tests/unit/test_processutils.py
+++ b/oslo_concurrency/tests/unit/test_processutils.py
@@ -19,6 +19,7 @@ import errno
 import logging
 import multiprocessing
 import os
+import resource
 import stat
 import subprocess
 import sys
@@ -724,3 +725,111 @@ class SshExecuteTestCase(test_base.BaseTestCase):
 
     def test_compromising_ssh6(self):
         self._test_compromising_ssh(rc=-1, check=False)
+
+
+class PrlimitTestCase(test_base.BaseTestCase):
+    # Simply program that does nothing and returns an exit code 0.
+    # Use Python to be portable.
+    SIMPLE_PROGRAM = [sys.executable, '-c', 'pass']
+
+    def soft_limit(self, res, substract, default_limit):
+        # 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:
+            soft_limit = default_limit
+        else:
+            soft_limit -= substract
+        return soft_limit
+
+    def memory_limit(self, res):
+        # Substract 1 kB just to get a different limit. Don't substract too
+        # much to avoid memory allocation issues.
+        #
+        # Use 1 GB by default. Limit high enough to be able to load shared
+        # libraries. Limit low enough to be work on 32-bit platforms.
+        return self.soft_limit(res, 1024, 1024 ** 3)
+
+    def limit_address_space(self):
+        max_memory = self.memory_limit(resource.RLIMIT_AS)
+        return processutils.ProcessLimits(address_space=max_memory)
+
+    def test_simple(self):
+        # Simple test running a program (/bin/true) with no parameter
+        prlimit = self.limit_address_space()
+        stdout, stderr = processutils.execute(*self.SIMPLE_PROGRAM,
+                                              prlimit=prlimit)
+        self.assertEqual(stdout.rstrip(), '')
+        self.assertEqual(stderr.rstrip(), '')
+
+    def check_limit(self, prlimit, resource, value):
+        code = ';'.join(('import resource',
+                         'print(resource.getrlimit(resource.%s))' % resource))
+        args = [sys.executable, '-c', code]
+        stdout, stderr = processutils.execute(*args, prlimit=prlimit)
+        expected = (value, value)
+        self.assertEqual(stdout.rstrip(), str(expected))
+
+    def test_address_space(self):
+        prlimit = self.limit_address_space()
+        self.check_limit(prlimit, 'RLIMIT_AS', prlimit.address_space)
+
+    def test_resident_set_size(self):
+        max_memory = self.memory_limit(resource.RLIMIT_RSS)
+        prlimit = processutils.ProcessLimits(resident_set_size=max_memory)
+        self.check_limit(prlimit, 'RLIMIT_RSS', max_memory)
+
+    def test_number_files(self):
+        nfiles = self.soft_limit(resource.RLIMIT_NOFILE, 1, 1024)
+        prlimit = processutils.ProcessLimits(number_files=nfiles)
+        self.check_limit(prlimit, 'RLIMIT_NOFILE', nfiles)
+
+    def test_unsupported_prlimit(self):
+        self.assertRaises(ValueError, processutils.ProcessLimits, xxx=33)
+
+    def test_relative_path(self):
+        prlimit = self.limit_address_space()
+        program = sys.executable
+
+        env = dict(os.environ)
+        env['PATH'] = os.path.dirname(program)
+        args = [os.path.basename(program), '-c', 'pass']
+        processutils.execute(*args, prlimit=prlimit, env_variables=env)
+
+    def test_execv_error(self):
+        prlimit = self.limit_address_space()
+        args = ['/missing_path/dont_exist/program']
+        try:
+            processutils.execute(*args, prlimit=prlimit)
+        except processutils.ProcessExecutionError as exc:
+            self.assertEqual(exc.exit_code, 1)
+            self.assertEqual(exc.stdout, '')
+            expected = ('%s -m oslo_concurrency.prlimit: '
+                        'failed to execute /missing_path/dont_exist/program: '
+                        % os.path.basename(sys.executable))
+            self.assertIn(expected, exc.stderr)
+        else:
+            self.fail("ProcessExecutionError not raised")
+
+    def test_setrlimit_error(self):
+        prlimit = self.limit_address_space()
+
+        # trying to set a limit higher than the current hard limit
+        # with setrlimit() should fail.
+        higher_limit = prlimit.address_space + 1024
+
+        args = [sys.executable, '-m', 'oslo_concurrency.prlimit',
+                '--as=%s' % higher_limit,
+                '--']
+        args.extend(self.SIMPLE_PROGRAM)
+        try:
+            processutils.execute(*args, prlimit=prlimit)
+        except processutils.ProcessExecutionError as exc:
+            self.assertEqual(exc.exit_code, 1)
+            self.assertEqual(exc.stdout, '')
+            expected = ('%s -m oslo_concurrency.prlimit: '
+                        'failed to set the AS resource limit: '
+                        % os.path.basename(sys.executable))
+            self.assertIn(expected, exc.stderr)
+        else:
+            self.fail("ProcessExecutionError not raised")