tricu

An interpreted language for exploring Tree Calculus
Log | Files | Refs | README | LICENSE

run-bounded.py (2715B)

      1 #!/usr/bin/env python3
      2 """Run site checks with per-process data limits, aggregate RSS guard and deadline.
      3 
      4 Use inside a Nix tooling shell. Chromium reserves a huge PROT_NONE address space,
      5 so bound writable data rather than virtual address reservations. RLIMIT_DATA also
      6 covers anonymous writable mmap on Linux. The aggregate /proc guard covers the
      7 whole process group, including Chromium and Workers; it is sampled, not a cgroup.
      8 """
      9 import argparse
     10 import os
     11 from pathlib import Path
     12 import resource
     13 import signal
     14 import subprocess
     15 import time
     16 
     17 
     18 def main():
     19     p = argparse.ArgumentParser(description=__doc__)
     20     p.add_argument('--seconds', type=int, default=300)
     21     p.add_argument('--memory-mib', type=int, default=4096)
     22     p.add_argument('command', nargs=argparse.REMAINDER)
     23     args = p.parse_args()
     24     command = args.command[1:] if args.command[:1] == ['--'] else args.command
     25     if not command or args.seconds <= 0 or args.memory_mib <= 0:
     26         p.error('positive limits and a command are required')
     27     limit = args.memory_mib * 1024 * 1024
     28 
     29     def limits():
     30         resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
     31         resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
     32 
     33     child = subprocess.Popen(command, start_new_session=True, preexec_fn=limits)
     34     started = time.monotonic()
     35     peak = 0
     36     reason = None
     37     try:
     38         while child.poll() is None:
     39             resident = 0
     40             for path in Path('/proc').glob('[0-9]*/status'):
     41                 try:
     42                     fields = dict(line.split(':', 1) for line in path.read_text().splitlines() if ':' in line)
     43                     # /proc can be mounted from the outer PID namespace. The
     44                     # last NSpgid is the ID seen by this sandbox's subprocess.
     45                     if int(fields.get('NSpgid', '-1').split()[-1]) == child.pid:
     46                         resident += int(fields.get('VmRSS', '0 kB').split()[0]) * 1024
     47                 except (OSError, ValueError, IndexError):
     48                     pass
     49             peak = max(peak, resident)
     50             if resident > limit:
     51                 reason = 'aggregate RSS guard exceeded'
     52                 break
     53             if time.monotonic() - started > args.seconds:
     54                 reason = 'whole-job deadline exceeded'
     55                 break
     56             time.sleep(.1)
     57     finally:
     58         try:
     59             os.killpg(child.pid, signal.SIGKILL)
     60         except ProcessLookupError:
     61             pass
     62         child.wait()
     63     print(f'Bounded check: peak sampled group RSS {peak / 1024**2:.1f} MiB; {time.monotonic()-started:.1f}s; {reason or "completed"}', flush=True)
     64     raise SystemExit(1 if reason else child.returncode)
     65 
     66 
     67 if __name__ == '__main__':
     68     main()