process_rss.py (6318B)
1 """Linux direct-child RSS monitoring, including mismatched procfs PID namespaces. 2 3 Never treat an inaccessible/wrong process as zero RSS. This measures the direct 4 child, not descendant totals; callers retain their heap/time/process-group guards. 5 """ 6 import os 7 import subprocess 8 import signal 9 import tempfile 10 import time 11 from pathlib import Path 12 13 14 def status_fields(text): 15 return dict(line.split(':', 1) for line in text.splitlines() if ':' in line) 16 17 18 class ChildRSS: 19 def __init__(self, proc): 20 self.proc = proc 21 own = status_fields(Path('/proc/self/status').read_text()) 22 self.parent = int(own['Pid']) 23 self.namespace_index = len(own.get('NSpid', own['Pid']).split()) - 1 24 if int(own.get('NSpid', own['Pid']).split()[-1]) != os.getpid(): 25 raise RuntimeError('cannot relate current PID namespace to procfs') 26 self.pid = None 27 28 def matches(self, fields): 29 ids = fields.get('NSpid', fields.get('Pid', '')).split() 30 return (int(fields.get('PPid', '-1')) == self.parent and 31 len(ids) > self.namespace_index and 32 int(ids[self.namespace_index]) == self.proc.pid) 33 34 def candidates(self): 35 # Fast path when procfs is mounted in this process's PID namespace. 36 yield self.proc.pid 37 # children reports IDs in the procfs mount's namespace. Use all parent 38 # threads, since Popen need not have been called on the main thread. 39 for path in Path('/proc/self/task').glob('*/children'): 40 try: 41 for value in path.read_text().split(): 42 yield int(value) 43 except FileNotFoundError: 44 pass 45 46 def sample(self): 47 if self.proc.poll() is not None: 48 return None 49 candidates = [self.pid] if self.pid is not None else self.candidates() 50 diagnostic = 'no visible candidate' 51 for pid in candidates: 52 try: 53 fields = status_fields(Path(f'/proc/{pid}/status').read_text()) 54 except FileNotFoundError: 55 continue 56 diagnostic = repr({key: fields.get(key) for key in ('Pid', 'PPid', 'NSpid', 'State', 'VmRSS')}) 57 if not self.matches(fields): 58 continue 59 self.pid = pid 60 if 'VmRSS' in fields: 61 value, unit = fields['VmRSS'].split() 62 if unit != 'kB': 63 raise RuntimeError('unexpected RSS unit') 64 return int(value) 65 # A thread-group leader may release its own mm before surviving 66 # threads finish. Their status still reports the shared address 67 # space; validate Tgid before accepting a thread sample. 68 for path in Path(f'/proc/{pid}/task').glob('*/status'): 69 try: 70 task = status_fields(path.read_text()) 71 except FileNotFoundError: 72 continue 73 if int(task.get('Tgid', '-1')) == pid and 'VmRSS' in task: 74 value, unit = task['VmRSS'].split() 75 if unit != 'kB': 76 raise RuntimeError('unexpected thread RSS unit') 77 return int(value) 78 # exit_mm can clear mm while the task remains R during kernel 79 # teardown. statm explicitly reports seven zero fields for no mm; 80 # unlike missing status, this is a verified zero-memory observation. 81 try: 82 statm = Path(f'/proc/{pid}/statm').read_text().split() 83 except FileNotFoundError: 84 statm = [] 85 if len(statm) == 7 and all(value == '0' for value in statm): 86 return 0 87 # Zombie leaders can spend significant time releasing a large 88 # address space before waitpid becomes ready. Identity is already 89 # validated and no surviving thread above reports an address space. 90 # Keep polling/reaping under the caller's time guard, not an 91 # arbitrary exit-grace deadline; this is not a live-memory sample. 92 if fields.get('State', '').strip().startswith('Z'): 93 return None 94 # A child may exit between poll and status reads. Otherwise absence of 95 # a sample is a monitoring failure, not a valid zero-memory observation. 96 if self.proc.poll() is not None: 97 return None 98 # During exit, the kernel can release the address space before waitpid 99 # reports termination. Permit only a bounded reap, not unmonitored work. 100 try: 101 self.proc.wait(timeout=0.25) 102 except subprocess.TimeoutExpired: 103 raise RuntimeError('cannot monitor child RSS; refusing an unguarded run: ' + diagnostic) 104 return None 105 106 107 def run_guarded(command, *, memory_mib, seconds): 108 """Capture a bounded direct-child job without pipe backpressure. 109 110 RSS sampling is Linux-only; monitoring failure aborts, never grants an 111 unmonitored run. Returns stdout and sampled peak KiB. Error output is merged. 112 """ 113 peak = 0 114 with tempfile.TemporaryFile(mode='w+') as output: 115 process = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT, 116 start_new_session=True) 117 start = time.monotonic() 118 try: 119 monitor = ChildRSS(process) 120 while process.poll() is None: 121 rss = monitor.sample() 122 if rss is not None: 123 peak = max(peak, rss) 124 if rss > memory_mib * 1024: 125 raise RuntimeError('RSS guard exceeded') 126 if time.monotonic() - start > seconds: 127 raise RuntimeError('wall-time guard exceeded') 128 time.sleep(0.05) 129 finally: 130 if process.poll() is None: 131 os.killpg(process.pid, signal.SIGTERM) 132 try: 133 process.wait(timeout=10) 134 except subprocess.TimeoutExpired: 135 os.killpg(process.pid, signal.SIGKILL) 136 process.wait() 137 output.seek(0) 138 reply = output.read() 139 if process.returncode: 140 raise subprocess.CalledProcessError(process.returncode, command, output=reply) 141 return reply, peak