[Clang] Port ulimit tests to work with internal shell

clang/test/PCH/leakfiles.test

@@ -2,7 +2,6 @@
 // https://bugs.chromium.org/p/chromium/issues/detail?id=924225
 //
 // This test requires bash loops and ulimit.
-// REQUIRES: shell
 // UNSUPPORTED: target={{.*win32.*}}
 //
 // Set up source files. lib/lib.h includes lots of lib*.h files in that dir.
@@ -12,10 +11,10 @@
 // RUN: mkdir %t
 // RUN: cd %t
 // RUN: mkdir lib
-// RUN: for i in {1..300}; do touch lib/lib$i.h; done
-// RUN: for i in {1..300}; do echo "#include \"lib$i.h\"" >> lib/lib.h; done
+// RUN: %python -c "from pathlib import Path; list(map(lambda i: Path(f'lib/lib{i}.h').touch(), range(1, 301)))"
+// RUN: %python -c "for i in range(1, 301): print(f'#include \"lib{i}.h\"')" > lib/lib.h
 // RUN: echo "#include \"lib/lib.h\"" > client.c
-// RUN: for i in {1..300}; do echo "#include \"lib/lib$i.h\"" >> client.c; done
+// RUN: %python -c "for i in range(1, 301): print(f'#include \"lib/lib{i}.h\"')" > client.c
 //
 // We want to verify that we don't hold all the files open at the same time.
 // This is important e.g. on mac, which has a low default FD limit.

clang/test/SemaCXX/PR51712-large-array-constexpr-check-oom.cpp

@@ -1,7 +1,6 @@
 // Only run this test where ulimit is known to work well.
 // (There's nothing really platform-specific being tested, this is just ulimit).
 //
-// REQUIRES: shell
 // REQUIRES: system-linux
 // UNSUPPORTED: msan
 // UNSUPPORTED: asan

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

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)

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)))

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]))

llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit-bad-arg.txt

@@ -0,0 +1 @@
+# RUN: ulimit -n

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

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