kstyagi
3e17e3e72f
- Added new API endpoints for tool catalog and installation management. - Updated usage documentation to include new tool installation features. - Implemented asynchronous tool installation with job tracking. - Introduced integration tests for tool installation in disposable VMs. - Enhanced dashboard UI to reflect tool installation status and usage instructions. - Improved error handling and logging for tool installation processes. - Added support for multiple Python versions and Conda environments. - Refactored tool management logic to ensure better state handling and job persistence.
339 lines
16 KiB
Python
339 lines
16 KiB
Python
"""Bounded background installation workers using the existing guest exec channel."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import fcntl
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import shlex
|
||
import stat
|
||
import threading
|
||
import time
|
||
from functools import lru_cache
|
||
from pathlib import Path
|
||
|
||
from . import guest_tools
|
||
from .errors import ConflictError, GuestAgentUnavailableError, GuestExecutionError, NotFoundError, UvmError, ValidationError
|
||
from .guest_exec import validate_exec
|
||
from .system import require_root
|
||
from .tool_catalog import all_selections, normalize_tools, recipe
|
||
from .tool_store import ToolJobStore
|
||
|
||
|
||
@lru_cache(maxsize=1)
|
||
def _guest_source():
|
||
return Path(guest_tools.__file__).read_text(encoding="utf-8")
|
||
|
||
|
||
def guest_command(payload: dict) -> str:
|
||
command = shlex.join(["/usr/bin/python3", "-I", "-c", _guest_source(), json.dumps(payload, ensure_ascii=True)])
|
||
validate_exec(command, True, 60)
|
||
return command
|
||
|
||
|
||
def _identity(vm):
|
||
return vm.pid, vm.process_start_time, vm.boot_id
|
||
|
||
|
||
def _text(value, limit=16384):
|
||
return str(value or "").encode("utf-8", errors="replace")[-limit:].decode("utf-8", errors="replace")
|
||
|
||
|
||
class ToolService:
|
||
def __init__(self, settings, state_store, lifecycle) -> None:
|
||
self.settings = settings
|
||
self.state = state_store
|
||
self.lifecycle = lifecycle
|
||
self.store = ToolJobStore(settings, state_store)
|
||
self._mutex = threading.RLock()
|
||
self._wake = threading.Condition()
|
||
self._stop = threading.Event()
|
||
self._threads = []
|
||
self._sessions = {}
|
||
self._owner = None
|
||
self._workers_left = 0
|
||
|
||
def start(self):
|
||
with self._mutex:
|
||
if self._stop.is_set():
|
||
raise ConflictError("the tool installer is shutting down")
|
||
if self._owner is not None:
|
||
return
|
||
with self.state.operation_lock():
|
||
descriptor = os.open(self.settings.base / "tool-worker.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
|
||
try:
|
||
metadata = os.fstat(descriptor)
|
||
if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1 or metadata.st_uid != os.geteuid():
|
||
raise UvmError("unsafe tool worker ownership file")
|
||
os.fchmod(descriptor, 0o600)
|
||
fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
except BaseException as error:
|
||
os.close(descriptor)
|
||
if isinstance(error, BlockingIOError):
|
||
raise ConflictError("another management server is already processing tool installations for this registry") from error
|
||
raise
|
||
self._owner = descriptor
|
||
try:
|
||
self.store.recover()
|
||
for index in range(self.settings.tool_workers):
|
||
thread = threading.Thread(target=self._worker, name=f"uvm-tools-{index}", daemon=True)
|
||
self._workers_left += 1
|
||
try:
|
||
thread.start()
|
||
except BaseException:
|
||
self._workers_left -= 1
|
||
raise
|
||
self._threads.append(thread)
|
||
except BaseException:
|
||
self._stop.set()
|
||
self.notify()
|
||
if not self._workers_left:
|
||
self._release_owner()
|
||
raise
|
||
|
||
def _release_owner(self):
|
||
if self._owner is not None:
|
||
os.close(self._owner)
|
||
self._owner = None
|
||
|
||
def notify(self):
|
||
with self._wake:
|
||
self._wake.notify_all()
|
||
|
||
def shutdown(self):
|
||
self._stop.set()
|
||
self.notify()
|
||
with self._mutex:
|
||
sessions = list(self._sessions.values())
|
||
for session in sessions:
|
||
session.close()
|
||
deadline = time.monotonic() + 8
|
||
for thread in self._threads:
|
||
thread.join(max(0, deadline - time.monotonic()))
|
||
# A worker finishing a host start keeps ownership until that start ends.
|
||
# Its guest command is never submitted after shutdown was requested.
|
||
with self._mutex:
|
||
if not self._workers_left:
|
||
self._release_owner()
|
||
|
||
def submit(self, vm_id, selections, *, start_if_stopped=False, restart=False, request_id=None):
|
||
require_root()
|
||
selections = normalize_tools(selections)
|
||
if type(start_if_stopped) is not bool or type(restart) is not bool:
|
||
raise ValidationError("start_if_stopped and restart must be booleans")
|
||
if request_id is not None and (not isinstance(request_id, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,80}", request_id)):
|
||
raise ValidationError("request_id must contain 1–80 letters, numbers, underscores, or hyphens")
|
||
self.start()
|
||
job = self.lifecycle.queue_tools(vm_id, selections, start_if_stopped=start_if_stopped,
|
||
restart=restart, request_id=request_id)
|
||
self.notify()
|
||
return job
|
||
|
||
def _execute(self, vm, payload, timeout):
|
||
if self._stop.is_set():
|
||
raise InterruptedError("management server is shutting down")
|
||
thread_id = threading.get_ident()
|
||
|
||
def connected(session):
|
||
with self._mutex:
|
||
self._sessions[thread_id] = session
|
||
if self._stop.is_set():
|
||
session.close()
|
||
|
||
try:
|
||
result = self.lifecycle.execute(vm.id, guest_command(payload), sudo=True, timeout_seconds=timeout,
|
||
expected_identity=_identity(vm), on_connect=connected)
|
||
finally:
|
||
with self._mutex:
|
||
self._sessions.pop(thread_id, None)
|
||
if result.timed_out:
|
||
return {"ok": False, "error": f"installation/verification timed out after {timeout:g} seconds; retry to finish it",
|
||
"output": _text(result.stderr), "truncated": result.truncated}
|
||
try:
|
||
value = json.loads(result.stdout)
|
||
if result.exit_code != 0 or result.truncated or not isinstance(value, dict) or type(value.get("ok")) is not bool:
|
||
raise ValueError("invalid result")
|
||
return value
|
||
except (ValueError, TypeError) as error:
|
||
raise GuestExecutionError("guest returned an incomplete tool result; refresh status before retrying",
|
||
outcome_unknown=payload["action"] == "install") from error
|
||
|
||
def _payload(self, vm, action, selections, *, boot_id=None):
|
||
return {"action": action, "username": vm.username, "architecture": self.settings.architecture,
|
||
"boot_id": boot_id, "tools": [recipe(selection, self.settings.architecture) for selection in selections]}
|
||
|
||
@staticmethod
|
||
def _inventory(result, selections):
|
||
expected = {item.key: item for item in selections}
|
||
records = result.get("tools")
|
||
if not isinstance(records, list) or len(records) != len(expected):
|
||
raise GuestExecutionError("invalid guest tool inventory")
|
||
inventory = []
|
||
for value in records:
|
||
if (not isinstance(value, dict) or value.get("key") not in expected
|
||
or type(value.get("installed")) is not bool):
|
||
raise GuestExecutionError("invalid guest tool inventory")
|
||
selection = expected.pop(value["key"])
|
||
inventory.append({**selection.to_dict(), "key": selection.key, "installed": value["installed"],
|
||
"detected_version": _text(value.get("detected_version"), 120) or None,
|
||
"error": _text(value.get("error"), 512) or None,
|
||
"usage": _text(value.get("usage"), 512)})
|
||
return inventory
|
||
|
||
def snapshot(self, vm_id, *, refresh=False):
|
||
require_root()
|
||
vm = self.state.get(vm_id)
|
||
data = self.store.read(vm_id)
|
||
running = self.lifecycle.inspect(vm.id).observed_status == "running"
|
||
if refresh and running:
|
||
supported = data["supported"]
|
||
try:
|
||
result = self._execute(vm, self._payload(vm, "probe", all_selections()), 90)
|
||
if type(result.get("supported")) is bool:
|
||
supported = result["supported"]
|
||
if not result["ok"]:
|
||
raise UvmError(_text(result.get("error"), 2048))
|
||
inventory = self._inventory(result, all_selections())
|
||
with self.store.transaction(vm_id) as current:
|
||
current.update(inventory=inventory, checked_at=time.time(), error=None, agent_ready=True,
|
||
guest_busy=result.get("busy") is True, supported=result.get("supported") is True)
|
||
data = self.store.read(vm_id)
|
||
except UvmError as error:
|
||
with self.store.transaction(vm_id) as current:
|
||
current.update(error=_text(error, 2048), agent_ready=not isinstance(error, GuestAgentUnavailableError),
|
||
supported=supported)
|
||
data = self.store.read(vm_id)
|
||
if not running:
|
||
data.update(agent_ready=False, guest_busy=False)
|
||
return {"vm_id": vm_id, **{key: data[key] for key in (
|
||
"inventory", "checked_at", "agent_ready", "guest_busy", "supported", "error")},
|
||
"stale": not running or data["checked_at"] is None or time.time() - data["checked_at"] > 60,
|
||
"latest_job": data["jobs"][-1] if data["jobs"] else None,
|
||
"jobs": [{key: job[key] for key in ("id", "status", "created_at", "completed_at")}
|
||
for job in reversed(data["jobs"])]}
|
||
|
||
def _worker(self):
|
||
try:
|
||
while not self._stop.is_set():
|
||
try:
|
||
job = self.store.claim()
|
||
if job is not None:
|
||
self._run_job(job)
|
||
continue
|
||
except (UvmError, OSError):
|
||
logging.getLogger(__name__).exception("Could not process a tool installation job")
|
||
with self._wake:
|
||
self._wake.wait(timeout=1)
|
||
finally:
|
||
with self._mutex:
|
||
self._workers_left -= 1
|
||
if self._workers_left == 0 and self._stop.is_set():
|
||
self._release_owner()
|
||
|
||
def _update(self, job, update):
|
||
return self.store.update_job(job["vm_id"], job["id"], update)
|
||
|
||
def _wait_ready(self, vm, selections):
|
||
deadline = time.monotonic() + self.settings.tool_ready_timeout_s
|
||
while not self._stop.is_set():
|
||
try:
|
||
result = self._execute(vm, self._payload(vm, "probe", selections), 90)
|
||
if not result["ok"]:
|
||
raise UvmError(_text(result.get("error"), 2048))
|
||
if result.get("busy"):
|
||
raise ConflictError("an earlier installation is still running inside this VM; wait and refresh before retrying")
|
||
if not isinstance(result.get("boot_id"), str) or not result["boot_id"]:
|
||
raise GuestExecutionError("guest did not report its boot identity")
|
||
return result
|
||
except GuestAgentUnavailableError:
|
||
if time.monotonic() >= deadline:
|
||
raise GuestAgentUnavailableError("Guest agent did not become ready. Restart & retry to apply agent setup, "
|
||
"then check the VM log if startup still fails.")
|
||
self._stop.wait(min(2, max(0, deadline - time.monotonic())))
|
||
raise InterruptedError("management server is shutting down")
|
||
|
||
def _run_job(self, job):
|
||
submitted = False
|
||
try:
|
||
if self._stop.is_set():
|
||
raise InterruptedError("management server is shutting down")
|
||
selections = normalize_tools(job["selections"])
|
||
vm = self.lifecycle.prepare_tool_install(job["vm_id"], start_if_stopped=job["start_if_stopped"],
|
||
restart=job.get("restart", False))
|
||
self._update(job, lambda current, data: current.update(status="waiting"))
|
||
info = self._wait_ready(vm, selections)
|
||
initial = self._inventory(info, selections)
|
||
|
||
def ready(current, data):
|
||
observed = {item["key"]: item for item in data["inventory"]}
|
||
observed.update({item["key"]: item for item in initial})
|
||
data.update(inventory=list(observed.values()), agent_ready=True,
|
||
guest_busy=False, supported=True, error=None)
|
||
self._update(job, ready)
|
||
for index, selection in enumerate(selections):
|
||
if self._stop.is_set():
|
||
raise InterruptedError("management server is shutting down")
|
||
|
||
def begin(current, data):
|
||
current["status"] = "installing"
|
||
current["items"][index]["status"] = "installing"
|
||
self._update(job, begin)
|
||
submitted = True
|
||
result = self._execute(vm, self._payload(vm, "install", [selection], boot_id=info["boot_id"]),
|
||
self.settings.tool_install_timeout_s)
|
||
observed = None
|
||
if result["ok"]:
|
||
observed, = self._inventory({"tools": [result.get("tool")]}, [selection])
|
||
if not observed["installed"] or result.get("status") not in {"installed", "already_installed"}:
|
||
raise GuestExecutionError("guest installation did not verify the selected tool", outcome_unknown=True)
|
||
|
||
def finish_item(current, data):
|
||
current["items"][index].update(
|
||
status=result["status"] if result["ok"] else "failed",
|
||
detected_version=observed["detected_version"] if observed else None,
|
||
error=None if result["ok"] else _text(result.get("error"), 2048),
|
||
output=_text(result.get("output")), truncated=result.get("truncated") is True)
|
||
inventory = {item["key"]: item for item in data["inventory"]}
|
||
if observed:
|
||
inventory[selection.key] = observed
|
||
elif selection.key in inventory:
|
||
inventory[selection.key].update(installed=False, error="Refresh to verify after the failed installation.")
|
||
data.update(inventory=list(inventory.values()))
|
||
self._update(job, finish_item)
|
||
submitted = False
|
||
|
||
def finish(current, data):
|
||
failures = sum(item["status"] == "failed" for item in current["items"])
|
||
current.update(status="failed" if failures else "succeeded", completed_at=time.time(),
|
||
error=f"{failures} tool installation(s) failed; see the per-tool output and retry." if failures else None)
|
||
self._update(job, finish)
|
||
except NotFoundError:
|
||
return # Deletion owns removal of the whole runtime, including history.
|
||
except (InterruptedError, GuestExecutionError, ConflictError) as error:
|
||
try:
|
||
self._update(job, lambda current, data: self.store.interrupt(
|
||
current, _text(error, 2048), time.time(),
|
||
outcome_unknown=getattr(error, "outcome_unknown", submitted)))
|
||
except NotFoundError:
|
||
pass
|
||
except (UvmError, OSError, ValueError) as error:
|
||
def failed(current, data):
|
||
self.store.interrupt(current, _text(error, 2048), time.time(), outcome_unknown=submitted)
|
||
current["status"] = "failed" if not submitted else "interrupted"
|
||
if isinstance(error, GuestAgentUnavailableError):
|
||
data.update(agent_ready=False, error=_text(error, 2048))
|
||
try:
|
||
self._update(job, failed)
|
||
except NotFoundError:
|
||
pass
|
||
except Exception:
|
||
logging.getLogger(__name__).exception("Unexpected tool installation failure")
|
||
try:
|
||
self._update(job, lambda current, data: self.store.interrupt(
|
||
current, "Unexpected installer failure; check the management server log and refresh before retrying.",
|
||
time.time(), outcome_unknown=submitted))
|
||
except NotFoundError:
|
||
pass
|