runner: drain in-flight jobs on SIGTERM instead of dying

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__":

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