Files

141 lines
4.8 KiB
Python

"""Locked, atomic persistence for the local JSON VM inventory."""
from __future__ import annotations
import fcntl
import json
import os
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from .config import Settings
from .domain import State, VmRecord
from .errors import StateError, UvmError
from .system import ensure_data_directories
class StateStore:
"""Own the JSON state file so callers cannot race IP and VM allocation."""
def __init__(self, settings: Settings) -> None:
self._settings = settings
def initialize(self) -> None:
ensure_data_directories(self._settings)
with self._locked():
if not self._settings.state_path.exists():
self._save_unlocked(State())
else:
try:
os.chmod(self._settings.state_path, 0o600)
except OSError as error:
raise StateError(
f"cannot update permissions on {self._settings.state_path}: {error}"
) from error
def load(self) -> State:
# Reads do not create /var/lib/uvm, so `uvm list` remains usable before install.
# Atomic replacement makes an unlocked reader see either the old or new full document.
return self._load_unlocked()
@contextmanager
def transaction(self) -> Iterator[State]:
"""Load and commit one state mutation while holding an exclusive lock."""
ensure_data_directories(self._settings)
with self._locked():
state = self._load_unlocked()
yield state
self._save_unlocked(state)
@contextmanager
def operation_lock(self) -> Iterator[None]:
"""Serialize external VM lifecycle work across concurrent CLI processes."""
ensure_data_directories(self._settings)
with self._locked_path(self._settings.operation_lock_path):
yield
def get(self, vm_id: str) -> VmRecord:
state = self.load()
try:
return state.vms[vm_id]
except KeyError as error:
raise UvmError(f"VM not found: {vm_id}") from error
def find_by_id_or_ip(self, identifier: str) -> VmRecord:
state = self.load()
if identifier in state.vms:
return state.vms[identifier]
for vm in state.vms.values():
if vm.guest_ip == identifier:
return vm
raise UvmError(f"VM not found: {identifier}")
@contextmanager
def _locked(self) -> Iterator[None]:
with self._locked_path(self._settings.state_lock_path):
yield
@contextmanager
def _locked_path(self, lock_path: Path) -> Iterator[None]:
lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def _load_unlocked(self) -> State:
path = self._settings.state_path
if not path.exists():
return State()
try:
if path.stat().st_mode & 0o077:
os.chmod(path, 0o600)
with path.open(encoding="utf-8") as state_file:
raw = json.load(state_file)
except PermissionError as error:
raise StateError(
f"cannot access protected VM registry {path}. Run this command with sudo."
) from error
except (OSError, json.JSONDecodeError) as error:
raise StateError(f"cannot read {path}: {error}") from error
return State.from_dict(raw)
def _save_unlocked(self, state: State) -> None:
path = self._settings.state_path
path.parent.mkdir(parents=True, exist_ok=True)
fd, temporary_path = tempfile.mkstemp(
prefix=f".{path.name}.",
suffix=".tmp",
dir=path.parent,
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as temporary_file:
json.dump(state.to_dict(), temporary_file, indent=2, sort_keys=True)
temporary_file.write("\n")
temporary_file.flush()
os.fsync(temporary_file.fileno())
os.replace(temporary_path, path)
os.chmod(path, 0o600)
self._fsync_parent(path)
except OSError as error:
raise StateError(f"cannot write {path}: {error}") from error
finally:
try:
os.unlink(temporary_path)
except FileNotFoundError:
pass
@staticmethod
def _fsync_parent(path: Path) -> None:
directory = os.open(str(path.parent), os.O_RDONLY)
try:
os.fsync(directory)
finally:
os.close(directory)