__init__

uvm/lifecycle.py

@@ -0,0 +1,284 @@
+"""Ordered VM lifecycle operations and compensation for partial failures."""
+
+from __future__ import annotations
+
+import shutil
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+from .config import Settings
+from .domain import VmRecord, VmSpec, new_vm_id
+from .errors import TapCreationError, UvmError
+from .firecracker.api import FirecrackerClient
+from .firecracker.config import build_config, write_config
+from .firecracker.process import FirecrackerProcessManager
+from .images import ImageStore
+from .integrity import load_manifest, verify_file
+from .network import NetworkManager
+from .state import StateStore
+from .system import check_kvm, ensure_data_directories, require_root
+
+
+@dataclass(frozen=True, slots=True)
+class ListedVm:
+    vm: VmRecord
+    observed_status: str
+
+
+class LifecycleService:
+    """The only module allowed to coordinate multiple VM host resources."""
+
+    def __init__(
+        self,
+        settings: Settings,
+        state_store: StateStore,
+        images: ImageStore,
+        network: NetworkManager,
+        process: FirecrackerProcessManager,
+        *,
+        client_factory: Callable[[Path, float], FirecrackerClient] = FirecrackerClient,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        self._settings = settings
+        self._state_store = state_store
+        self._images = images
+        self._network = network
+        self._process = process
+        self._client_factory = client_factory
+        self._clock = clock
+
+    def create(self, spec: VmSpec) -> VmRecord:
+        """Reserve identity first, then create all VM resources in dependency order."""
+
+        require_root()
+        ensure_data_directories(self._settings)
+        check_kvm()
+        if not self._settings.firecracker_binary.exists():
+            raise UvmError("Firecracker is not installed. Run: sudo uvm install")
+        with self._state_store.operation_lock():
+            manifest = load_manifest(self._settings.integrity_manifest_path)
+            firecracker_checksum = self._settings.firecracker_binary_sha256
+            if manifest.verified or self._settings.allow_unverified_downloads:
+                firecracker_checksum = firecracker_checksum or manifest.checksums.get("firecracker")
+            verify_file(
+                self._settings.firecracker_binary,
+                firecracker_checksum,
+                "Firecracker binary",
+                "UVM_FIRECRACKER_BINARY_SHA256",
+                allow_unverified=self._settings.allow_unverified_downloads,
+            )
+            return self._create_locked(spec)
+
+    def _create_locked(self, spec: VmSpec) -> VmRecord:
+        assets = self._images.installed_assets()
+        vm = self._reserve_vm(spec)
+
+        tap_created = False
+        try:
+            runtime_dir = self._settings.vm_dir(vm.id)
+            runtime_dir.mkdir(parents=True, exist_ok=False)
+            runtime_dir.chmod(0o700)
+            disk = self._images.create_vm_disk(assets.rootfs, Path(vm.disk))
+            self._images.provision_credentials(disk, vm.username, vm.password)
+            self._network.ensure_bridge()
+            try:
+                self._network.create_tap(vm.tap)
+            except TapCreationError as error:
+                tap_created = error.tap_created
+                raise
+            else:
+                tap_created = True
+
+            config = build_config(self._settings, vm, assets.kernel, disk)
+            write_config(Path(vm.config), config)
+
+            process_info = self._process.start(vm)
+            vm.pid = process_info.pid
+            vm.process_start_time = process_info.start_time
+            vm.updated_at = self._now()
+            self._replace_vm(vm)
+
+            client = self._client_factory(Path(vm.socket), self._settings.api_timeout_s)
+            client.configure_and_start(config)
+        except BaseException as error:
+            cleanup_errors = self._rollback_create(vm, tap_created)
+            message = f"failed to create {vm.id}: {error}"
+            if cleanup_errors:
+                message = f"{message}; cleanup failed: {'; '.join(cleanup_errors)}"
+                vm.status = "failed"
+                vm.last_error = message
+                vm.updated_at = self._now()
+                self._replace_vm_if_present(vm)
+            else:
+                self._remove_vm_if_present(vm.id)
+            if isinstance(error, (KeyboardInterrupt, SystemExit)):
+                raise
+            if isinstance(error, UvmError):
+                raise error
+            raise UvmError(message) from error
+
+        vm.status = "running"
+        vm.last_error = None
+        vm.updated_at = self._now()
+        self._replace_vm(vm)
+        return vm
+
+    def list_vms(self) -> list[ListedVm]:
+        state = self._state_store.load()
+        listed: list[ListedVm] = []
+        for vm in state.vms.values():
+            observed_status = vm.status
+            if vm.status in {"starting", "running", "stopping", "terminating"}:
+                if not self._process.is_alive(vm):
+                    observed_status = "dead"
+            listed.append(ListedVm(vm=vm, observed_status=observed_status))
+        return listed
+
+    def find_for_ssh(self, identifier: str) -> VmRecord:
+        return self._state_store.find_by_id_or_ip(identifier)
+
+    def stop(self, vm_id: str) -> VmRecord:
+        """Stop one VM while retaining its private writable disk for future use."""
+
+        require_root()
+        with self._state_store.operation_lock():
+            return self._stop_locked(vm_id)
+
+    def _stop_locked(self, vm_id: str) -> VmRecord:
+        vm = self._begin_transition(vm_id, "stopping")
+        try:
+            self._process.terminate(vm)
+            self._network.delete_tap(vm.tap)
+        except BaseException as error:
+            vm.status = "failed"
+            vm.last_error = f"failed to stop VM: {error}"
+            vm.updated_at = self._now()
+            self._replace_vm_if_present(vm)
+            if isinstance(error, (KeyboardInterrupt, SystemExit)):
+                raise
+            if isinstance(error, UvmError):
+                raise error
+            raise UvmError(vm.last_error) from error
+
+        vm.status = "stopped"
+        vm.pid = None
+        vm.process_start_time = None
+        vm.last_error = None
+        vm.updated_at = self._now()
+        self._replace_vm(vm)
+        return vm
+
+    def destroy(self, vm_id: str) -> VmRecord:
+        """Stop a VM, remove its private resources, then release its allocation."""
+
+        require_root()
+        with self._state_store.operation_lock():
+            return self._destroy_locked(vm_id)
+
+    def _destroy_locked(self, vm_id: str) -> VmRecord:
+        vm = self._begin_transition(vm_id, "terminating")
+        try:
+            self._process.terminate(vm)
+            self._network.delete_tap(vm.tap)
+            try:
+                shutil.rmtree(self._settings.vm_dir(vm.id))
+            except FileNotFoundError:
+                pass
+        except BaseException as error:
+            vm.status = "failed"
+            vm.last_error = f"failed to destroy VM: {error}"
+            vm.updated_at = self._now()
+            self._replace_vm_if_present(vm)
+            if isinstance(error, (KeyboardInterrupt, SystemExit)):
+                raise
+            if isinstance(error, UvmError):
+                raise error
+            raise UvmError(vm.last_error) from error
+
+        self._remove_vm_if_present(vm.id)
+        return vm
+
+    def _reserve_vm(self, spec: VmSpec) -> VmRecord:
+        with self._state_store.transaction() as state:
+            vm_id = new_vm_id()
+            used_taps = {existing_vm.tap for existing_vm in state.vms.values()}
+            while vm_id in state.vms or self._network.tap_name(vm_id) in used_taps:
+                vm_id = new_vm_id()
+            guest_ip = self._network.allocate_ip(state.vms.values(), spec.guest_ip)
+            mac = self._network.mac_for(state.next_mac_index)
+            state.next_mac_index += 1
+            now = self._now()
+            runtime_dir = self._settings.vm_dir(vm_id)
+            vm = VmRecord(
+                id=vm_id,
+                cpu=spec.cpu,
+                ram_mib=spec.ram_mib,
+                guest_ip=str(guest_ip),
+                gateway=str(self._settings.gateway),
+                tap=self._network.tap_name(vm_id),
+                mac=mac,
+                socket=str(runtime_dir / "firecracker.sock"),
+                config=str(runtime_dir / "config.json"),
+                log=str(runtime_dir / "firecracker.log"),
+                disk=str(runtime_dir / "rootfs.ext4"),
+                username=spec.username,
+                password=spec.password,
+                status="starting",
+                created_at=now,
+                updated_at=now,
+            )
+            state.vms[vm.id] = vm
+        return vm
+
+    def _begin_transition(self, vm_id: str, target_status: str) -> VmRecord:
+        with self._state_store.transaction() as state:
+            try:
+                vm = state.vms[vm_id]
+            except KeyError as error:
+                raise UvmError(f"VM not found: {vm_id}") from error
+            if vm.status in {"stopping", "terminating"}:
+                raise UvmError(f"VM operation is already in progress: {vm_id}")
+            vm.status = target_status
+            vm.updated_at = self._now()
+            state.vms[vm_id] = vm
+            return vm
+
+    def _replace_vm(self, vm: VmRecord) -> None:
+        with self._state_store.transaction() as state:
+            if vm.id not in state.vms:
+                raise UvmError(f"VM not found: {vm.id}")
+            state.vms[vm.id] = vm
+
+    def _replace_vm_if_present(self, vm: VmRecord) -> None:
+        with self._state_store.transaction() as state:
+            if vm.id in state.vms:
+                state.vms[vm.id] = vm
+
+    def _remove_vm_if_present(self, vm_id: str) -> None:
+        with self._state_store.transaction() as state:
+            state.vms.pop(vm_id, None)
+
+    def _rollback_create(self, vm: VmRecord, tap_created: bool) -> list[str]:
+        errors: list[str] = []
+        if vm.pid is not None:
+            try:
+                self._process.terminate(vm)
+            except Exception as error:
+                errors.append(f"process: {error}")
+        if tap_created:
+            try:
+                self._network.delete_tap(vm.tap)
+            except Exception as error:
+                errors.append(f"network: {error}")
+        try:
+            shutil.rmtree(self._settings.vm_dir(vm.id))
+        except FileNotFoundError:
+            pass
+        except OSError as error:
+            errors.append(f"files: {error}")
+        return errors
+
+    def _now(self) -> int:
+        return int(self._clock())