[lit] Implement ulimit builtin

[boomanaiden154]

This patch implements ulimit inside the lit internal shell.
Implementation wise, this functions similar to umask. But instead of
setting the limits within the lit test worker process, we set
environment variables and add a wrapper around the command to be
executed. The wrapper then sets the limits. This is because we cannot
increase the limits after lowering them, so we would otherwise end up
with a lit test worker stuck with a lower limit.

There are several tests where the use of ulimit is essential to the
semantics of the test (two in clang, ~7 in compiler-rt), so we need to
implement this in order to switch on the internal shell by default
without losing test coverage.

[llvmbot]

@llvm/pr-subscribers-testing-tools

Author: Aiden Grossman (boomanaiden154)

Changes

This patch implements ulimit inside the lit internal shell.
Implementation wise, this functions similar to umask. But instead of
setting the limits within the lit test worker process, we set
environment variables and add a wrapper around the command to be
executed. The wrapper then sets the limits. This is because we cannot
increase the limits after lowering them, so we would otherwise end up
with a lit test worker stuck with a lower limit.

There are several tests where the use of ulimit is essential to the
semantics of the test (two in clang, ~7 in compiler-rt), so we need to
implement this in order to switch on the internal shell by default
without losing test coverage.

---

Full diff: https://github.com/llvm/llvm-project/pull/157958.diff

7 Files Affected:

• (modified) llvm/utils/lit/lit/TestRunner.py (+31-1)
• (added) llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py (+24)
• (added) llvm/utils/lit/tests/Inputs/shtest-ulimit/lit.cfg (+8)
• (added) llvm/utils/lit/tests/Inputs/shtest-ulimit/print_limits.py (+4)
• (added) llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit-bad-arg.txt (+1)
• (added) llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit_okay.txt (+5)
• (added) llvm/utils/lit/tests/shtest-ulimit.py (+18)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 69ca80008e2f9..ec44c26c29afc 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -98,11 +98,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):
@@ -600,6 +601,24 @@ def executeBuiltinUmask(cmd, shenv):
         raise InternalShellError(cmd, "Error: 'umask': %s" % str(err))
     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."""
@@ -755,6 +774,7 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
         "popd": executeBuiltinPopd,
         "pushd": executeBuiltinPushd,
         "rm": executeBuiltinRm,
+        "ulimit": executeBuiltinUlimit,
         "umask": executeBuiltinUmask,
         ":": executeBuiltinColon,
     }
@@ -925,6 +945,16 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
         # with some core utility distributions.
         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:
+            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
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..0ccce8982400e
--- /dev/null
+++ b/llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py
@@ -0,0 +1,24 @@
+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

[dyung]

Hi @boomanaiden154, not sure why this PR isn't being pinged, but this change seems to be failing on my MacOS bots:

https://lab.llvm.org/buildbot/#/builders/190/builds/27332

Can you take a look?

[boomanaiden154]

Yeah, just got a couple of the emails. I'll revert this for now and take a look locally when I have a chance to get a build setup.