Viewing files at commit 1a441edafc71Commit detailsView latest

suites/robustness/robustness/procs.py 178 lines · 5.8 KB · python
1
"""Process control for robustness cases: isolated env, timeouts, kill -9.
2
3
Only PIDs this module starts are ever signalled (never pkill by pattern).
4
"""
5
6
from __future__ import annotations
7
8
import os
9
import platform
10
import shutil
11
import signal
12
import socket
13
import subprocess
14
import threading
15
import time
16
from dataclasses import dataclass, field
17
from pathlib import Path
18
from typing import Callable, Optional
19
20
from . import REPO_ROOT  # noqa: F401  (ensures scripts/ on sys.path)
21
from oakbench.environment import base_env
22
23
DEFAULT_TIMEOUT_S = 90.0
24
25
26
@dataclass
27
class Result:
28
    argv: list[str]
29
    returncode: Optional[int]  # None = timed out (killed by harness)
30
    stdout: str
31
    stderr: str
32
    elapsed_ms: float
33
    killed_by_harness: bool = False
34
    timed_out: bool = False
35
36
    @property
37
    def ok(self) -> bool:
38
        return self.returncode == 0
39
40
    def text(self) -> str:
41
        return (self.stdout or "") + "\n" + (self.stderr or "")
42
43
    def brief(self, limit: int = 600) -> dict:
44
        return {
45
            "argv": self.argv,
46
            "returncode": self.returncode,
47
            "elapsed_ms": round(self.elapsed_ms, 3),
48
            "killed_by_harness": self.killed_by_harness,
49
            "timed_out": self.timed_out,
50
            "stdout_tail": (self.stdout or "")[-limit:],
51
            "stderr_tail": (self.stderr or "")[-limit:],
52
        }
53
54
55
def subject_env(home: Path, extra: Optional[dict] = None) -> dict:
56
    """Isolated HOME + deterministic identity; no inherited Oak remote/repo."""
57
    home.mkdir(parents=True, exist_ok=True)
58
    env = base_env(home=home)
59
    for key in ("OAK_REMOTE", "OAK_REPO", "OAK_TOKEN", "OAK_API_KEY", "DYLD_INSERT_LIBRARIES",
60
                "OAK_LOG", "OAK_VERBOSE", "OAK_PROGRESS"):
61
        env.pop(key, None)
62
    env["OAK_PROGRESS"] = "never"
63
    if extra:
64
        env.update(extra)
65
    return env
66
67
68
def run(argv: list[str], cwd: Path, env: dict, timeout: float = DEFAULT_TIMEOUT_S,
69
        stdin: Optional[str] = None) -> Result:
70
    start = time.perf_counter()
71
    proc = subprocess.Popen(argv, cwd=str(cwd), env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
72
                            stdin=subprocess.PIPE if stdin is not None else subprocess.DEVNULL,
73
                            start_new_session=True)
74
    try:
75
        out, err = proc.communicate(input=stdin.encode() if stdin is not None else None, timeout=timeout)
76
        timed_out = False
77
    except subprocess.TimeoutExpired:
78
        _kill_group(proc.pid)
79
        out, err = proc.communicate()
80
        timed_out = True
81
    elapsed = (time.perf_counter() - start) * 1000.0
82
    return Result(argv, None if timed_out else proc.returncode, out.decode("utf-8", "replace"),
83
                  err.decode("utf-8", "replace"), elapsed, timed_out=timed_out, killed_by_harness=timed_out)
84
85
86
def _kill_group(pid: int) -> None:
87
    try:
88
        os.killpg(pid, signal.SIGKILL)
89
    except (ProcessLookupError, PermissionError):
90
        try:
91
            os.kill(pid, signal.SIGKILL)
92
        except ProcessLookupError:
93
            pass
94
95
96
def run_killed_after(argv: list[str], cwd: Path, env: dict, delay_s: float,
97
                     timeout: float = DEFAULT_TIMEOUT_S, whole_group: bool = False) -> Result:
98
    """Start argv and SIGKILL it after delay_s (time-based crash point).
99
100
    whole_group=False kills only the client PID: children (e.g. Git's local
101
    receive-pack) survive, like a server that keeps going after its client died.
102
    """
103
    start = time.perf_counter()
104
    proc = subprocess.Popen(argv, cwd=str(cwd), env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
105
                            stdin=subprocess.DEVNULL, start_new_session=True)
106
    killed = threading.Event()
107
108
    def killer() -> None:
109
        if proc.poll() is None:
110
            try:
111
                if whole_group:
112
                    os.killpg(proc.pid, signal.SIGKILL)
113
                else:
114
                    os.kill(proc.pid, signal.SIGKILL)
115
                killed.set()
116
            except ProcessLookupError:
117
                pass
118
119
    timer = threading.Timer(delay_s, killer)
120
    timer.start()
121
    try:
122
        out, err = proc.communicate(timeout=timeout)
123
        timed_out = False
124
    except subprocess.TimeoutExpired:
125
        _kill_group(proc.pid)
126
        out, err = proc.communicate()
127
        timed_out = True
128
    finally:
129
        timer.cancel()
130
    elapsed = (time.perf_counter() - start) * 1000.0
131
    rc = None if timed_out else proc.returncode
132
    return Result(argv, rc, out.decode("utf-8", "replace"), err.decode("utf-8", "replace"), elapsed,
133
                  killed_by_harness=killed.is_set() or timed_out, timed_out=timed_out)
134
135
136
def free_port() -> int:
137
    with socket.socket() as sock:
138
        sock.bind(("127.0.0.1", 0))
139
        return int(sock.getsockname()[1])
140
141
142
def copy_tree(src: Path, dst: Path) -> None:
143
    if dst.exists():
144
        shutil.rmtree(dst)
145
    if platform.system() == "Darwin":
146
        # APFS clonefile copy: byte-identical, copy-on-write, ~10x faster per trial.
147
        proc = subprocess.run(["/bin/cp", "-cR", str(src), str(dst)], capture_output=True)
148
        if proc.returncode == 0:
149
            return
150
        shutil.rmtree(dst, ignore_errors=True)
151
    shutil.copytree(src, dst, symlinks=True)
152
153
154
def wait_until(predicate: Callable[[], bool], timeout: float, interval: float = 0.05) -> bool:
155
    deadline = time.monotonic() + timeout
156
    while time.monotonic() < deadline:
157
        if predicate():
158
            return True
159
        time.sleep(interval)
160
    return predicate()
161
162
163
@dataclass
164
class Spawned:
165
    """A long-lived helper process we own (loopback serve)."""
166
167
    proc: subprocess.Popen
168
    log_path: Path
169
    extra: dict = field(default_factory=dict)
170
171
    def stop(self) -> None:
172
        if self.proc.poll() is None:
173
            self.proc.terminate()
174
            try:
175
                self.proc.wait(timeout=5)
176
            except subprocess.TimeoutExpired:
177
                self.proc.kill()
178
                self.proc.wait(timeout=5)