kstyagi
e1b934e35e
- Added `guest_tools.py` for standalone tool installer functionality. - Introduced `ToolJobStore` in `tool_store.py` to manage per-VM tool job records. - Created `ToolService` in `tools.py` for background installation workers using the guest execution channel. - Enhanced `LifecycleService` in `lifecycle.py` to support tool installation queuing and preparation. - Updated `guest_exec.py` to handle connection shutdown gracefully. - Added validation and normalization functions for tool specifications. - Implemented logging and error handling for tool installation processes. - Ensured compatibility with Ubuntu 22.04 and ARM64/x86_64 architectures.
320 lines
16 KiB
Python
320 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()
|
||
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
|