@@ -92,11 +92,12 @@ class ShellEnvironment(object):
9292 we maintain a dir stack for pushd/popd.
9393 """
9494
95- def __init__ (self , cwd , env , umask = - 1 ):
95+ def __init__ (self , cwd , env , umask = - 1 , ulimit = {} ):
9696 self .cwd = cwd
9797 self .env = dict (env )
9898 self .umask = umask
9999 self .dirStack = []
100+ self .ulimit = ulimit
100101
101102 def change_dir (self , newdir ):
102103 if os .path .isabs (newdir ):
@@ -595,6 +596,27 @@ def executeBuiltinUmask(cmd, shenv):
595596 return ShellCommandResult (cmd , "" , "" , 0 , False )
596597
597598
599+ def executeBuiltinUlimit (cmd , shenv ):
600+ """executeBuiltinUlimit - Change the current limits."""
601+ if os .name != "posix" :
602+ raise InternalShellError (cmd , "'ulimit' not supported on this system" )
603+ if len (cmd .args ) != 3 :
604+ raise InternalShellError (cmd , "'ulimit' requires two arguments" )
605+ try :
606+ new_limit = int (cmd .args [2 ])
607+ except ValueError as err :
608+ raise InternalShellError (cmd , "Error: 'ulimit': %s" % str (err ))
609+ if cmd .args [1 ] == "-v" :
610+ shenv .ulimit ["RLIMIT_AS" ] = new_limit * 1024
611+ elif cmd .args [1 ] == "-n" :
612+ shenv .ulimit ["RLIMIT_NOFILE" ] = new_limit
613+ else :
614+ raise InternalShellError (
615+ cmd , "'ulimit' does not support option: %s" % cmd .args [1 ]
616+ )
617+ return ShellCommandResult (cmd , "" , "" , 0 , False )
618+
619+
598620def executeBuiltinColon (cmd , cmd_shenv ):
599621 """executeBuiltinColon - Discard arguments and exit with status 0."""
600622 return ShellCommandResult (cmd , "" , "" , 0 , False )
@@ -749,6 +771,7 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
749771 "popd" : executeBuiltinPopd ,
750772 "pushd" : executeBuiltinPushd ,
751773 "rm" : executeBuiltinRm ,
774+ "ulimit" : executeBuiltinUlimit ,
752775 "umask" : executeBuiltinUmask ,
753776 ":" : executeBuiltinColon ,
754777 }
@@ -914,6 +937,19 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
914937 if kIsWindows :
915938 args = quote_windows_command (args )
916939
940+ # Handle any resource limits. We do this by launching the command with
941+ # a wrapper that sets the necessary limits. We use a wrapper rather than
942+ # setting the limits in process as we cannot reraise the limits back to
943+ # their defaults without elevated permissions.
944+ if cmd_shenv .ulimit :
945+ executable = sys .executable
946+ args .insert (0 , sys .executable )
947+ args .insert (1 , os .path .join (builtin_commands_dir , "_launch_with_limit.py" ))
948+ for limit in cmd_shenv .ulimit :
949+ cmd_shenv .env ["LIT_INTERNAL_ULIMIT_" + limit ] = str (
950+ cmd_shenv .ulimit [limit ]
951+ )
952+
917953 try :
918954 # TODO(boomanaiden154): We currently wrap the subprocess.Popen with
919955 # os.umask as the umask argument in subprocess.Popen is not
0 commit comments