"""Image-resident queue workers. Requires strongly consistent S3 GET/PUT/LIST. Per-task Lamport bakery tickets use only ordinary object operations. A choosing marker precedes ticket selection; contenders wait for choosing peers and lower (ticket, worker) pairs. Late arrivals cannot displace a running owner. Leases assume bounded clock skew (less than MARGIN / 2) and a worker killed by its deadline. """ from __future__ import annotations import argparse import contextlib import ctypes import io import json import logging import multiprocessing import os import signal import socket import tarfile import tempfile import threading import time import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from engine_trace_worker import unit_record, worker from openultrasast.plane.memory import S3Client, s3_settings PENDING = "engine-queue/pending/" CLAIMED = "engine-queue/claimed/" RESULTS = "engine-results/" MARGIN = 60 LOG = logging.getLogger(__name__) def client_from_env(): settings = s3_settings(os.environ) bucket = os.environ.get("S3_BUCKET") if not bucket: raise ValueError("S3_BUCKET is required") return S3Client(bucket=bucket, **settings) def read(client, key): value = client.get(key) return json.loads(value[0]) if value else None def put(client, key, value): client.put(key, json.dumps(value).encode(), {}) def result_key(task): return f"{RESULTS}{task['run']}/{task['id']}.json" def pack(root): stream = io.BytesIO() with tarfile.open(fileobj=stream, mode="w") as archive: archive.add(root, arcname=".") return stream.getvalue() def unpack(data, root): with tarfile.open(fileobj=io.BytesIO(data)) as archive: # Source snapshots need only regular files and directories, never links/devices. for member in archive.getmembers(): target = (root / member.name).resolve() if not target.is_relative_to(root.resolve()) or not (member.isfile() or member.isdir()): raise ValueError("unsafe source archive") archive.extractall(root, filter="data") class Queue: def __init__(self, client, identity=None, *, now=time.time, pause=time.sleep, stop=None): self.client = client self.identity = identity or f"{socket.gethostname()}-{uuid.uuid4().hex}" self.now, self.pause = now, pause self.stop = stop or threading.Event() def markers(self, task_id=None): result = [] for key in self.client.keys(CLAIMED): if task_id is None or key.rsplit("/", 1)[-1] == task_id: marker = read(self.client, key) if marker: result.append((key, marker)) return result def tasks(self): tasks = {} for key in self.client.keys(PENDING): task = read(self.client, key) if task: tasks[task["task_id"]] = task # A dead worker may have removed pending already. Its task remains in the lease. for _, marker in self.markers(): if marker["expires"] < self.now(): task = marker["task"] tasks.setdefault(task["task_id"], task) return list(tasks.values()) def claim(self, task): if read(self.client, result_key(task)) is not None or self.stop.is_set(): return None key = f"{CLAIMED}{self.identity}/{task['task_id']}" marker = dict(task=task, ticket=0, expires=self.now() + task["deadline"] + MARGIN) put(self.client, key, marker) try: live = [m for _, m in self.markers(task["task_id"]) if m["expires"] > self.now()] marker["ticket"] = 1 + max((m["ticket"] for m in live), default=0) put(self.client, key, marker) while not self.stop.is_set(): # Do not start once the execution budget no longer fits within our lease. if self.now() >= marker["expires"] - task["deadline"] - MARGIN / 2: break peers = [(k, m) for k, m in self.markers(task["task_id"]) if k != key and m["expires"] > self.now()] if any(m["ticket"] == 0 for _, m in peers): self.pause(0.1) continue if any((m["ticket"], k) < (marker["ticket"], key) for k, m in peers): break # loser backs off and can try another pending task if read(self.client, key) != marker or read(self.client, result_key(task)) is not None: break self.client.delete(PENDING + task["task_id"]) return key except Exception: self.client.delete(key) raise self.client.delete(key) return None def process_one(self, execute=None): for task in self.tasks(): key = self.claim(task) if key is None: continue try: lease = read(self.client, key) if lease is None: continue budget = min(task["deadline"], lease["expires"] - self.now() - MARGIN / 2) if budget <= 0: continue # Leave the expired lease's task available to recovery. result = (execute or execute_task)(self.client, {**task, "runtime_budget": budget}) result.update(executor="queue", worker=self.identity, pod=socket.gethostname(), image=os.environ.get("ENGINE_IMAGE", "")) put(self.client, result_key(task), result) # Publish the result before removing either the input or the claim. self.client.delete(task["input"]) for old_key, old in self.markers(task["task_id"]): if old_key == key or old["expires"] < self.now(): self.client.delete(old_key) return True except Exception: # Keep the lease and task for recovery, without logging credential-bearing errors. LOG.error("task failed; lease retained for recovery") return False return False def _child(pin, root, out, deadline, question_deadline): os.setsid() signal.signal(signal.SIGTERM, signal.SIG_DFL) worker(pin, root, out, deadline, question_deadline) def adopt_orphans(): # Linux PR_SET_CHILD_SUBREAPER: escaped Joern sessions reparent here, not # outside the service. There is exactly one task process tree per worker. if ctypes.CDLL(None, use_errno=True).prctl(36, 1, 0, 0, 0) != 0: raise OSError("cannot enable task descendant cleanup") def kill_children(): """Freeze parents before walking children, including new-session JVMs.""" seen = set() def freeze(parent): for entry in Path("/proc").iterdir(): if not entry.name.isdigit() or int(entry.name) in seen: continue try: # comm can contain spaces or parentheses. Fields after it begin at state. fields = (entry / "stat").read_text().rsplit(")", 1)[1].split() if int(fields[1]) != parent: continue pid = int(entry.name) os.kill(pid, signal.SIGSTOP) seen.add(pid) freeze(pid) except (FileNotFoundError, ProcessLookupError, PermissionError): continue # A child that exited during discovery can reparent descendants to this # subreaper. Repeat until every remaining descendant is frozen. while True: count = len(seen) freeze(os.getpid()) if len(seen) == count: break for pid in seen: with contextlib.suppress(ProcessLookupError): os.kill(pid, signal.SIGKILL) def execute_task(client, task): adopt_orphans() pin = task["pin"] started = time.monotonic() budget = task.get("runtime_budget", task["deadline"]) status, reason, code = "failed", "worker produced no result", None with tempfile.TemporaryDirectory(prefix="engine-task-") as scratch: root, out = Path(scratch) / "case", Path(scratch) / "result.json" root.mkdir() try: if task["deadline"] > 900 or task["deadline"] <= 0: raise ValueError("deadline exceeds deployment grace") if task.get("image") != os.environ.get("ENGINE_IMAGE"): raise ValueError("task image differs from deployed image") source = client.get(task["input"]) if source is None: raise ValueError("missing source") unpack(source[0], root) # The shipped worker also checks frontend input and graph output for every language. files = [p for p in root.rglob("*") if p.is_file()] if not files: raise ValueError("empty source") LOG.info("input proof: first file bytes=%d", len(files[0].read_bytes())) if time.monotonic() - started >= budget: status, reason = "timeout", "task input deadline" raise TimeoutError(reason) process = multiprocessing.get_context("fork").Process(target=_child, args=(pin, root, out, budget, task["question_deadline"])) process.start() try: process.join(max(0, budget - (time.monotonic() - started))) if process.is_alive(): status, reason = "timeout", "task deadline" else: reason = f"worker exit {process.exitcode}" code = process.exitcode finally: # Joern creates separate sessions. Kill the whole descendant tree, # not merely the Python wrapper's process group. kill_children() process.join() # Reap adopted JVMs and their launchers after the direct child. while True: try: os.waitpid(-1, 0) except ChildProcessError: break record = json.loads(out.read_text()) if out.exists() else {**pin, "units": []} except Exception: record = {**pin, "units": []} reason = "worker input or execution failed" if not record.get("done") or code != 0: previous = {u["unit"]: u for u in record["units"]} record["units"] = [ previous[u["unit"]] if code is None and previous.get(u["unit"], {}).get("status") in {"path", "asked-nothing"} else unit_record(u, status if u["supported"] else "unsupported", reason if u["supported"] else u["reason"]) for u in pin["units"] ] record.update(done=True, container_exit=code, seconds=time.monotonic() - started) return record def status(client): return { "pending": len(client.keys(PENDING)), "claimed": len(client.keys(CLAIMED)), "done": len([k for k in client.keys(RESULTS) if k.endswith(".json")]), } def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--port", type=int, default=8080) args = parser.parse_args() logging.basicConfig(level=logging.INFO) stop = threading.Event() signal.signal(signal.SIGTERM, lambda *_: stop.set()) signal.signal(signal.SIGINT, lambda *_: stop.set()) client = client_from_env() queue = Queue(client, stop=stop) class Health(BaseHTTPRequestHandler): def do_GET(self): self.send_response(200 if self.path == "/healthz" and not stop.is_set() else 503) self.end_headers() def log_message(self, *_): pass server = ThreadingHTTPServer(("0.0.0.0", args.port), Health) threading.Thread(target=server.serve_forever, daemon=True).start() try: while not stop.is_set(): try: if not queue.process_one(): stop.wait(2) except Exception: LOG.error("queue unavailable; retrying") stop.wait(5) finally: server.shutdown() server.server_close() if __name__ == "__main__": main()