perf(delivery): use bounded native cold restore copies
This commit is contained in:
@@ -0,0 +1,220 @@
|
||||
"""One bounded native data-copy child in the existing exclusive restore unit."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import selectors
|
||||
import signal
|
||||
import stat
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from kf_oj.contracts.restore_io import RestoreIOReply
|
||||
|
||||
from . import file_copy, preflight, restore_io
|
||||
from .policy import DeliveryError, invalid_constant, unique_object
|
||||
|
||||
MAX_REQUEST_BYTES = 256 * 1024
|
||||
MAX_REPLY_BYTES = 64 * 1024
|
||||
MAX_STDERR_BYTES = 16 * 1024
|
||||
|
||||
|
||||
class Session:
|
||||
"""No fallback after partial writes; all verified witnesses stay in the child."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
payload: Path,
|
||||
backup_id: str,
|
||||
recovery_id: str,
|
||||
expected_files: int,
|
||||
expected_bytes: int,
|
||||
deadline: float,
|
||||
):
|
||||
self.deadline = deadline
|
||||
self.batch_id = self.files = self.bytes = 0
|
||||
self.expected_files, self.expected_bytes = expected_files, expected_bytes
|
||||
self.finished = False
|
||||
self.process = None
|
||||
self.stderr_bytes = 0
|
||||
helper = payload / "bin/kf-restore-io"
|
||||
preflight.protected(helper)
|
||||
info = helper.lstat()
|
||||
if (
|
||||
not stat.S_ISREG(info.st_mode)
|
||||
or info.st_nlink != 1
|
||||
or stat.S_IMODE(info.st_mode) != 0o555
|
||||
):
|
||||
raise DeliveryError("selected native recovery helper differs")
|
||||
witness = file_copy._identity(info)
|
||||
self.selector = selectors.DefaultSelector()
|
||||
try:
|
||||
restore_io.remaining(self.deadline)
|
||||
self.process = subprocess.Popen(
|
||||
[str(helper), "--stdio-v1"],
|
||||
cwd="/",
|
||||
env={
|
||||
"PATH": "/usr/bin:/bin",
|
||||
"LC_ALL": "C.UTF-8",
|
||||
"LD_LIBRARY_PATH": f"{payload}/lib:{payload}/python/lib",
|
||||
},
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
start_new_session=True,
|
||||
close_fds=True,
|
||||
)
|
||||
if file_copy._identity(helper.lstat()) != witness:
|
||||
raise DeliveryError("selected native recovery helper changed")
|
||||
assert self.process.stdin and self.process.stdout and self.process.stderr
|
||||
for handle in (self.process.stdin, self.process.stdout, self.process.stderr):
|
||||
os.set_blocking(handle.fileno(), False)
|
||||
self.selector.register(self.process.stdout, selectors.EVENT_READ, "stdout")
|
||||
self.selector.register(self.process.stderr, selectors.EVENT_READ, "stderr")
|
||||
self._exchange(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"kind": "restore.io.open",
|
||||
"backup_id": backup_id,
|
||||
"recovery_id": recovery_id,
|
||||
"expected_files": expected_files,
|
||||
"expected_logical_bytes": expected_bytes,
|
||||
"deadline_monotonic_ns": int(deadline * 1_000_000_000),
|
||||
},
|
||||
"opened",
|
||||
)
|
||||
except BaseException:
|
||||
self.close()
|
||||
raise
|
||||
|
||||
def _exchange(self, value: dict, stage: str) -> None:
|
||||
restore_io.remaining(self.deadline)
|
||||
if self.process is None or self.finished:
|
||||
raise DeliveryError("native recovery session is closed")
|
||||
raw = json.dumps(value, separators=(",", ":"), ensure_ascii=True).encode() + b"\n"
|
||||
if len(raw) > MAX_REQUEST_BYTES:
|
||||
raise DeliveryError("native recovery request exceeds bound")
|
||||
assert self.process.stdin
|
||||
self.selector.register(self.process.stdin, selectors.EVENT_WRITE, "stdin")
|
||||
offset = 0
|
||||
output = bytearray()
|
||||
try:
|
||||
while True:
|
||||
restore_io.remaining(self.deadline)
|
||||
ready = self.selector.select(max(0, self.deadline - time.monotonic()))
|
||||
if not ready:
|
||||
raise DeliveryError("native recovery copy deadline exceeded")
|
||||
for key, _ in ready:
|
||||
if key.data == "stdin":
|
||||
try:
|
||||
count = os.write(key.fd, memoryview(raw)[offset:])
|
||||
except BlockingIOError:
|
||||
continue
|
||||
if count <= 0:
|
||||
raise DeliveryError("native recovery request pipe closed")
|
||||
offset += count
|
||||
if offset == len(raw):
|
||||
self.selector.unregister(self.process.stdin)
|
||||
continue
|
||||
try:
|
||||
block = os.read(key.fd, 8192)
|
||||
except BlockingIOError:
|
||||
continue
|
||||
if key.data == "stderr":
|
||||
self.stderr_bytes += len(block)
|
||||
if self.stderr_bytes > MAX_STDERR_BYTES:
|
||||
raise DeliveryError("native recovery diagnostic exceeds bound")
|
||||
if not block:
|
||||
self.selector.unregister(key.fileobj)
|
||||
continue
|
||||
if not block:
|
||||
raise DeliveryError("native recovery stopped before verification")
|
||||
output.extend(block)
|
||||
if len(output) > MAX_REPLY_BYTES:
|
||||
raise DeliveryError("native recovery reply exceeds bound")
|
||||
if b"\n" in output:
|
||||
if offset != len(raw) or output[-1:] != b"\n" or output.count(b"\n") != 1:
|
||||
raise DeliveryError("native recovery reply framing differs")
|
||||
reply = RestoreIOReply.model_validate(
|
||||
json.loads(
|
||||
output,
|
||||
object_pairs_hook=unique_object,
|
||||
parse_constant=invalid_constant,
|
||||
)
|
||||
)
|
||||
if (
|
||||
reply.stage != stage
|
||||
or reply.error != "none"
|
||||
or (reply.batch_id, reply.verified_files, reply.verified_logical_bytes)
|
||||
!= (self.batch_id, self.files, self.bytes)
|
||||
):
|
||||
raise DeliveryError("native recovery file verification refused")
|
||||
restore_io.remaining(self.deadline)
|
||||
return
|
||||
except (OSError, ValueError, RecursionError) as error:
|
||||
raise DeliveryError("native recovery transport or reply differs") from error
|
||||
finally:
|
||||
if self.process.stdin in self.selector.get_map():
|
||||
self.selector.unregister(self.process.stdin)
|
||||
|
||||
def materialize(self, rows: list[dict]) -> None:
|
||||
for batch in restore_io.batches(rows):
|
||||
self.batch_id += 1
|
||||
self.files += len(batch)
|
||||
self.bytes += sum(item["bytes"] for item in batch)
|
||||
if self.files > self.expected_files or self.bytes > self.expected_bytes:
|
||||
raise DeliveryError("native recovery inventory exceeds original receipt")
|
||||
# The complete private manifest was already strictly validated.
|
||||
# The native decoder independently validates each row before use.
|
||||
self._exchange(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"kind": "restore.io.batch",
|
||||
"batch_id": self.batch_id,
|
||||
"files": batch,
|
||||
},
|
||||
"verified",
|
||||
)
|
||||
|
||||
def finish(self) -> None:
|
||||
if (self.files, self.bytes) != (self.expected_files, self.expected_bytes):
|
||||
raise DeliveryError("native recovery inventory incomplete")
|
||||
self._exchange({"schema_version": 1, "kind": "restore.io.finish"}, "finished")
|
||||
assert self.process is not None and self.process.stdin
|
||||
self.process.stdin.close()
|
||||
restore_io.remaining(self.deadline)
|
||||
try:
|
||||
code = self.process.wait(timeout=max(0, self.deadline - time.monotonic()))
|
||||
except subprocess.TimeoutExpired as error:
|
||||
raise DeliveryError("native recovery child did not join") from error
|
||||
if code != 0:
|
||||
raise DeliveryError("native recovery child failed after verification")
|
||||
restore_io.remaining(self.deadline)
|
||||
self.finished = True
|
||||
|
||||
def close(self) -> None:
|
||||
process = self.process
|
||||
try:
|
||||
if process is not None and process.poll() is None:
|
||||
# This child starts its own session and never spawns processes.
|
||||
# Its only descendants are its two native threads.
|
||||
try:
|
||||
if os.getpgid(process.pid) != process.pid:
|
||||
raise DeliveryError("native recovery child ownership differs")
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
process.wait(timeout=3)
|
||||
except subprocess.TimeoutExpired:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
process.wait(timeout=2)
|
||||
finally:
|
||||
self.selector.close()
|
||||
if process is not None:
|
||||
for handle in (process.stdin, process.stdout, process.stderr):
|
||||
if handle is not None:
|
||||
handle.close()
|
||||
Reference in New Issue
Block a user