Python subprocess 超時治理：POSIX Process Group、升級終止與輸出收尾

supervisor（我們的 Python）
└── parent（新 session，PID = PGID）
    └── worker（繼承 parent 的 process group 與輸出 pipe）

mkdir pypy-shutdown-lab
cd pypy-shutdown-lab
python3 --version

import math
import os
import selectors
import signal
import subprocess
import time
from dataclasses import dataclass


@dataclass
class Outcome:
    returncode: int
    timed_out: bool
    eof: bool
    discarded: int
    output: bytes
    elapsed: float


def supervise(argv, timeout=1.0, grace=0.3, drain=0.5, limit=65536):
    required = ('waitid', 'WNOWAIT', 'WEXITED', 'WNOHANG')
    if os.name != 'posix' or not all(hasattr(os, x) for x in required):
        raise NotImplementedError('requires POSIX waitid / WNOWAIT')
    if any(not math.isfinite(x) or x <= 0 for x in (timeout, grace, drain)):
        raise ValueError('durations must be finite and positive')
    if not isinstance(limit, int) or limit < 0:
        raise ValueError('limit must be a nonnegative integer')
    started = time.monotonic()
    proc = subprocess.Popen(
        argv, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT, start_new_session=True, bufsize=0,
    )
    sel = selectors.DefaultSelector()
    kept = bytearray()
    discarded = 0
    eof = False

    def signal_group(sig):
        try:
            os.killpg(proc.pid, sig)
        except ProcessLookupError:
            pass
        except PermissionError:
            # Accept EPERM only after observing a zombie-only owned group.
            rows = subprocess.run(
                ['ps', '-axo', 'pgid=,stat='], capture_output=True,
                text=True, check=True, timeout=0.5,
            ).stdout.splitlines()
            states = [r.split()[1] for r in rows
                      if r.split() and r.split()[0] == str(proc.pid)]
            if not (states and all(s.startswith('Z') for s in states)
                    and leader_exited()):
                raise

    def leader_exited():
        info = os.waitid(os.P_PID, proc.pid,
                         os.WEXITED | os.WNOHANG | os.WNOWAIT)
        return info is not None and info.si_pid == proc.pid

    def pump(until, early=True):
        nonlocal discarded, eof
        while time.monotonic() < until:
            if early and eof and leader_exited():
                return True
            delay = min(0.02, max(0, until - time.monotonic()))
            if not sel.get_map():
                time.sleep(delay)
                continue
            for key, _ in sel.select(delay):
                try:
                    chunk = os.read(key.fd, 8192)
                except BlockingIOError:
                    continue
                if not chunk:
                    sel.unregister(key.fileobj)
                    eof = True
                    continue
                room = max(0, limit - len(kept))
                kept.extend(chunk[:room])
                discarded += max(0, len(chunk) - room)
        return eof and leader_exited()

    try:
        os.set_blocking(proc.stdout.fileno(), False)
        sel.register(proc.stdout, selectors.EVENT_READ)
        timed_out = not pump(started + timeout)
        if timed_out:
            signal_group(signal.SIGTERM)
            pump(time.monotonic() + grace, early=False)
    finally:
        # Keep the leader unreaped until the final group signal.
        try:
            signal_group(signal.SIGKILL)
            pump(time.monotonic() + drain)
        finally:
            sel.close()
            proc.stdout.close()
            proc.wait(timeout=drain)
    return Outcome(proc.returncode, timed_out, eof, discarded,
                   bytes(kept), time.monotonic() - started)

import os
import signal
import subprocess
import sys
import time

mode = sys.argv[1]
if len(sys.argv) == 3 and sys.argv[2] == 'worker':
    def on_term(*_):
        print('CHILD_TERM', flush=True)
        raise SystemExit(0)

    handler = signal.SIG_IGN if mode == 'ignore' else on_term
    signal.signal(signal.SIGTERM, handler)
    print(f'CHILD {os.getpid()} PGID {os.getpgrp()}', flush=True)
    print('PARTIAL_STDERR', file=sys.stderr, flush=True)
    time.sleep(30)  # Finite lab failsafe, not the supervisor deadline.
else:
    signal.signal(signal.SIGTERM,
                  lambda *_: print('PARENT_TERM', flush=True))
    child = subprocess.Popen([sys.executable, '-u', __file__, mode, 'worker'])
    print(f'PARENT {os.getpid()} CHILD_PID {child.pid}', flush=True)
    if mode == 'holder':
        raise SystemExit(0)  # Child still owns the inherited output pipe.
    raise SystemExit(child.wait())

import re
import signal
import subprocess
import sys
import time
import unittest
from pathlib import Path
from supervisor import supervise

FIXTURE = str(Path(__file__).with_name('fixture.py'))


def assert_not_running(test, pid):
    # A zombie is dead, but only its own parent/adopter can reap it.
    until = time.monotonic() + 2
    while time.monotonic() < until:
        state = subprocess.run(
            ['ps', '-p', str(pid), '-o', 'stat='],
            capture_output=True, text=True, timeout=1,
        ).stdout.strip()
        if not state or state.startswith('Z'):
            return
        time.sleep(0.02)
    test.fail(f'fixture process still running: {pid} {state}')


class ShutdownTests(unittest.TestCase):
    def test_tree_modes(self):
        for mode in ('cooperate', 'ignore', 'holder'):
            with self.subTest(mode=mode):
                result = supervise([sys.executable, '-u', FIXTURE, mode])
                text = result.output.decode('utf-8', errors='replace')
                self.assertTrue(result.timed_out)
                self.assertTrue(result.eof)
                self.assertEqual(result.discarded, 0)
                self.assertIn('PARTIAL_STDERR', text)
                self.assertLess(result.elapsed, 4)
                parent, child = map(int, re.search(
                    r'PARENT (\d+) CHILD_PID (\d+)', text).groups())
                pgid = int(re.search(r'PGID (\d+)', text).group(1))
                self.assertEqual(pgid, parent)
                self.assertNotEqual(pgid, __import__('os').getpgrp())
                expected = -signal.SIGKILL if mode == 'ignore' else 0
                self.assertEqual(result.returncode, expected)
                if mode != 'ignore':
                    self.assertIn('CHILD_TERM', text)
                assert_not_running(self, parent)
                assert_not_running(self, child)
                print(mode, result.returncode, result.eof)

    def test_normal_and_nonzero(self):
        for code in (0, 7):
            with self.subTest(code=code):
                result = supervise([sys.executable, '-c',
                                    f"print('拍拍君'); raise SystemExit({code})"])
                self.assertFalse(result.timed_out)
                self.assertTrue(result.eof)
                self.assertEqual(result.returncode, code)
                self.assertIn('拍拍君', result.output.decode())

    def test_retention_limit_still_drains(self):
        result = supervise([sys.executable, '-c',
                            "import os; os.write(1, b'x' * 100000)"], limit=100)
        self.assertEqual(len(result.output), 100)
        self.assertEqual(result.discarded, 99900)
        self.assertTrue(result.eof)
        self.assertFalse(result.timed_out)


if __name__ == '__main__':
    unittest.main()

python3 -m unittest -v test_lab