diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py index cecbae61a3d73..342770af80228 100644 --- a/llvm/utils/lit/lit/TestRunner.py +++ b/llvm/utils/lit/lit/TestRunner.py @@ -92,11 +92,12 @@ class ShellEnvironment(object): we maintain a dir stack for pushd/popd. """ - def __init__(self, cwd, env, umask=-1): + def __init__(self, cwd, env, umask=-1, ulimit={}): self.cwd = cwd self.env = dict(env) self.umask = umask self.dirStack = [] + self.ulimit = ulimit def change_dir(self, newdir): if os.path.isabs(newdir): @@ -595,6 +596,27 @@ def executeBuiltinUmask(cmd, shenv): return ShellCommandResult(cmd, "", "", 0, False) +def executeBuiltinUlimit(cmd, shenv): + """executeBuiltinUlimit - Change the current limits.""" + if os.name != "posix": + raise InternalShellError(cmd, "'ulimit' not supported on this system") + if len(cmd.args) != 3: + raise InternalShellError(cmd, "'ulimit' requires two arguments") + try: + new_limit = int(cmd.args[2]) + except ValueError as err: + raise InternalShellError(cmd, "Error: 'ulimit': %s" % str(err)) + if cmd.args[1] == "-v": + shenv.ulimit["RLIMIT_AS"] = new_limit * 1024 + elif cmd.args[1] == "-n": + shenv.ulimit["RLIMIT_NOFILE"] = new_limit + else: + raise InternalShellError( + cmd, "'ulimit' does not support option: %s" % cmd.args[1] + ) + return ShellCommandResult(cmd, "", "", 0, False) + + def executeBuiltinColon(cmd, cmd_shenv): """executeBuiltinColon - Discard arguments and exit with status 0.""" return ShellCommandResult(cmd, "", "", 0, False) @@ -749,6 +771,7 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper): "popd": executeBuiltinPopd, "pushd": executeBuiltinPushd, "rm": executeBuiltinRm, + "ulimit": executeBuiltinUlimit, "umask": executeBuiltinUmask, ":": executeBuiltinColon, } @@ -920,6 +943,19 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper): if kIsWindows: args = quote_windows_command(args) + # Handle any resource limits. We do this by launching the command with + # a wrapper that sets the necessary limits. We use a wrapper rather than + # setting the limits in process as we cannot reraise the limits back to + # their defaults without elevated permissions. + if cmd_shenv.ulimit: + executable = sys.executable + args.insert(0, sys.executable) + args.insert(1, os.path.join(builtin_commands_dir, "_launch_with_limit.py")) + for limit in cmd_shenv.ulimit: + cmd_shenv.env["LIT_INTERNAL_ULIMIT_" + limit] = str( + cmd_shenv.ulimit[limit] + ) + try: # TODO(boomanaiden154): We currently wrap the subprocess.Popen with # os.umask as the umask argument in subprocess.Popen is not diff --git a/llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py b/llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py new file mode 100644 index 0000000000000..33d2d59ff0dbe --- /dev/null +++ b/llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py @@ -0,0 +1,25 @@ +import sys +import subprocess +import resource +import os + +ULIMIT_ENV_VAR_PREFIX = "LIT_INTERNAL_ULIMIT_" + + +def main(argv): + command_args = argv[1:] + for env_var in os.environ: + if env_var.startswith(ULIMIT_ENV_VAR_PREFIX): + limit_str = env_var[len(ULIMIT_ENV_VAR_PREFIX) :] + limit_value = int(os.environ[env_var]) + limit = (limit_value, limit_value) + if limit_str == "RLIMIT_AS": + resource.setrlimit(resource.RLIMIT_AS, limit) + elif limit_str == "RLIMIT_NOFILE": + resource.setrlimit(resource.RLIMIT_NOFILE, limit) + process_output = subprocess.run(command_args) + sys.exit(process_output.returncode) + + +if __name__ == "__main__": + main(sys.argv) diff --git a/llvm/utils/lit/tests/Inputs/shtest-ulimit/lit.cfg b/llvm/utils/lit/tests/Inputs/shtest-ulimit/lit.cfg new file mode 100644 index 0000000000000..c7bdc7e7b6bc0 --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/shtest-ulimit/lit.cfg @@ -0,0 +1,8 @@ +import lit.formats + +config.name = "shtest-ulimit" +config.suffixes = [".txt"] +config.test_format = lit.formats.ShTest(execute_external=False) +config.test_source_root = None +config.test_exec_root = None +config.substitutions.append(("%{python}", '"%s"' % (sys.executable))) diff --git a/llvm/utils/lit/tests/Inputs/shtest-ulimit/print_limits.py b/llvm/utils/lit/tests/Inputs/shtest-ulimit/print_limits.py new file mode 100644 index 0000000000000..632f954fa8fde --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/shtest-ulimit/print_limits.py @@ -0,0 +1,4 @@ +import resource + +print("RLIMIT_AS=" + str(resource.getrlimit(resource.RLIMIT_AS)[0])) +print("RLIMIT_NOFILE=" + str(resource.getrlimit(resource.RLIMIT_NOFILE)[0])) diff --git a/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit-bad-arg.txt b/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit-bad-arg.txt new file mode 100644 index 0000000000000..efa22881047e9 --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit-bad-arg.txt @@ -0,0 +1 @@ +# RUN: ulimit -n diff --git a/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit_okay.txt b/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit_okay.txt new file mode 100644 index 0000000000000..ad353b5d7c459 --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit_okay.txt @@ -0,0 +1,5 @@ +# RUN: ulimit -v 1048576 +# RUN: ulimit -n 50 +# RUN: %{python} %S/print_limits.py +# Fail the test so that we can assert on the output. +# RUN: not echo return diff --git a/llvm/utils/lit/tests/shtest-ulimit.py b/llvm/utils/lit/tests/shtest-ulimit.py new file mode 100644 index 0000000000000..8d7f436dc8af2 --- /dev/null +++ b/llvm/utils/lit/tests/shtest-ulimit.py @@ -0,0 +1,18 @@ +# Check the ulimit command + +# ulimit does not work on non-POSIX platforms. +# UNSUPPORTED: system-windows + +# RUN: not %{lit} -a -v %{inputs}/shtest-ulimit | FileCheck %s + +# CHECK: -- Testing: 2 tests{{.*}} + +# CHECK-LABEL: FAIL: shtest-ulimit :: ulimit-bad-arg.txt ({{[^)]*}}) +# CHECK: ulimit -n +# CHECK: 'ulimit' requires two arguments + +# CHECK-LABEL: FAIL: shtest-ulimit :: ulimit_okay.txt ({{[^)]*}}) +# CHECK: ulimit -v 1048576 +# CHECK: ulimit -n 50 +# CHECK: RLIMIT_AS=1073741824 +# CHECK: RLIMIT_NOFILE=50