111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
|
"""Detached process host for one project-bound DocForge visualization."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import socket
|
||
|
|
import time
|
||
|
|
from contextlib import suppress
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from .errors import DocForgeError
|
||
|
|
from .visualization import VisualizationRunner
|
||
|
|
|
||
|
|
|
||
|
|
def _parser() -> argparse.ArgumentParser:
|
||
|
|
parser = argparse.ArgumentParser(prog="docforge-visualization-worker")
|
||
|
|
parser.add_argument("--control-fd", type=int, required=True)
|
||
|
|
return parser
|
||
|
|
|
||
|
|
|
||
|
|
def _read_request(control: socket.socket) -> dict[str, Any]:
|
||
|
|
chunks: list[bytes] = []
|
||
|
|
size = 0
|
||
|
|
while True:
|
||
|
|
chunk = control.recv(65536)
|
||
|
|
if not chunk:
|
||
|
|
break
|
||
|
|
chunks.append(chunk)
|
||
|
|
size += len(chunk)
|
||
|
|
if size > 1_000_000:
|
||
|
|
raise ValueError("Visualization launch request exceeded its fixed boundary")
|
||
|
|
if b"\n" in chunk:
|
||
|
|
break
|
||
|
|
payload = b"".join(chunks).split(b"\n", 1)[0]
|
||
|
|
request = json.loads(payload)
|
||
|
|
if not isinstance(request, dict):
|
||
|
|
raise ValueError("Visualization launch request is invalid")
|
||
|
|
return request
|
||
|
|
|
||
|
|
|
||
|
|
def _pid_exists(pid: int) -> bool:
|
||
|
|
if pid <= 1:
|
||
|
|
return False
|
||
|
|
try:
|
||
|
|
os.kill(pid, 0)
|
||
|
|
except ProcessLookupError:
|
||
|
|
return False
|
||
|
|
except PermissionError:
|
||
|
|
return True
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> int:
|
||
|
|
arguments = _parser().parse_args(argv)
|
||
|
|
control = socket.socket(fileno=arguments.control_fd)
|
||
|
|
runner: VisualizationRunner | None = None
|
||
|
|
try:
|
||
|
|
request = _read_request(control)
|
||
|
|
target = request["target"]
|
||
|
|
if not isinstance(target, dict):
|
||
|
|
raise ValueError("Visualization target is invalid")
|
||
|
|
runner = VisualizationRunner(
|
||
|
|
None,
|
||
|
|
snapshot_spec=request["snapshot"],
|
||
|
|
token=str(request["token"]),
|
||
|
|
register_atexit=False,
|
||
|
|
initial_grace_seconds=float(request["initial_grace_seconds"]),
|
||
|
|
lease_seconds=float(request["lease_seconds"]),
|
||
|
|
monitor_interval_seconds=float(request["monitor_interval_seconds"]),
|
||
|
|
)
|
||
|
|
visualization = runner.start(
|
||
|
|
node_id=target.get("node_id"),
|
||
|
|
query=target.get("query"),
|
||
|
|
depth=int(target["depth"]),
|
||
|
|
)
|
||
|
|
control.sendall(
|
||
|
|
json.dumps(
|
||
|
|
{"status": "ok", "visualization": visualization},
|
||
|
|
sort_keys=True,
|
||
|
|
separators=(",", ":"),
|
||
|
|
).encode("utf-8")
|
||
|
|
+ b"\n"
|
||
|
|
)
|
||
|
|
owner_pid = int(request["owner_pid"])
|
||
|
|
except (DocForgeError, KeyError, OSError, TypeError, ValueError) as error:
|
||
|
|
with suppress(OSError):
|
||
|
|
control.sendall(
|
||
|
|
json.dumps(
|
||
|
|
{"status": "error", "error": type(error).__name__},
|
||
|
|
sort_keys=True,
|
||
|
|
separators=(",", ":"),
|
||
|
|
).encode("utf-8")
|
||
|
|
+ b"\n"
|
||
|
|
)
|
||
|
|
if runner is not None:
|
||
|
|
runner.stop()
|
||
|
|
return 2
|
||
|
|
finally:
|
||
|
|
control.close()
|
||
|
|
|
||
|
|
while runner.is_running() and _pid_exists(owner_pid):
|
||
|
|
time.sleep(0.25)
|
||
|
|
runner.stop()
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|