| 1 | """Subprocess helpers: safe timeout + process-group cleanup. |
| 2 | |
| 3 | Used by bird_x.py (Node.js Bird search) and youtube_yt.py (yt-dlp search |
| 4 | and transcript download). Both need the same os.setsid/killpg cleanup |
| 5 | dance on timeout to avoid orphaning child processes. |
| 6 | """ |
| 7 | |
| 8 | from __future__ import annotations |
| 9 | |
| 10 | import os |
| 11 | import signal |
| 12 | import subprocess |
| 13 | from dataclasses import dataclass |
| 14 | from typing import Optional, Sequence |
| 15 | |
| 16 | |
| 17 | class SubprocTimeout(Exception): |
| 18 | """Raised when a subprocess exceeds its timeout and is killed.""" |
| 19 | |
| 20 | |
| 21 | @dataclass |
| 22 | class SubprocResult: |
| 23 | """Result of a subprocess run that captured stdout and stderr.""" |
| 24 | |
| 25 | returncode: int |
| 26 | stdout: str |
| 27 | stderr: str |
| 28 | |
| 29 | |
| 30 | def run_with_timeout( |
| 31 | cmd: Sequence[str], |
| 32 | *, |
| 33 | timeout: int, |
| 34 | env: Optional[dict] = None, |
| 35 | on_pid: Optional[callable] = None, |
| 36 | ) -> SubprocResult: |
| 37 | """Run a subprocess with process-group cleanup on timeout. |
| 38 | |
| 39 | Spawns ``cmd`` inside its own process group via ``os.setsid`` where |
| 40 | available. If ``communicate(timeout=...)`` raises ``TimeoutExpired``, |
| 41 | signals ``SIGTERM`` to the entire group, falls back to ``proc.kill()`` |
| 42 | if the signal fails, then waits up to 5 seconds for cleanup, and |
| 43 | raises ``SubprocTimeout``. |
| 44 | |
| 45 | Args: |
| 46 | cmd: Command and arguments to spawn. |
| 47 | timeout: Timeout in seconds passed to ``communicate()``. |
| 48 | env: Optional environment dict. If None, inherits parent env. |
| 49 | on_pid: Optional callable invoked with the child PID right after |
| 50 | spawn. Used by bird_x.py to register child PIDs for cleanup |
| 51 | tracking. Exceptions raised by the callback are suppressed. |
| 52 | |
| 53 | Returns: |
| 54 | SubprocResult with returncode, stdout, and stderr as strings. |
| 55 | |
| 56 | Raises: |
| 57 | SubprocTimeout: If the process exceeded ``timeout``. |
| 58 | FileNotFoundError: If the executable is not found. |
| 59 | OSError: For other spawn failures. |
| 60 | """ |
| 61 | preexec = os.setsid if hasattr(os, "setsid") else None |
| 62 | |
| 63 | proc = subprocess.Popen( |
| 64 | list(cmd), |
| 65 | stdout=subprocess.PIPE, |
| 66 | stderr=subprocess.PIPE, |
| 67 | text=True, |
| 68 | encoding="utf-8", |
| 69 | errors="replace", |
| 70 | preexec_fn=preexec, |
| 71 | env=env, |
| 72 | ) |
| 73 | |
| 74 | if on_pid is not None: |
| 75 | try: |
| 76 | on_pid(proc.pid) |
| 77 | except Exception: |
| 78 | pass |
| 79 | |
| 80 | try: |
| 81 | stdout, stderr = proc.communicate(timeout=timeout) |
| 82 | except subprocess.TimeoutExpired: |
| 83 | try: |
| 84 | if hasattr(os, "killpg") and hasattr(os, "getpgid"): |
| 85 | os.killpg(os.getpgid(proc.pid), signal.SIGTERM) |
| 86 | else: |
| 87 | proc.kill() |
| 88 | except (ProcessLookupError, PermissionError, OSError, AttributeError): |
| 89 | proc.kill() |
| 90 | try: |
| 91 | proc.wait(timeout=5) |
| 92 | except subprocess.TimeoutExpired: |
| 93 | # Child ignored SIGTERM (or our killpg lost the race); escalate. |
| 94 | # Guard killpg/getpgid the same way the SIGTERM path above does: |
| 95 | # they are POSIX-only and raise AttributeError on Windows. The |
| 96 | # primary path was hardened in #552; this mirrors that guard on the |
| 97 | # escalation path (added later in #433) so the same crash can't |
| 98 | # re-surface here (#588). |
| 99 | try: |
| 100 | if hasattr(os, "killpg") and hasattr(os, "getpgid"): |
| 101 | os.killpg(os.getpgid(proc.pid), signal.SIGKILL) |
| 102 | else: |
| 103 | proc.kill() |
| 104 | except (ProcessLookupError, PermissionError, OSError, AttributeError): |
| 105 | proc.kill() |
| 106 | try: |
| 107 | proc.wait(timeout=5) |
| 108 | except subprocess.TimeoutExpired: |
| 109 | pass # process unkillable (e.g. D-state); leave as zombie |
| 110 | raise SubprocTimeout(f"Command {cmd[0]} timed out after {timeout}s") |
| 111 | |
| 112 | return SubprocResult( |
| 113 | returncode=proc.returncode, |
| 114 | stdout=stdout or "", |
| 115 | stderr=stderr or "", |
| 116 | ) |
| 117 |