Automate host preparation and speed up game file downloads
This commit is contained in:
+138
-43
@@ -7,6 +7,7 @@ unknown and is reported for reconciliation, never automatically executed twice.
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import contextlib
|
||||
import hashlib
|
||||
import fcntl
|
||||
@@ -29,9 +30,10 @@ import urllib.parse
|
||||
import urllib.request
|
||||
import uuid
|
||||
|
||||
VERSION = "1.2.0"
|
||||
VERSION = "1.2.1"
|
||||
LABEL = "io.electric-cherry"
|
||||
MAX_FILE = 32 * 1024 * 1024
|
||||
LIBRARY_DOWNLOAD_WORKERS = 4
|
||||
ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,127}$")
|
||||
AGENT_CONFIG = Path("/etc/electric-cherry/agent.json")
|
||||
AGENT_EXECUTABLE = Path("/usr/local/libexec/electric-cherry-agent.py")
|
||||
@@ -788,13 +790,46 @@ class Agent:
|
||||
pass
|
||||
libraries, caches = [], []
|
||||
for marker in (self.root / "libraries").glob("*/*/manifest.json"):
|
||||
# A crash between writing a manifest and its atomic promotion must
|
||||
# never expose the staging directory as a ready library.
|
||||
if marker.parent.name.startswith("."):
|
||||
continue
|
||||
try:
|
||||
value = json.loads(marker.read_text())
|
||||
libraries.append({"id": value["libraryId"], "state": "ready"})
|
||||
caches.append({"libraryVersionId": value["libraryId"], "status": "ready", "bytesDownloaded": value["size"], "bytesTotal": value["size"], "manifestSha256": value.get("manifestSha256")})
|
||||
except (OSError, ValueError, KeyError):
|
||||
continue
|
||||
progress = list(self.library_progress.values())
|
||||
ready = {item["id"] for item in libraries}
|
||||
progress = {key: value.copy() for key, value in list(self.library_progress.items()) if key not in ready}
|
||||
for marker in (self.root / "libraries").glob("*/.*.partial/source.json"):
|
||||
try:
|
||||
value = json.loads(marker.read_text())
|
||||
library_id = identifier(value["libraryId"])
|
||||
if library_id in ready or library_id in progress:
|
||||
continue
|
||||
expected = self.library_dir(value["game"], library_id).with_name("." + library_id + ".partial")
|
||||
if marker.parent != expected or marker.is_symlink() or marker.parent.is_symlink():
|
||||
continue
|
||||
files = value["files"]
|
||||
canonical = json.dumps([{ "path": item["path"], "size": item["size"], "sha256": item["sha256"] } for item in files], separators=(",", ":"), ensure_ascii=False)
|
||||
if digest(canonical.encode("utf-8")) != value["manifestSha256"]:
|
||||
continue
|
||||
retained, size, paths = 0, 0, set()
|
||||
for item in files:
|
||||
if item["path"] in paths or type(item["size"]) is not int or item["size"] < 0:
|
||||
raise AgentError("Invalid partial library manifest")
|
||||
paths.add(item["path"])
|
||||
target = scoped(marker.parent / "content", item["path"])
|
||||
size += item["size"]
|
||||
if target.exists() and stat_module.S_ISREG(target.stat().st_mode):
|
||||
retained += min(target.stat().st_size, item["size"])
|
||||
# These are resumable bytes, not verified content. Reporting
|
||||
# them survives restarts without requiring space for two copies.
|
||||
progress[library_id] = {"libraryVersionId": library_id, "status": "missing", "bytesDownloaded": retained, "bytesTotal": size, "manifestSha256": value["manifestSha256"]}
|
||||
except (AgentError, OSError, ValueError, KeyError, TypeError):
|
||||
continue
|
||||
progress = list(progress.values())
|
||||
libraries.extend({"id": item["libraryVersionId"], "state": item["status"]} for item in progress)
|
||||
caches.extend(progress)
|
||||
images = []
|
||||
@@ -1195,6 +1230,65 @@ class Agent:
|
||||
raise AgentError("Artifact redirects are not allowed")
|
||||
return urllib.request.build_opener(NoRedirect()).open(req, timeout=60)
|
||||
|
||||
def download_library_file(self, item, target, files_url, update_progress, aborted):
|
||||
def check_active():
|
||||
if aborted.is_set():
|
||||
raise AgentError("Library download stopped after a transfer failed")
|
||||
|
||||
check_active()
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
offset = target.stat().st_size if target.exists() else 0
|
||||
if offset > item["size"]:
|
||||
target.unlink()
|
||||
offset = 0
|
||||
checksum = hashlib.sha256()
|
||||
# Hash the saved prefix once so resumed transfers still verify the whole
|
||||
# file. New bytes are hashed as they arrive, avoiding a second disk read.
|
||||
if offset:
|
||||
with target.open("rb") as stream:
|
||||
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
check_active()
|
||||
checksum.update(chunk)
|
||||
if offset == item["size"] and target.exists():
|
||||
if checksum.hexdigest() == item["sha256"]:
|
||||
update_progress(item["path"], offset)
|
||||
return
|
||||
target.unlink()
|
||||
offset = 0
|
||||
checksum = hashlib.sha256()
|
||||
update_progress(item["path"], offset)
|
||||
check_active()
|
||||
url = files_url.rstrip("/") + "/" + urllib.parse.quote(item["path"], safe="/")
|
||||
with self.artifact_open("GET", url, extra_headers={"Range": f"bytes={offset}-"} if offset else {}) as response:
|
||||
if response.status == 206:
|
||||
expected = f"bytes {offset}-{item['size'] - 1}/{item['size']}"
|
||||
if response.headers.get("Content-Range") != expected:
|
||||
raise AgentError("Invalid resumed download response")
|
||||
elif response.status != 200:
|
||||
raise AgentError("Invalid library download response")
|
||||
elif offset:
|
||||
# A server may ignore Range and send the whole file instead.
|
||||
offset = 0
|
||||
checksum = hashlib.sha256()
|
||||
update_progress(item["path"], 0)
|
||||
with target.open("ab" if offset else "wb") as output:
|
||||
received = offset
|
||||
while True:
|
||||
check_active()
|
||||
chunk = response.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
received += len(chunk)
|
||||
if received > item["size"]:
|
||||
raise AgentError("Library download exceeded its manifest size")
|
||||
output.write(chunk)
|
||||
checksum.update(chunk)
|
||||
update_progress(item["path"], received)
|
||||
output.flush()
|
||||
os.fsync(output.fileno())
|
||||
if target.stat().st_size != item["size"] or checksum.hexdigest() != item["sha256"]:
|
||||
raise AgentError("Downloaded library file failed integrity verification")
|
||||
|
||||
def library_download(self, payload):
|
||||
library_id = identifier(payload.get("libraryVersionId", payload.get("libraryId")))
|
||||
with self.artifact_open("GET", payload["manifestUrl"]) as response:
|
||||
@@ -1221,55 +1315,46 @@ class Agent:
|
||||
source_marker = staging / "source.json"
|
||||
if source_marker.exists() and json.loads(source_marker.read_text()).get("manifestSha256") != manifest_sha:
|
||||
raise Rejected("Partial library download belongs to a different manifest")
|
||||
atomic_json(source_marker, {"libraryId": library_id, "manifestSha256": manifest_sha})
|
||||
size = sum(item["size"] for item in files)
|
||||
downloaded = 0
|
||||
paths = set()
|
||||
atomic_json(source_marker, {"libraryId": library_id, "game": manifest["game"], "size": size, "files": files, "manifestSha256": manifest_sha})
|
||||
targets = {}
|
||||
progress = {"libraryVersionId": library_id, "status": "downloading", "bytesDownloaded": 0, "bytesTotal": size, "manifestSha256": manifest_sha}
|
||||
self.library_progress[library_id] = progress
|
||||
try:
|
||||
if shutil.disk_usage(staging).free < max(0, size - sum(path.stat().st_size for path in staging.rglob("*") if path.is_file())):
|
||||
raise AgentError("Insufficient library storage")
|
||||
# Validate every path before starting concurrent writes. In addition
|
||||
# to duplicates, a manifest cannot treat a file as a parent folder.
|
||||
for item in files:
|
||||
if item["path"] in paths or not isinstance(item["size"], int) or item["size"] < 0 or not re.fullmatch("[0-9a-f]{64}", item["sha256"]):
|
||||
if item["path"] in targets or not isinstance(item["size"], int) or item["size"] < 0 or not re.fullmatch("[0-9a-f]{64}", item["sha256"]):
|
||||
raise AgentError("Invalid library file manifest")
|
||||
paths.add(item["path"])
|
||||
target = scoped(staging / "content", item["path"])
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
offset = target.stat().st_size if target.exists() else 0
|
||||
if offset > item["size"]:
|
||||
target.unlink()
|
||||
offset = 0
|
||||
if offset == item["size"] and target.exists() and file_digest(target) != item["sha256"]:
|
||||
target.unlink()
|
||||
offset = 0
|
||||
if offset < item["size"] or not target.exists():
|
||||
url = payload["filesBaseUrl"].rstrip("/") + "/" + urllib.parse.quote(item["path"], safe="/")
|
||||
with self.artifact_open("GET", url, extra_headers={"Range": f"bytes={offset}-"} if offset else {}) as response:
|
||||
if offset and response.status == 206:
|
||||
content_range = response.headers.get("Content-Range", "")
|
||||
if not content_range.startswith(f"bytes {offset}-"):
|
||||
raise AgentError("Invalid resumed download response")
|
||||
elif offset:
|
||||
offset = 0
|
||||
with target.open("ab" if offset else "wb") as output:
|
||||
received = offset
|
||||
while True:
|
||||
chunk = response.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
received += len(chunk)
|
||||
if received > item["size"]:
|
||||
raise AgentError("Library download exceeded its manifest size")
|
||||
output.write(chunk)
|
||||
progress["bytesDownloaded"] = downloaded + received
|
||||
output.flush()
|
||||
os.fsync(output.fileno())
|
||||
progress["status"] = "verifying"
|
||||
if target.stat().st_size != item["size"] or file_digest(target) != item["sha256"]:
|
||||
raise AgentError("Downloaded library file failed integrity verification")
|
||||
downloaded += item["size"]
|
||||
progress.update(status="downloading", bytesDownloaded=downloaded)
|
||||
targets[item["path"]] = scoped(staging / "content", item["path"])
|
||||
for path in targets:
|
||||
if any(str(parent) in targets for parent in PurePosixPath(path).parents):
|
||||
raise AgentError("Library file paths overlap")
|
||||
progress_lock = threading.Lock()
|
||||
offsets = {}
|
||||
aborted = threading.Event()
|
||||
|
||||
def update_progress(path, offset):
|
||||
with progress_lock:
|
||||
progress["bytesDownloaded"] += offset - offsets.get(path, 0)
|
||||
offsets[path] = offset
|
||||
|
||||
def transfer(item):
|
||||
try:
|
||||
self.download_library_file(item, targets[item["path"]], payload["filesBaseUrl"], update_progress, aborted)
|
||||
except Exception:
|
||||
aborted.set()
|
||||
raise
|
||||
|
||||
# Bound sockets, memory and disk pressure while allowing other files
|
||||
# to progress during network round trips and per-file fsyncs.
|
||||
with ThreadPoolExecutor(max_workers=min(LIBRARY_DOWNLOAD_WORKERS, len(files))) as pool:
|
||||
futures = [pool.submit(transfer, item) for item in files]
|
||||
for future in as_completed(futures):
|
||||
future.result()
|
||||
progress["status"] = "verifying"
|
||||
saved = {"libraryId": library_id, "game": manifest["game"], "size": size, "files": files, "manifestSha256": manifest_sha}
|
||||
atomic_json(staging / "manifest.json", saved)
|
||||
for path in (staging / "content").rglob("*"):
|
||||
@@ -1446,6 +1531,16 @@ class Agent:
|
||||
if not self.config.get("pendingCredential") or self.recover_credential():
|
||||
return self.journal.update(operation["id"], "completed", {"rotated": True, "reconciled": True})
|
||||
return row
|
||||
if kind == "runtime.pull":
|
||||
# Docker resumes already downloaded layers. Pull again under
|
||||
# the same operation ID: an existing tag can still refer to an
|
||||
# older image and cannot prove an interrupted pull completed.
|
||||
self.journal.update(operation["id"], "running")
|
||||
try:
|
||||
result = self.dispatch(operation)
|
||||
return self.journal.update(operation["id"], "completed", {**result, "reconciled": True})
|
||||
except Exception:
|
||||
return self.journal.update(operation["id"], "failed", error="Interrupted runtime download could not be resumed")
|
||||
if kind == "library.import":
|
||||
try:
|
||||
with self.artifact_open("GET", payload["completeUrl"]) as response:
|
||||
|
||||
Reference in New Issue
Block a user