diff --git a/servers/_process.py b/servers/_process.py index 21ad12a37..35b794898 100644 --- a/servers/_process.py +++ b/servers/_process.py @@ -1,4 +1,4 @@ -"""Give a Lean child a clean environment, and stop and reap its dedicated POSIX process group.""" +"""Give a Lean child a clean environment, find the processes carrying a token the caller set, and stop and reap its dedicated POSIX process group.""" from __future__ import annotations @@ -7,6 +7,7 @@ import signal import subprocess import time +from typing import Any REPL_ABORT_TERM_SECONDS = 0.5 REPL_ABORT_KILL_SECONDS = 1.0 @@ -28,7 +29,7 @@ def _process_group_exists(process_group_id: int) -> bool: return True -def _process_group_has_live_members(process_group_id: int) -> bool: +def process_group_has_live_members(process_group_id: int) -> bool: """Treat an all-zombie process group as terminated after SIGKILL.""" if not _process_group_exists(process_group_id): return False @@ -49,7 +50,7 @@ def _process_group_has_live_members(process_group_id: int) -> bool: def _wait_for_live_process_group_exit(process_group_id: int, deadline: float) -> bool: - while _process_group_has_live_members(process_group_id): + while process_group_has_live_members(process_group_id): remaining = deadline - time.monotonic() if remaining <= 0: return False @@ -77,7 +78,7 @@ def _kill_subprocesses( ) -> None: """Terminate and verify one dedicated POSIX process group.""" if getattr(process, "returncode", None) is not None: - if _process_group_has_live_members(process_group_id): + if process_group_has_live_members(process_group_id): raise RuntimeError( "refusing to signal a process group after its leader was reaped" ) @@ -98,7 +99,7 @@ def _terminate_process_group( term_deadline = min(term_deadline, deadline) kill_deadline = min(kill_deadline, deadline) - if _process_group_has_live_members(process_group_id): + if process_group_has_live_members(process_group_id): try: os.killpg(process_group_id, signal.SIGTERM) except ProcessLookupError: @@ -107,7 +108,7 @@ def _terminate_process_group( raise RuntimeError("failed to terminate the Lean process group") from error _wait_for_live_process_group_exit(process_group_id, term_deadline) - if _process_group_has_live_members(process_group_id): + if process_group_has_live_members(process_group_id): try: os.killpg(process_group_id, signal.SIGKILL) except ProcessLookupError: @@ -143,8 +144,44 @@ def _reap_process( raise RuntimeError("timed out reaping the Lean process") -def _inherit_clean_env() -> dict[str, str]: +def inherit_clean_env() -> dict[str, str]: """Return a copy of the current environment without PYTHONPATH noise.""" env = os.environ.copy() env.pop("PYTHONPATH", None) return env + + +def processes_with_token(name: str, token: str) -> list[Any]: + """Return the processes whose environment sets ``name`` to ``token``. + + A process whose environment cannot be read is left out, so it is never + signalled: nothing shows that it carries the token. Lean and Lake run as + the user, who can read their environments; setuid programs the user runs, + such as ``login``, are unreadable at all times, so counting them would + fail every close. A process table that cannot be read fails the scan. + """ + if not token: + # A missing token would match every process without the variable. + raise ValueError("a process token must be a nonempty string") + import psutil + + try: + pids = psutil.pids() + except (OSError, psutil.Error) as error: + raise RuntimeError(f"cannot read the process table to find the processes carrying {name}") from error + found = [] + for pid in pids: + if pid == os.getpid(): + continue + try: + # The handle records the creation time before the environment is + # read, so it can only name the process that was read or an + # earlier one, which psutil then refuses to signal. + candidate = psutil.Process(pid) + environment = candidate.environ() + except (psutil.NoSuchProcess, psutil.AccessDenied): + # Exited or unreadable; a zombie reads as exited or empty. + continue + if environment.get(name) == token: + found.append(candidate) + return found diff --git a/servers/batch_diagnostics.py b/servers/batch_diagnostics.py index 1a593c234..b10b60e7e 100644 --- a/servers/batch_diagnostics.py +++ b/servers/batch_diagnostics.py @@ -21,9 +21,9 @@ from typing import Any from servers._process import ( - _inherit_clean_env, _kill_subprocesses, _wait_for_live_process_group_exit, + inherit_clean_env, ) logger = logging.getLogger(__name__) @@ -118,7 +118,7 @@ def _run(argv: list[str], root: Path, timeout: float) -> tuple[int, bytes | None process = subprocess.Popen( argv, cwd=root, - env={**_inherit_clean_env(), "TMPDIR": scratch}, + env={**inherit_clean_env(), "TMPDIR": scratch}, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, diff --git a/servers/lsp/server.py b/servers/lsp/server.py index 2cd161d35..fba8677f5 100644 --- a/servers/lsp/server.py +++ b/servers/lsp/server.py @@ -44,11 +44,12 @@ from servers._process import ( REPL_ABORT_KILL_SECONDS, REPL_ABORT_TERM_SECONDS, - _inherit_clean_env, - _process_group_has_live_members, _reap_process, _terminate_process_group, _wait_for_live_process_group_exit, + inherit_clean_env, + process_group_has_live_members, + processes_with_token, ) from servers.lean_client import LeanRuntimeClient @@ -120,42 +121,6 @@ class LspConfig: timeout: float = DEFAULT_LSP_TIMEOUT -def _processes_with_token(token: str) -> list[Any]: - """Return the processes whose environment carries a session's token. - - A process whose environment cannot be read is left out, so it is never - signalled: nothing shows that it is the session's. Lean and Lake run as - the user, who can read their environments; setuid programs the user runs, - such as ``login``, are unreadable at all times, so counting them would - fail every close. A process table that cannot be read fails the scan. - """ - if not token: - # A missing token would match every process without the variable. - raise ValueError("a Lean LSP session token must be a nonempty string") - import psutil - - try: - pids = psutil.pids() - except (OSError, psutil.Error) as error: - raise RuntimeError("cannot read the process table to find the Lean LSP session's processes") from error - found = [] - for pid in pids: - if pid == os.getpid(): - continue - try: - # The handle records the creation time before the environment is - # read, so it can only name the process that was read or an - # earlier one, which psutil then refuses to signal. - candidate = psutil.Process(pid) - environment = candidate.environ() - except (psutil.NoSuchProcess, psutil.AccessDenied): - # Exited or unreadable; a zombie reads as exited or empty. - continue - if environment.get(_SESSION_TOKEN_ENV) == token: - found.append(candidate) - return found - - def _running(process: Any) -> bool: """Whether a psutil process is still running; one that cannot be read counts as running.""" import psutil @@ -197,7 +162,7 @@ def start(self) -> None: """Start the language server process.""" if os.name != "posix": raise RuntimeError("Lean LSP transport requires a POSIX platform") - env = _inherit_clean_env() + env = inherit_clean_env() self._token = secrets.token_hex(16) env[_SESSION_TOKEN_ENV] = self._token process: subprocess.Popen | None = None @@ -255,7 +220,7 @@ def _retire_once(self) -> None: process_group_id = process.pid self._process_group_id = process_group_id if getattr(process, "returncode", None) is not None: - if _process_group_has_live_members(process_group_id): + if process_group_has_live_members(process_group_id): raise RuntimeError( "refusing to signal an LSP process group after its leader was reaped" ) @@ -309,7 +274,7 @@ def _stop_processes_with_token(self) -> None: (signal.SIGTERM, REPL_ABORT_TERM_SECONDS), (signal.SIGKILL, REPL_ABORT_KILL_SECONDS), ): - processes = _processes_with_token(token) + processes = processes_with_token(_SESSION_TOKEN_ENV, token) if not processes: return for process in processes: @@ -320,7 +285,7 @@ def _stop_processes_with_token(self) -> None: except (OSError, psutil.Error): pass _wait_for_exit(processes, time.monotonic() + seconds) - if _processes_with_token(token): + if processes_with_token(_SESSION_TOKEN_ENV, token): raise RuntimeError("processes started by the Lean LSP server are still running") def _abort_process(self) -> None: @@ -539,7 +504,7 @@ def _left_running(self, deadline: float) -> bool: if token is None: return False try: - carriers = _processes_with_token(token) + carriers = processes_with_token(_SESSION_TOKEN_ENV, token) except RuntimeError: return True outside = [] diff --git a/servers/repl/core.py b/servers/repl/core.py index a093b86dd..301d2f0e1 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -25,7 +25,7 @@ from logging import getLogger from typing import Any, Callable -from servers._process import _inherit_clean_env, _kill_subprocesses +from servers._process import _kill_subprocesses, inherit_clean_env logger = getLogger(__name__) @@ -737,7 +737,7 @@ def _spawn(self, command: list[str]) -> subprocess.Popen: The child leads its own session, and close() retires that whole process group through the id published here. """ - env = _inherit_clean_env() + env = inherit_clean_env() env.update(self.config.env) self.process = subprocess.Popen( command, diff --git a/tests/test_lsp_protocol.py b/tests/test_lsp_protocol.py index f16bf6bd4..a7ee43a69 100644 --- a/tests/test_lsp_protocol.py +++ b/tests/test_lsp_protocol.py @@ -18,6 +18,7 @@ import psutil import pytest +from servers._process import processes_with_token from servers.lsp import server as lsp from tests.lean_fixture import _real_lean_project, _requires_real_lean @@ -336,9 +337,11 @@ def test_abort_retries_the_token_sweep_before_reaping_the_leader(monkeypatch): # Each scan of the first attempt finds a process carrying the token, so # one is left after SIGKILL; the retry's first scan finds none. scans = iter([[_FakeTokenCarrier(events)]] * 3) + # Asserting inside the scan would only be retried by abort, so record. + calls: list[tuple[str, str]] = [] - def scan(token): - assert token == "session-token" + def scan(name, token): + calls.append((name, token)) events.append("scan") return next(scans, []) @@ -346,7 +349,7 @@ def reap(owned): events.append("reap") owned.wait() - monkeypatch.setattr(lsp, "_processes_with_token", scan) + monkeypatch.setattr(lsp, "processes_with_token", scan) monkeypatch.setattr(lsp, "_terminate_process_group", lambda process_group_id: events.append("group")) monkeypatch.setattr(lsp, "_reap_process", reap) monkeypatch.setattr(lsp.time, "sleep", lambda delay: None) @@ -357,6 +360,7 @@ def reap(owned): # the retired group again, and the leader was reaped only after a sweep # found nothing left. assert events == ["group", "scan", "SIGTERM", "scan", "SIGKILL", "scan", "scan", "reap"] + assert calls == [(lsp._SESSION_TOKEN_ENV, "session-token")] * 4 assert session.process is None @@ -366,7 +370,7 @@ def test_token_sweep_signals_only_what_its_latest_scan_found(monkeypatch): session = _owned_session(_FakeProcess(), "session-token") signals: list[str] = [] scans = iter([[_FakeTokenCarrier(signals)], []]) - monkeypatch.setattr(lsp, "_processes_with_token", lambda token: next(scans, [])) + monkeypatch.setattr(lsp, "processes_with_token", lambda name, token: next(scans, [])) monkeypatch.setattr(lsp, "_terminate_process_group", lambda process_group_id: None) monkeypatch.setattr(lsp, "_reap_process", lambda owned: owned.wait()) @@ -539,7 +543,7 @@ class StopRetry(BaseException): monkeypatch.setattr( lsp, - "_process_group_has_live_members", + "process_group_has_live_members", lambda process_group_id: True, ) monkeypatch.setattr( @@ -575,7 +579,9 @@ def test_token_scan_finds_only_the_processes_carrying_the_token(): } try: # Popen returns once the child has exec'd, so its environment is set. - assert [process.pid for process in lsp._processes_with_token(token)] == [processes["this session"].pid] + assert [process.pid for process in processes_with_token(lsp._SESSION_TOKEN_ENV, token)] == [ + processes["this session"].pid + ] finally: for process in processes.values(): process.kill() @@ -605,7 +611,25 @@ def environ(self): monkeypatch.setattr(psutil, "pids", lambda: list(candidates)) monkeypatch.setattr(psutil, "Process", lambda pid: candidates[pid]) - assert lsp._processes_with_token("session-token") == [candidates[1]] + assert processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token") == [candidates[1]] + + +def test_token_scan_reads_the_variable_it_is_given(monkeypatch): + class Candidate: + def __init__(self, environment): + self.environment = environment + + def environ(self): + return self.environment + + candidates = { + 1: Candidate({"OTHER_TOKEN": "session-token"}), + 2: Candidate({lsp._SESSION_TOKEN_ENV: "session-token"}), + } + monkeypatch.setattr(psutil, "pids", lambda: list(candidates)) + monkeypatch.setattr(psutil, "Process", lambda pid: candidates[pid]) + + assert processes_with_token("OTHER_TOKEN", "session-token") == [candidates[1]] def test_token_scan_fails_when_the_process_table_cannot_be_read(monkeypatch): @@ -616,7 +640,7 @@ def unreadable(): # The attempt fails, so close and abort retry instead of reaping the leader. with pytest.raises(RuntimeError, match="cannot read the process table"): - lsp._processes_with_token("session-token") + processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token") @pytest.mark.parametrize( @@ -635,14 +659,14 @@ def environ(self): # Only exited and unreadable processes are passed over; anything else # fails the attempt, so close and abort retry. with pytest.raises(type(error)): - lsp._processes_with_token("session-token") + processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token") @pytest.mark.parametrize("token", [None, ""], ids=["none", "empty"]) def test_token_scan_refuses_a_missing_token(token): # A missing token would match every process without the variable. with pytest.raises(ValueError, match="nonempty"): - lsp._processes_with_token(token) + processes_with_token(lsp._SESSION_TOKEN_ENV, token) def test_concurrent_abort_and_close_serialize_owned_cleanup(monkeypatch): diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index 27ed254ca..e69eadcb1 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -206,7 +206,7 @@ def test_process_group_cleanup_escalates_and_reports_timeout(monkeypatch): ) monkeypatch.setattr( process_group, - "_process_group_has_live_members", + "process_group_has_live_members", lambda process_group_id: next(live_results), ) monkeypatch.setattr( @@ -242,7 +242,7 @@ def test_process_group_cleanup_does_not_signal_an_all_zombie_group(monkeypatch): ) monkeypatch.setattr( process_group, - "_process_group_has_live_members", + "process_group_has_live_members", lambda process_group_id: next(live_results), ) monkeypatch.setattr( @@ -267,7 +267,7 @@ def test_process_group_cleanup_refuses_a_pre_reaped_leader_with_live_members( process = SimpleNamespace(returncode=0) monkeypatch.setattr( process_group, - "_process_group_has_live_members", + "process_group_has_live_members", lambda process_group_id: True, ) monkeypatch.setattr( @@ -1417,7 +1417,7 @@ def test_header_check_reaps_descendants_after_success(tmp_path): assert _header_modules([sys.executable, "-c", script]) == ["Mathlib"] process_ids = json.loads(process_record.read_text()) - assert not process_group._process_group_has_live_members(process_ids["group"]) + assert not process_group.process_group_has_live_members(process_ids["group"]) # Answers one frame, then waits for the next as the REPL does.