139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Atomically persist the Python profile operation journal."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
|
|
class JournalError(ValueError):
|
|
pass
|
|
|
|
|
|
TRANSITIONS = {
|
|
None: {"profile_selected"},
|
|
"profile_selected": {"checkpoint_ready", "rolled_back"},
|
|
"checkpoint_ready": {"installing", "rolled_back"},
|
|
"installing": {"installed", "rollback_required"},
|
|
"rollback_required": {"rolled_back"},
|
|
"installed": set(),
|
|
"rolled_back": set(),
|
|
}
|
|
|
|
|
|
def canonical(record: dict) -> bytes:
|
|
return (json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n").encode()
|
|
|
|
|
|
def digest_record(record: dict) -> str:
|
|
unsigned = {key: value for key, value in record.items() if key != "record_sha256"}
|
|
return hashlib.sha256(canonical(unsigned)).hexdigest()
|
|
|
|
|
|
def read_verified(path: Path) -> dict | None:
|
|
if not path.exists():
|
|
return None
|
|
record = json.loads(path.read_text(encoding="utf-8"))
|
|
if record.get("record_sha256") != digest_record(record):
|
|
raise JournalError("JOURNAL_DIGEST_INVALID")
|
|
return record
|
|
|
|
|
|
def next_record(current: dict | None, expected_generation: int, operation_id: str,
|
|
profile: str, manifest_sha256: str, state: str,
|
|
checkpoint_id: str | None) -> dict:
|
|
generation = current.get("generation", 0) if current else 0
|
|
if generation != expected_generation:
|
|
raise JournalError("JOURNAL_GENERATION_STALE")
|
|
prior_state = current.get("state") if current else None
|
|
if state not in TRANSITIONS.get(prior_state, set()):
|
|
raise JournalError("JOURNAL_TRANSITION_INVALID")
|
|
if current and (current["operation_id"] != operation_id or current["selected_profile"] != profile
|
|
or current["manifest_sha256"] != manifest_sha256):
|
|
raise JournalError("JOURNAL_IDENTITY_CHANGED")
|
|
inherited_checkpoint = current.get("checkpoint_id") if current else None
|
|
checkpoint = checkpoint_id or inherited_checkpoint
|
|
if state in {"checkpoint_ready", "installing", "installed", "rollback_required", "rolled_back"} \
|
|
and prior_state != "profile_selected" and not checkpoint:
|
|
raise JournalError("CHECKPOINT_ID_REQUIRED")
|
|
if state in {"checkpoint_ready", "installing", "installed", "rollback_required"} and not checkpoint:
|
|
raise JournalError("CHECKPOINT_ID_REQUIRED")
|
|
record = {
|
|
"schema": "com.cezen.nexus-one.python-profile-journal.v1",
|
|
"generation": generation + 1,
|
|
"operation_id": operation_id,
|
|
"selected_profile": profile,
|
|
"manifest_sha256": manifest_sha256,
|
|
"state": state,
|
|
"checkpoint_id": checkpoint,
|
|
"previous_record_sha256": current.get("record_sha256") if current else None,
|
|
}
|
|
record["record_sha256"] = digest_record(record)
|
|
return record
|
|
|
|
|
|
def atomic_write(path: Path, record: dict) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
|
try:
|
|
os.fchmod(fd, 0o600)
|
|
with os.fdopen(fd, "wb", closefd=True) as stream:
|
|
stream.write(canonical(record))
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(temporary, path)
|
|
directory_fd = os.open(path.parent, os.O_RDONLY)
|
|
try:
|
|
os.fsync(directory_fd)
|
|
finally:
|
|
os.close(directory_fd)
|
|
finally:
|
|
if os.path.exists(temporary):
|
|
os.unlink(temporary)
|
|
|
|
|
|
def update(path: Path, **values) -> dict:
|
|
lock_path = path.with_suffix(path.suffix + ".lock")
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with lock_path.open("a+b") as lock:
|
|
os.chmod(lock_path, 0o600)
|
|
fcntl.flock(lock, fcntl.LOCK_EX)
|
|
current = read_verified(path)
|
|
record = next_record(current, **values)
|
|
atomic_write(path, record)
|
|
return record
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("path", type=Path)
|
|
parser.add_argument("--expected-generation", type=int, required=True)
|
|
parser.add_argument("--operation-id", required=True)
|
|
parser.add_argument("--profile", choices=("cpu-only", "nvidia-cu128"), required=True)
|
|
parser.add_argument("--manifest-sha256", required=True)
|
|
parser.add_argument("--state", choices=tuple(state for state in TRANSITIONS if state), required=True)
|
|
parser.add_argument("--checkpoint-id")
|
|
args = parser.parse_args()
|
|
try:
|
|
record = update(
|
|
args.path, expected_generation=args.expected_generation,
|
|
operation_id=args.operation_id, profile=args.profile,
|
|
manifest_sha256=args.manifest_sha256, state=args.state,
|
|
checkpoint_id=args.checkpoint_id,
|
|
)
|
|
except (JournalError, json.JSONDecodeError) as error:
|
|
print(json.dumps({"written": False, "code": str(error)}, sort_keys=True))
|
|
return 1
|
|
print(json.dumps({"written": True, **record}, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|