diff options
| -rwxr-xr-x | research/guard.py | 110 |
1 files changed, 110 insertions, 0 deletions
diff --git a/research/guard.py b/research/guard.py new file mode 100755 index 0000000..5e9760f --- /dev/null +++ b/research/guard.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""A command run so that it cannot take the machine down with it: in a +session of its own, stopped whole (SIGKILL) when the session's processes +hold more than MOST MiB, or when the machine has less than FLOOR MiB +available (what it writes to a tmpfs such as /tmp counts there). The Pi has +8 GiB, swap in zram and no OOM daemon: a run that outgrows its memory has +every program's code evicted before the kernel kills anything, and the +machine stops answering for minutes (it was reset twice so, 2026-10-01: a +run held 13 GiB of output it should have streamed). + +guard.py [--most MiB] [--floor MiB] [--peak] -- COMMAND [ARG ...] + --most the session's resident memory at most (default 3072) + --floor the machine's available memory at least (default 1024) + --peak a JSON line on stderr at the end: the session's peak + The command's exit, or 137 when stopped (a JSON line on stderr says why). + What the session leaves running when the command ends is stopped too. + Stopped itself (SIGTERM, SIGINT, SIGHUP), it stops the session first. +""" +import ctypes +import json +import os +import signal +import subprocess +import sys +import time + +PAGE = os.sysconf('SC_PAGE_SIZE') +MIB = 1 << 20 + + +def available(): + with open('/proc/meminfo') as info: + for line in info: + if line.startswith('MemAvailable:'): + return int(line.split()[1]) * 1024 + return 0 + + +def members(session): + """The session's processes and their resident bytes.""" + found = {} + for entry in os.scandir('/proc'): + if not entry.name.isdigit(): + continue + try: + with open(f'/proc/{entry.name}/stat', 'rb') as stat: + fields = stat.read().rsplit(b')', 1)[1].split() + except OSError: # gone meanwhile + continue + if int(fields[3]) == session: + found[int(entry.name)] = int(fields[21]) * PAGE + return found + + +def stop(session): + for pid in members(session): + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +def main(): + args, most, floor, peak_wanted = sys.argv[1:], 3072, 1024, False + while args and args[0] != '--': + option = args.pop(0) + if option == '--peak': + peak_wanted = True + elif option in ('--most', '--floor') and args: + value = int(args.pop(0)) + most, floor = (value, floor) if option == '--most' else (most, value) + else: + sys.exit(__doc__) + command = args[1:] + if not command: + sys.exit(__doc__) + libc = ctypes.CDLL(None, use_errno=True) + # The command dies with the guard, however the guard ends (its own + # children are stopped through the session). + child = subprocess.Popen(command, start_new_session=True, preexec_fn=lambda: libc.prctl(1, signal.SIGKILL)) + session = child.pid + + def stopped_from_outside(number, frame): + stop(session) + sys.exit(128 + number) + + for number in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP): + signal.signal(number, stopped_from_outside) + peak, began = 0, time.monotonic() + while child.poll() is None: + held, free = sum(members(session).values()), available() + peak = max(peak, held) + if held > most * MIB or free < floor * MIB: + stop(session) + child.wait() + print(json.dumps({'guard': 'stopped', 'why': 'its memory' if held > most * MIB else "the machine's memory", + 'held_mib': held // MIB, 'most_mib': most, 'available_mib': free // MIB, + 'floor_mib': floor, 'after_s': round(time.monotonic() - began, 1), + 'command': ' '.join(command)[:200]}), file=sys.stderr, flush=True) + return 137 + time.sleep(0.25) + stop(session) # what it left running + if peak_wanted: + print(json.dumps({'guard': 'done', 'peak_mib': peak // MIB, 'exit': child.returncode, + 'seconds': round(time.monotonic() - began, 1)}), file=sys.stderr, flush=True) + return child.returncode + + +if __name__ == '__main__': + sys.exit(main()) |