"""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() self._workers_left = self.settings.tool_workers self._threads = [threading.Thread(target=self._worker, name=f"uvm-tools-{index}", daemon=True) for index in range(self.settings.tool_workers)] for thread in self._threads: thread.start() except BaseException: 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: try: result = self._execute(vm, self._payload(vm, "probe", all_selections()), 90) 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)) 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()), checked_at=time.time(), 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()), checked_at=time.time()) 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