diff --git a/orchestrator/orchestrator/runner.py b/orchestrator/orchestrator/runner.py index aba74be..29d5e13 100644 --- a/orchestrator/orchestrator/runner.py +++ b/orchestrator/orchestrator/runner.py @@ -29,13 +29,16 @@ from __future__ import annotations import os +import signal import socket import subprocess import sys +import threading import time import httpx from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import wait as futures_wait from .hostnames import validate_short @@ -211,25 +214,48 @@ def main() -> None: f"reprovision-runner [{cfg.runner_id}] → {cfg.api} " f"(auth: {auth}, poll {cfg.poll}s, max_concurrent {cfg.max_concurrent})" ) + # SIGTERM drains instead of dying. launchd sends it on every restart -- the step-ca + # renewal's `kickstart -k`, a puppet reload's bootout -- and a job's `reprovision` + # subprocess dying with the runner strands its host mid-EACS. launchd waits up to the + # plist's ExitTimeOut before SIGKILL, so that must cover a whole reprovision. + stop = threading.Event() + + def _on_sigterm(_signum, _frame) -> None: + print("SIGTERM: no new claims; draining in-flight jobs", flush=True) + stop.set() + + signal.signal(signal.SIGTERM, _on_sigterm) # httpx.Client is thread-safe (pooled connections), so one client is shared across job # threads. Claim up to max_concurrent jobs and run them in parallel — a whole pool of # workers reprovisions "at once" instead of serially. with httpx.Client(**cfg.client_kwargs) as client, \ ThreadPoolExecutor(max_workers=cfg.max_concurrent, thread_name_prefix="job") as pool: - active: set = set() - while True: - active = {f for f in active if not f.done()} # reap finished jobs - while len(active) < cfg.max_concurrent: # fill idle capacity, FIFO - try: - job = _claim(client, cfg) - except httpx.HTTPError as e: - print(f"claim failed: {e}") - break - if not job: - break # queue empty - print(f"claimed job {job['id']} → {job['short']} ({len(active) + 1}/{cfg.max_concurrent} active)") - active.add(pool.submit(_run_job_guarded, client, cfg, job)) - time.sleep(cfg.poll) + _serve(client, cfg, pool, stop) + print("drained: no jobs in flight, exiting", flush=True) + + +def _serve(client: httpx.Client, cfg: Config, pool: ThreadPoolExecutor, stop: threading.Event) -> None: + """Claim and run jobs until `stop` is set, then return once every in-flight job is done.""" + active: set = set() + while True: + active = {f for f in active if not f.done()} # reap finished jobs + if stop.is_set(): + if not active: + return + futures_wait(active, timeout=cfg.poll) + continue + # fill idle capacity, FIFO; re-check stop per claim so a SIGTERM mid-pass takes nothing new + while len(active) < cfg.max_concurrent and not stop.is_set(): + try: + job = _claim(client, cfg) + except httpx.HTTPError as e: + print(f"claim failed: {e}") + break + if not job: + break # queue empty + print(f"claimed job {job['id']} → {job['short']} ({len(active) + 1}/{cfg.max_concurrent} active)") + active.add(pool.submit(_run_job_guarded, client, cfg, job)) + stop.wait(cfg.poll) if __name__ == "__main__": diff --git a/orchestrator/tests/test_runner.py b/orchestrator/tests/test_runner.py index 34cbe1b..758491e 100644 --- a/orchestrator/tests/test_runner.py +++ b/orchestrator/tests/test_runner.py @@ -157,3 +157,79 @@ def test_max_concurrent_defaults_and_reads_env(monkeypatch): assert runner.Config().max_concurrent == 3 # default = staging pool size monkeypatch.setenv("RUNNER_MAX_CONCURRENT", "5") assert runner.Config().max_concurrent == 5 + + +# --- graceful drain: SIGTERM stops claims but lets in-flight jobs finish --- + +class _PollCfg(_Cfg): + poll = 0.05 + max_concurrent = 3 + + +def test_serve_on_stop_claims_nothing_more_and_waits_for_in_flight_job(): + """launchd SIGTERMs the runner on every restart (cert renewal, puppet reload). Returning + while a job runs would let launchd kill its `reprovision` subprocess mid-EACS.""" + import threading + from concurrent.futures import ThreadPoolExecutor + + started, release, stop = threading.Event(), threading.Event(), threading.Event() + jobs = iter([{"id": 1, "short": "macmini-m4-111"}]) + + def fake_job(*_a): + started.set() + release.wait(5) + + def fake_claim(*_a): + # Stop lands while the fill pass is still claiming: the next claim must not happen. + job = next(jobs, None) + if job is None: + started.wait(2) + stop.set() + return job + + with patch("orchestrator.runner._claim", side_effect=fake_claim) as claim, \ + patch("orchestrator.runner._run_job_guarded", side_effect=fake_job), \ + ThreadPoolExecutor(max_workers=3) as pool: + t = threading.Thread(target=runner._serve, args=(MagicMock(), _PollCfg(), pool, stop)) + t.start() + assert started.wait(2) + stop.wait(2) + claims_at_stop = claim.call_count + t.join(0.3) + assert t.is_alive(), "returned while a job was still running" + assert claim.call_count == claims_at_stop, "claimed new work after stop" + release.set() + t.join(2) + assert not t.is_alive(), "did not return after the in-flight job finished" + + +def test_serve_returns_promptly_when_stopped_and_idle(): + import threading + from concurrent.futures import ThreadPoolExecutor + + stop = threading.Event() + stop.set() + with patch("orchestrator.runner._claim") as claim, ThreadPoolExecutor(max_workers=1) as pool: + runner._serve(MagicMock(), _PollCfg(), pool, stop) + claim.assert_not_called() + + +def test_main_installs_a_sigterm_drain_handler(monkeypatch): + import signal + + monkeypatch.setenv("HANGAR_API_URL", "http://hangar/api") + monkeypatch.setenv("REPROVISION_RUNNER_TOKEN", "t") + previous = signal.getsignal(signal.SIGTERM) + seen = {} + + def fake_serve(_client, _cfg, _pool, stop): + signal.getsignal(signal.SIGTERM)(signal.SIGTERM, None) + seen["stopped"] = stop.is_set() + + try: + with patch("orchestrator.runner._serve", side_effect=fake_serve), \ + patch("orchestrator.runner.httpx.Client"): + runner.main() + finally: + signal.signal(signal.SIGTERM, previous) + assert seen["stopped"] is True