[lit] Add support for setting limits to unlimited

llvm/utils/lit/lit/TestRunner.py

@@ -602,16 +602,30 @@ def executeBuiltinUlimit(cmd, shenv):
     """executeBuiltinUlimit - Change the current limits."""
     if os.name != "posix":
         raise InternalShellError(cmd, "'ulimit' not supported on this system")
+    # Import resource here after we confirm we are on a POSIX system as the
+    # module does not exist on Windows.
+    import resource
     if len(cmd.args) != 3:
         raise InternalShellError(cmd, "'ulimit' requires two arguments")
     try:
-        new_limit = int(cmd.args[2])
+        if cmd.args[2] == "unlimited":
+            new_limit = resource.RLIM_INFINITY
+        else:
+            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
+        if new_limit != resource.RLIM_INFINITY:
+            new_limit = new_limit * 1024
+        shenv.ulimit["RLIMIT_AS"] = new_limit
     elif cmd.args[1] == "-n":
         shenv.ulimit["RLIMIT_NOFILE"] = new_limit
+    elif cmd.args[1] == "-s":
+        if new_limit != resource.RLIM_INFINITY:
+            new_limit = new_limit * 1024
+        shenv.ulimit["RLIMIT_STACK"] = new_limit
+    elif cmd.args[1] == "-f":
+        shenv.ulimit["RLIMIT_FSIZE"] = new_limit
     else:
         raise InternalShellError(
             cmd, "'ulimit' does not support option: %s" % cmd.args[1]

llvm/utils/lit/lit/builtin_commands/_launch_with_limit.py

@@ -17,6 +17,10 @@ def main(argv):
                 resource.setrlimit(resource.RLIMIT_AS, limit)
             elif limit_str == "RLIMIT_NOFILE":
                 resource.setrlimit(resource.RLIMIT_NOFILE, limit)
+            elif limit_str == "RLIMIT_STACK":
+                resource.setrlimit(resource.RLIMIT_STACK, limit)
+            elif limit_str == "RLIMIT_FSIZE":
+                resource.setrlimit(resource.RLIMIT_FSIZE, limit)
     process_output = subprocess.run(command_args)
     sys.exit(process_output.returncode)
 

llvm/utils/lit/tests/Inputs/shtest-ulimit/print_limits.py

@@ -2,3 +2,5 @@
 
 print("RLIMIT_AS=" + str(resource.getrlimit(resource.RLIMIT_AS)[0]))
 print("RLIMIT_NOFILE=" + str(resource.getrlimit(resource.RLIMIT_NOFILE)[0]))
+print("RLIMIT_STACK=" + str(resource.getrlimit(resource.RLIMIT_STACK)[0]))
+print("RLIMIT_FSIZE=" + str(resource.getrlimit(resource.RLIMIT_FSIZE)[0]))

llvm/utils/lit/tests/Inputs/shtest-ulimit/ulimit_okay.txt

@@ -1,4 +1,6 @@
 # RUN: ulimit -n 50
+# RUN: ulimit -s 256
+# RUN: ulimit -f 5
 # 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/Inputs/shtest-ulimit/ulimit_unlimited.txt

@@ -0,0 +1,6 @@
+# RUN: ulimit -f 5
+# RUN: %{python} %S/print_limits.py
+# RUN: ulimit -f unlimited
+# 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

@@ -11,15 +11,25 @@
 # RUN: not %{lit} -a -v %{inputs}/shtest-ulimit --order=lexical \
 # RUN:   | FileCheck -DBASE_NOFILE_LIMIT=%{readfile:%t.nofile_limit} %s
 
-# CHECK: -- Testing: 3 tests{{.*}}
+# CHECK: -- Testing: 4 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 -n 50
+# CHECK: ulimit -s 256
+# CHECK: ulimit -f 5
 # CHECK: RLIMIT_NOFILE=50
+# CHECK: RLIMIT_STACK=262144
+# CHECK: RLIMIT_FSIZE=5
 
 # CHECK-LABEL: FAIL: shtest-ulimit :: ulimit_reset.txt ({{[^)]*}})
 # CHECK: RLIMIT_NOFILE=[[BASE_NOFILE_LIMIT]]
+
+# CHECK-LABEL: FAIL: shtest-ulimit :: ulimit_unlimited.txt ({{[^)]*}})
+# CHECK: ulimit -f 5
+# CHECK: RLIMIT_FSIZE=5
+# CHECK: ulimit -f unlimited
+# CHECK: RLIMIT_FSIZE=-1