Files
nest/uvm/domain.py
T
kstyagi fb80ce7c8a feat: Add storage management for VM creation and resizing
- Introduced storage parameter in VmCreateRequest and VmResizeRequest models.
- Implemented storage validation to ensure minimum and maximum limits.
- Enhanced CLI to support storage options during VM creation and resizing.
- Updated lifecycle management to handle pending resize operations for VM resources.
- Added methods for resizing VM disks and validating storage capacity.
- Updated API endpoints to accommodate new storage functionality.
- Refactored related validation and configuration logic to support storage management.
2026-09-21 19:16:33 +05:30

327 lines
12 KiB
Python

"""Internal data structures persisted by the uvm CLI."""
from __future__ import annotations
import math
import time
import uuid
from dataclasses import dataclass, field
from ipaddress import IPv4Address
from typing import Any
from .config import MAX_STORAGE_MIB, MIN_STORAGE_MIB
from .errors import StateError, ValidationError
STATE_SCHEMA_VERSION = 3
@dataclass(frozen=True, slots=True)
class PortRule:
"""One private guest TCP port; HTTP/HTTPS are descriptive service labels."""
port: int
protocol: str = "TCP"
def __post_init__(self) -> None:
if type(self.port) is not int or not 1 <= self.port <= 65535:
raise ValidationError("port must be an integer between 1 and 65535")
if not isinstance(self.protocol, str) or self.protocol.upper() not in {"TCP", "HTTP", "HTTPS"}:
raise ValidationError("port type must be TCP, HTTP, or HTTPS")
object.__setattr__(self, "protocol", self.protocol.upper())
def to_dict(self) -> dict[str, Any]:
return {"port": self.port, "protocol": self.protocol}
@classmethod
def from_dict(cls, value: object) -> "PortRule":
if not isinstance(value, dict) or set(value) - {"port", "protocol"}:
raise StateError("invalid port rule")
try:
return cls(port=value["port"], protocol=value.get("protocol", "TCP"))
except (KeyError, ValidationError) as error:
raise StateError(f"invalid port rule: {error}") from error
@dataclass(frozen=True, slots=True)
class VmSpec:
"""Validated resource request supplied to the create command."""
cpu: float
ram_mib: int
guest_ip: IPv4Address | None = None
username: str = "root"
password: str = field(default="root", repr=False)
name: str | None = None
ports: tuple[PortRule, ...] = (PortRule(22),)
storage_mib: int = MIN_STORAGE_MIB
@dataclass(frozen=True, slots=True)
class VmResources:
"""Complete resource target retained while a stop/resize/start is pending."""
cpu: float
ram_mib: int
storage_mib: int
def __post_init__(self) -> None:
try:
valid_cpu = type(self.cpu) in (int, float) and math.isfinite(self.cpu) and self.cpu > 0
except OverflowError:
valid_cpu = False
if not valid_cpu:
raise ValidationError("CPU must be a finite value greater than 0")
if type(self.ram_mib) is not int or self.ram_mib < 128:
raise ValidationError("RAM must be an integer of at least 128 MiB")
if type(self.storage_mib) is not int or not MIN_STORAGE_MIB <= self.storage_mib <= MAX_STORAGE_MIB:
raise ValidationError(f"storage must be between {MIN_STORAGE_MIB} and {MAX_STORAGE_MIB} MiB (minimum 30 GiB)")
def to_dict(self) -> dict[str, Any]:
return {"cpu": self.cpu, "ram_mib": self.ram_mib, "storage_mib": self.storage_mib}
@classmethod
def from_dict(cls, value: object) -> "VmResources":
if not isinstance(value, dict) or set(value) != {"cpu", "ram_mib", "storage_mib"}:
raise StateError("invalid pending resize resources")
try:
return cls(**value)
except ValidationError as error:
raise StateError(f"invalid pending resize resources: {error}") from error
@dataclass(slots=True)
class VmRecord:
"""Persistent record for one local Firecracker VM."""
id: str
cpu: float
ram_mib: int
guest_ip: str
gateway: str
tap: str
mac: str
socket: str
config: str
log: str
disk: str = ""
username: str = "root"
password: str | None = field(default=None, repr=False)
status: str = "starting"
pid: int | None = None
process_start_time: str | None = None
created_at: int = field(default_factory=lambda: int(time.time()))
updated_at: int = field(default_factory=lambda: int(time.time()))
last_error: str | None = None
name: str | None = None
ports: tuple[PortRule, ...] | None = None
disk_ready: bool = False
runtime_owned: bool = False
tap_owned: bool = False
operation: str | None = None
boot_id: str | None = None
spawn_pending: bool = False
network: str | None = None
bridge: str | None = None
storage_mib: int | None = None
pending_resize: VmResources | None = None
def to_dict(self) -> dict[str, Any]:
data: dict[str, Any] = {
"id": self.id,
"cpu": self.cpu,
"ram_mib": self.ram_mib,
"storage_mib": self.storage_mib,
"pending_resize": self.pending_resize.to_dict() if self.pending_resize is not None else None,
"guest_ip": self.guest_ip,
"gateway": self.gateway,
"tap": self.tap,
"mac": self.mac,
"socket": self.socket,
"config": self.config,
"log": self.log,
"disk": self.disk,
"username": self.username,
"password": self.password,
"status": self.status,
"created_at": self.created_at,
"updated_at": self.updated_at,
"name": self.name,
"ports": [port.to_dict() for port in self.ports] if self.ports is not None else None,
"disk_ready": self.disk_ready,
"runtime_owned": self.runtime_owned,
"tap_owned": self.tap_owned,
"operation": self.operation,
"boot_id": self.boot_id,
"spawn_pending": self.spawn_pending,
"network": self.network,
"bridge": self.bridge,
}
if self.pid is not None:
data["pid"] = self.pid
if self.process_start_time is not None:
data["process_start_time"] = self.process_start_time
if self.last_error is not None:
data["last_error"] = self.last_error
return data
@classmethod
def from_dict(cls, value: object) -> "VmRecord":
if not isinstance(value, dict):
raise StateError("VM record is not an object")
required = (
"id",
"cpu",
"ram_mib",
"guest_ip",
"gateway",
"tap",
"mac",
"socket",
"config",
"log",
)
missing = [name for name in required if name not in value]
if missing:
raise StateError(f"VM record is missing fields: {', '.join(missing)}")
try:
pid_value = value.get("pid")
raw_ports = value.get("ports")
if raw_ports is not None and not isinstance(raw_ports, list):
raise StateError("VM ports must be a list or null for a legacy VM")
ports = tuple(PortRule.from_dict(port) for port in raw_ports) if raw_ports is not None else None
if ports is not None and len({port.port for port in ports}) != len(ports):
raise StateError("VM has duplicate TCP port rules")
legacy_ready = bool(value.get("disk")) and value.get("status") in {"running", "stopped"}
storage_mib = value.get("storage_mib")
if storage_mib is not None and (
type(storage_mib) is not int or not MIN_STORAGE_MIB <= storage_mib <= MAX_STORAGE_MIB
):
raise StateError("VM storage_mib must be a supported integer capacity of at least 30 GiB or null")
pending_resize = value.get("pending_resize")
for name in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"):
if name in value and type(value[name]) is not bool:
raise StateError(f"VM {name} must be a boolean")
return cls(
id=str(value["id"]),
cpu=float(value["cpu"]),
ram_mib=int(value["ram_mib"]),
guest_ip=str(value["guest_ip"]),
gateway=str(value["gateway"]),
tap=str(value["tap"]),
mac=str(value["mac"]),
socket=str(value["socket"]),
config=str(value["config"]),
log=str(value["log"]),
disk=str(value.get("disk", "")),
username=str(value.get("username", "root")),
password=(
str(value["password"])
if value.get("password") is not None
else None
),
status=str(value.get("status", "unknown")),
pid=int(pid_value) if pid_value is not None else None,
process_start_time=(
str(value["process_start_time"])
if value.get("process_start_time") is not None
else None
),
created_at=int(value.get("created_at", int(time.time()))),
updated_at=int(value.get("updated_at", value.get("created_at", int(time.time())))),
last_error=(
str(value["last_error"])
if value.get("last_error") is not None
else None
),
name=str(value["name"]) if value.get("name") is not None else None,
ports=ports,
disk_ready=value.get("disk_ready", legacy_ready),
runtime_owned=value.get("runtime_owned", True),
tap_owned=value.get("tap_owned", bool(value.get("pid"))),
operation=str(value["operation"]) if value.get("operation") is not None else None,
boot_id=str(value["boot_id"]) if value.get("boot_id") is not None else None,
spawn_pending=value.get("spawn_pending", value.get("operation") == "spawning"),
network=str(value["network"]) if value.get("network") is not None else None,
bridge=str(value["bridge"]) if value.get("bridge") is not None else None,
storage_mib=storage_mib,
pending_resize=VmResources.from_dict(pending_resize) if pending_resize is not None else None,
)
except (TypeError, ValueError) as error:
raise StateError(f"invalid VM record: {error}") from error
@dataclass(slots=True)
class State:
"""Versioned JSON state document stored under the configured base path."""
vms: dict[str, VmRecord] = field(default_factory=dict)
next_mac_index: int = 1
schema_version: int = STATE_SCHEMA_VERSION
def to_dict(self) -> dict[str, Any]:
return {
"schema_version": self.schema_version,
"next_mac_index": self.next_mac_index,
"vms": {vm_id: vm.to_dict() for vm_id, vm in self.vms.items()},
}
@classmethod
def from_dict(cls, value: object) -> "State":
if not isinstance(value, dict):
raise StateError("state file is not an object")
try:
schema_version = int(value.get("schema_version", 0))
except (TypeError, ValueError) as error:
raise StateError("state file has an invalid schema_version") from error
if schema_version > STATE_SCHEMA_VERSION:
raise StateError(
f"state schema version {schema_version} is newer than this uvm version"
)
raw_vms = value.get("vms", {})
if not isinstance(raw_vms, dict):
raise StateError("state file has an invalid vms collection")
vms = {str(vm_id): VmRecord.from_dict(vm) for vm_id, vm in raw_vms.items()}
if any(vm_id != vm.id for vm_id, vm in vms.items()):
raise StateError("VM record ID does not match its registry key")
next_mac_index = value.get("next_mac_index")
if next_mac_index is None:
next_mac_index = _next_mac_index(vms.values())
try:
next_mac_index = max(1, int(next_mac_index))
except (TypeError, ValueError) as error:
raise StateError("state file has an invalid next_mac_index") from error
return cls(
schema_version=STATE_SCHEMA_VERSION,
vms=vms,
next_mac_index=next_mac_index,
)
def new_vm_id() -> str:
"""Generate an opaque local ID that does not depend on PID or timestamp reuse."""
return f"vm-{uuid.uuid4().hex}"
def _next_mac_index(vms: object) -> int:
maximum = 0
for vm in vms:
if not isinstance(vm, VmRecord):
continue
try:
parts = vm.mac.split(":")
if parts[:3] != ["02", "fc", "00"] or len(parts) != 6:
continue
maximum = max(maximum, int("".join(parts[3:]), 16))
except ValueError:
continue
return maximum + 1