Add prlimit parameter to execute()

Add a new oslo_concurrency.prlimit module which is written to be used
on the command line:

    python -m oslo_concurrency.prlimit --rss=RSS -- program arg1 ...

This module calls setrlimit() to restrict the resources and then
executes the program. Its command line is written to be the same than
the Linux prlimit system program.

Add a new ProcessLimits class processutils: resource limits on a
process.

Add an optional prlimit parameter to process_utils.execute(). If the
parameter is used, wrap the command through the new oslo_concurrency
prlimit wrapper.

Linux provides a prlimit command line tool which implements the same
feature (and even more), but it requires util-linux v2.21, and
OpenStack targets other operating systems like Solaris and FreeBSD.

Change-Id: Ib40aa62958ab9c157a2bd51d7ff3edb445556285
Related-Bug: 1449062
(cherry-pick from b2e78569c5cabc9582c02aacff1ce2a5e186c3ab)

3 files changed, 247 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 36ac9b0..8c6c9f9 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
 
 from oslo_utils import importutils
@@ -100,6 +101,38 @@ LOG_ALL_ERRORS = 1
 LOG_FINAL_ERROR = 2
 
 
+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.
 
@@ -169,11 +202,18 @@ 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.
     """
 
     cwd = kwargs.pop('cwd', None)
@@ -192,6 +232,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
@@ -219,6 +260,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 5ec5a9b..5186d9a 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
@@ -721,3 +722,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")