aios/active-runtime/app/aios_v2r1/cell_manager.py

"""AIOS V2R1 Tenant Cell lifecycle and atomic registry primitives."""
from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
import fcntl
import json
import os
from pathlib import Path
import tempfile
import threading
from typing import Iterator

from .tenant_context import ValidationError, require_ident

_THREAD_LOCKS_GUARD = threading.Lock()
_THREAD_LOCKS: dict[str, threading.RLock] = {}
_STATES = {"created", "starting", "healthy", "upgrading", "rollback", "failed"}
_TRANSITIONS = {
    "created": {"starting", "failed"},
    "starting": {"healthy", "failed"},
    "healthy": {"upgrading", "failed"},
    "upgrading": {"healthy", "rollback", "failed"},
    "rollback": {"healthy", "failed"},
    "failed": {"starting"},
}


def _lock_for(path: Path) -> threading.RLock:
    key = str(path.resolve())
    with _THREAD_LOCKS_GUARD:
        return _THREAD_LOCKS.setdefault(key, threading.RLock())


@dataclass(frozen=True)
class CellRecord:
    tenant_id: str
    cell_id: str
    version: str
    state: str
    config_digest: str
    created_at: str


class CellManager:
    def __init__(self, registry: str | Path) -> None:
        self.registry = Path(registry)
        self.registry.parent.mkdir(parents=True, exist_ok=True)
        self._lock_file = self.registry.with_name(self.registry.name + ".lock")
        self._thread_lock = _lock_for(self.registry)
        with self._exclusive():
            if not self.registry.exists():
                self._save({"cells": {}})
            else:
                self._load()

    @contextmanager
    def _exclusive(self) -> Iterator[None]:
        with self._thread_lock:
            with self._lock_file.open("a+", encoding="utf-8") as lock:
                os.chmod(self._lock_file, 0o600)
                fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
                try:
                    yield
                finally:
                    fcntl.flock(lock.fileno(), fcntl.LOCK_UN)

    def _load(self) -> dict[str, dict[str, dict[str, str]]]:
        failure: ValidationError | None = None
        try:
            value = json.loads(self.registry.read_text(encoding="utf-8"))
        except (OSError, UnicodeError, json.JSONDecodeError):
            failure = ValidationError("cell registry unavailable")
        if failure is not None:
            raise failure
        if not isinstance(value, dict) or set(value) != {"cells"}:
            raise ValidationError("cell registry unavailable")
        cells = value["cells"]
        if not isinstance(cells, dict):
            raise ValidationError("cell registry unavailable")
        expected = {
            "tenant_id", "cell_id", "version", "state",
            "config_digest", "created_at",
        }
        for tenant_id, item in cells.items():
            if (
                not isinstance(item, dict) or set(item) != expected
                or item.get("tenant_id") != tenant_id
                or item.get("state") not in _STATES
                or any(not isinstance(item.get(name), str) for name in expected)
            ):
                raise ValidationError("cell registry unavailable")
        return value

    def _save(self, value: dict[str, object]) -> None:
        descriptor, temporary_name = tempfile.mkstemp(
            prefix=f".{self.registry.name}.", dir=self.registry.parent, text=True
        )
        temporary = Path(temporary_name)
        try:
            with os.fdopen(descriptor, "w", encoding="utf-8") as output:
                json.dump(value, output, ensure_ascii=False, indent=2)
                output.write("\n")
                output.flush()
                os.fsync(output.fileno())
            os.chmod(temporary, 0o600)
            os.replace(temporary, self.registry)
            directory = os.open(self.registry.parent, os.O_RDONLY | os.O_DIRECTORY)
            try:
                os.fsync(directory)
            finally:
                os.close(directory)
        finally:
            if temporary.exists():
                temporary.unlink()

    def create(
        self, tenant_id: str, cell_id: str, version: str, digest: str,
    ) -> CellRecord:
        tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
        cell_id = require_ident("cell_id", cell_id)
        version = require_ident("version", version)
        digest = require_ident("digest", digest)
        with self._exclusive():
            data = self._load()
            if tenant_id in data["cells"]:
                raise ValidationError("tenant already has a cell")
            record = CellRecord(
                tenant_id, cell_id, version, "created", digest,
                datetime.now(timezone.utc).isoformat(),
            )
            data["cells"][tenant_id] = record.__dict__
            self._save(data)
            return record

    def status(self, tenant_id: str) -> CellRecord:
        tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
        with self._exclusive():
            data = self._load()
            item = data["cells"].get(tenant_id)
            if not item:
                raise ValidationError("unknown tenant")
            return CellRecord(**item)

    def transition(self, tenant_id: str, state: str) -> CellRecord:
        tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
        if state not in _STATES:
            raise ValidationError("invalid cell state")
        with self._exclusive():
            data = self._load()
            item = data["cells"].get(tenant_id)
            if not item:
                raise ValidationError("unknown tenant")
            current = item["state"]
            if state == current:
                return CellRecord(**item)
            if state not in _TRANSITIONS[current]:
                raise ValidationError(f"invalid cell transition: {current}->{state}")
            item["state"] = state
            self._save(data)
            return CellRecord(**item)