from __future__ import annotations import argparse import json import os import threading from collections.abc import Callable from pathlib import Path from typing import Any class PerformanceModeError(RuntimeError): """Raised when cpufreq policies cannot be changed safely.""" class PerformanceModeManager: RECORD_NAME = "cpufreq-governors.original.json" def __init__( self, runtime_root: Path, *, sysfs_root: Path | None = None, ) -> None: configured_root = os.environ.get("MATRIX_CPUFREQ_SYSFS_ROOT") self.sysfs_root = Path( sysfs_root if sysfs_root is not None else configured_root or "/sys/devices/system/cpu/cpufreq" ) self.runtime_root = Path(runtime_root) self.record_path = self.runtime_root / self.RECORD_NAME self._lock = threading.RLock() self._restore_governors: dict[str, str] = {} self._requested = False self._last_error: str | None = None self._load_or_capture_baseline() def _policy_paths(self) -> dict[str, Path]: if not self.sysfs_root.is_dir(): return {} return { path.name: path for path in sorted(self.sysfs_root.glob("policy*"), key=lambda item: item.name) if path.is_dir() } @staticmethod def _read_text(path: Path) -> str: return path.read_text(encoding="ascii").strip() @staticmethod def _write_text(path: Path, value: str) -> None: path.write_text(f"{value}\n", encoding="ascii") def _read_governors(self) -> dict[str, str]: return { name: self._read_text(path / "scaling_governor") for name, path in self._policy_paths().items() } def _available_governors(self) -> dict[str, set[str]]: result: dict[str, set[str]] = {} for name, path in self._policy_paths().items(): available_path = path / "scaling_available_governors" try: result[name] = set(self._read_text(available_path).split()) except OSError: result[name] = set() return result def _write_record(self, governors: dict[str, str]) -> None: self.runtime_root.mkdir(parents=True, exist_ok=True) temporary = self.record_path.with_suffix(".tmp") payload = { "schema_version": 1, "governors": governors, } with temporary.open("w", encoding="utf-8", newline="\n") as handle: json.dump(payload, handle, ensure_ascii=False, sort_keys=True, separators=(",", ":")) handle.write("\n") handle.flush() os.fsync(handle.fileno()) os.replace(temporary, self.record_path) def _read_record(self) -> dict[str, str] | None: if not self.record_path.is_file(): return None try: payload = json.loads(self.record_path.read_text(encoding="utf-8")) governors = payload.get("governors") if payload.get("schema_version") != 1 or not isinstance(governors, dict): raise ValueError("invalid governor record") checked = { str(name): str(value) for name, value in governors.items() if str(name).startswith("policy") and str(value).strip() } if checked != governors: raise ValueError("invalid governor record fields") return checked except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: self._last_error = f"saved governor record is invalid: {exc}" return None def _load_or_capture_baseline(self) -> None: with self._lock: policies = self._policy_paths() if not policies: return saved = self._read_record() if saved is not None: if set(saved) == set(policies): self._restore_governors = saved return self._last_error = "saved governor policies do not match current cpufreq policies" return try: current = self._read_governors() self._write_record(current) self._restore_governors = current except OSError as exc: self._last_error = f"could not capture original governors: {exc}" def _set_governors(self, requested: dict[str, str]) -> None: policies = self._policy_paths() if set(requested) != set(policies): raise PerformanceModeError("cpufreq policies changed while applying governor") before = self._read_governors() changed: list[str] = [] try: for name, value in requested.items(): self._write_text(policies[name] / "scaling_governor", value) changed.append(name) actual = self._read_text(policies[name] / "scaling_governor") if actual != value: raise PerformanceModeError( f"{name} governor is {actual!r} after requesting {value!r}" ) except Exception as exc: rollback_errors: list[str] = [] for name in reversed(changed): try: self._write_text(policies[name] / "scaling_governor", before[name]) except OSError as rollback_exc: rollback_errors.append(f"{name}: {rollback_exc}") suffix = f"; rollback failed for {', '.join(rollback_errors)}" if rollback_errors else "" raise PerformanceModeError(f"could not apply cpufreq governor: {exc}{suffix}") from exc def available(self) -> bool: policies = self._policy_paths() if not policies or set(self._restore_governors) != set(policies): return False available = self._available_governors() return all("performance" in values for values in available.values()) def start(self, requested: bool) -> None: with self._lock: self._requested = requested if not requested: if self._restore_governors and self._read_governors() != self._restore_governors: try: self._set_governors(self._restore_governors) self._last_error = None except PerformanceModeError as exc: self._last_error = str(exc) return try: self._apply_requested(True) self._last_error = None except PerformanceModeError as exc: self._last_error = str(exc) def _apply_requested(self, requested: bool) -> None: if requested: if not self.available(): raise PerformanceModeError("performance governor is unavailable") target = {name: "performance" for name in self._policy_paths()} else: if not self._restore_governors: raise PerformanceModeError("original governors are unavailable") target = dict(self._restore_governors) self._set_governors(target) def transact(self, requested: bool, persist: Callable[[bool], Any]) -> None: with self._lock: previous_requested = self._requested previous_governors = self._read_governors() try: self._apply_requested(requested) persist(requested) except Exception as exc: try: if previous_governors: self._set_governors(previous_governors) except Exception as rollback_exc: self._last_error = f"{exc}; governor rollback failed: {rollback_exc}" raise PerformanceModeError(self._last_error) from exc self._requested = previous_requested self._last_error = str(exc) if isinstance(exc, PerformanceModeError): raise raise PerformanceModeError(f"performance mode transaction failed: {exc}") from exc self._requested = requested self._last_error = None def restore(self, *, remove_record: bool = False) -> None: with self._lock: if self._restore_governors: try: self._set_governors(self._restore_governors) self._last_error = None except PerformanceModeError as exc: self._last_error = str(exc) raise if remove_record and self.record_path.exists(): self.record_path.unlink() def status(self, requested: bool | None = None) -> dict[str, Any]: with self._lock: current: dict[str, str] = {} try: current = self._read_governors() except OSError as exc: self._last_error = f"could not read current governors: {exc}" effective_requested = self._requested if requested is None else requested effective = bool( effective_requested and current and all(value == "performance" for value in current.values()) ) return { "requested": bool(effective_requested), "available": self.available(), "effective": effective, "current_governors": current, "restore_governors": dict(self._restore_governors), "last_error": self._last_error, } def _restore_cli() -> int: parser = argparse.ArgumentParser(description="Restore matrix service cpufreq governors") parser.add_argument("restore", nargs="?") parser.add_argument("--runtime-root", type=Path, required=True) parser.add_argument("--sysfs-root", type=Path) args = parser.parse_args() manager = PerformanceModeManager(args.runtime_root, sysfs_root=args.sysfs_root) try: manager.restore(remove_record=True) except PerformanceModeError as exc: parser.error(str(exc)) return 0 if __name__ == "__main__": raise SystemExit(_restore_cli())