feat: execute durable archive transfers
This commit is contained in:
@@ -0,0 +1,557 @@
|
||||
"""Control-commanded happy-path transfer execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Callable
|
||||
|
||||
from archive_clients.protocol import decode_message, encode_message
|
||||
from archive_clients.qbittorrent import (
|
||||
QBittorrentDownloadAttempt,
|
||||
QBittorrentReader,
|
||||
)
|
||||
from archive_clients.resources import NormalizedResource
|
||||
from archive_clients.state import ClientStore
|
||||
from archive_clients.syncthing import SyncthingTransferObserver
|
||||
from archive_clients.transfer import (
|
||||
TransferError,
|
||||
TransferIntegrityError,
|
||||
cleanup_transfer,
|
||||
load_published_transfer,
|
||||
materialize_transfer,
|
||||
stage_transfer,
|
||||
)
|
||||
from archive_control.v1 import (
|
||||
common_pb2,
|
||||
control_pb2,
|
||||
job_pb2,
|
||||
resource_pb2,
|
||||
transfer_pb2,
|
||||
)
|
||||
|
||||
|
||||
class JobExecutionError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class ClientJobExecutor:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client_id: str,
|
||||
qbittorrent: QBittorrentReader,
|
||||
store: ClientStore,
|
||||
qb_root: Path,
|
||||
qb_api_root: PurePosixPath,
|
||||
route_path: Callable[[str], Path],
|
||||
syncthing_transport: object,
|
||||
sparse_supported: bool,
|
||||
poll_interval: float = 1,
|
||||
):
|
||||
self.client_id = client_id
|
||||
self.qbittorrent = qbittorrent
|
||||
self.store = store
|
||||
self.qb_root = qb_root
|
||||
self.qb_api_root = qb_api_root
|
||||
self.route_path = route_path
|
||||
self.syncthing_transport = syncthing_transport
|
||||
self.sparse_supported = sparse_supported
|
||||
self.poll_interval = poll_interval
|
||||
|
||||
def assign(
|
||||
self, command: control_pb2.AssignJobCommand
|
||||
) -> list[control_pb2.JobEvent]:
|
||||
definition = command.job
|
||||
self._validate_definition(definition)
|
||||
replay = self._replay(
|
||||
definition.job_id, command.expected_last_event_sequence
|
||||
)
|
||||
if replay:
|
||||
return replay
|
||||
event = self._event(
|
||||
definition,
|
||||
sequence=command.expected_last_event_sequence + 1,
|
||||
revision=command.expected_job_revision,
|
||||
event_type=control_pb2.JOB_EVENT_TYPE_ASSIGNED,
|
||||
state=job_pb2.JOB_STATE_PREPARING,
|
||||
committed=False,
|
||||
)
|
||||
self._record(definition, event)
|
||||
return [event]
|
||||
|
||||
def execute(
|
||||
self, command: control_pb2.ExecuteStepCommand
|
||||
) -> list[control_pb2.JobEvent]:
|
||||
definition = self._definition(command.job_id)
|
||||
replay = self._replay(
|
||||
command.job_id, command.expected_last_event_sequence
|
||||
)
|
||||
if replay and replay[-1].type in {
|
||||
control_pb2.JOB_EVENT_TYPE_STEP_SUCCEEDED,
|
||||
control_pb2.JOB_EVENT_TYPE_COMMITTED,
|
||||
control_pb2.JOB_EVENT_TYPE_SUCCEEDED,
|
||||
control_pb2.JOB_EVENT_TYPE_FAILED,
|
||||
control_pb2.JOB_EVENT_TYPE_CLEANUP_REQUIRED,
|
||||
control_pb2.JOB_EVENT_TYPE_CANCELLED,
|
||||
}:
|
||||
return replay
|
||||
started = replay[0] if replay else self._event(
|
||||
definition,
|
||||
sequence=command.expected_last_event_sequence + 1,
|
||||
revision=command.expected_job_revision + 1,
|
||||
event_type=control_pb2.JOB_EVENT_TYPE_STEP_STARTED,
|
||||
state=job_pb2.JOB_STATE_RUNNING,
|
||||
committed=(
|
||||
self._committed(command.job_id)
|
||||
or command.step == job_pb2.JOB_STEP_KIND_STAGING_CLEANUP
|
||||
),
|
||||
step=command.step,
|
||||
step_state=job_pb2.STEP_STATE_RUNNING,
|
||||
)
|
||||
if not replay:
|
||||
self._record(definition, started)
|
||||
try:
|
||||
result = self._execute_step(definition, command.step)
|
||||
except Exception as error:
|
||||
cleanup = (
|
||||
command.step == job_pb2.JOB_STEP_KIND_STAGING_CLEANUP
|
||||
and started.committed
|
||||
)
|
||||
failed = self._event(
|
||||
definition,
|
||||
sequence=started.sequence + 1,
|
||||
revision=started.job_revision + 1,
|
||||
event_type=(
|
||||
control_pb2.JOB_EVENT_TYPE_CLEANUP_REQUIRED
|
||||
if cleanup else control_pb2.JOB_EVENT_TYPE_FAILED
|
||||
),
|
||||
state=(
|
||||
job_pb2.JOB_STATE_CLEANUP_REQUIRED
|
||||
if cleanup else job_pb2.JOB_STATE_FAILED
|
||||
),
|
||||
committed=started.committed,
|
||||
step=command.step,
|
||||
step_state=job_pb2.STEP_STATE_FAILED,
|
||||
)
|
||||
failed.error.code = _job_error_code(error)
|
||||
failed.error.message = str(error) or type(error).__name__
|
||||
failed.error.retryable = cleanup
|
||||
self._record(definition, failed)
|
||||
return [started, failed]
|
||||
sequence = started.sequence + 1
|
||||
revision = started.job_revision + 1
|
||||
committed = started.committed
|
||||
event_type = control_pb2.JOB_EVENT_TYPE_STEP_SUCCEEDED
|
||||
state = job_pb2.JOB_STATE_RUNNING
|
||||
if command.step == job_pb2.JOB_STEP_KIND_QB_VERIFY:
|
||||
committed = True
|
||||
event_type = control_pb2.JOB_EVENT_TYPE_COMMITTED
|
||||
elif command.step == job_pb2.JOB_STEP_KIND_STAGING_CLEANUP:
|
||||
event_type = control_pb2.JOB_EVENT_TYPE_SUCCEEDED
|
||||
state = job_pb2.JOB_STATE_SUCCEEDED
|
||||
succeeded = self._event(
|
||||
definition,
|
||||
sequence=sequence,
|
||||
revision=revision,
|
||||
event_type=event_type,
|
||||
state=state,
|
||||
committed=committed,
|
||||
step=command.step,
|
||||
step_state=job_pb2.STEP_STATE_SUCCEEDED,
|
||||
)
|
||||
if result is not None:
|
||||
succeeded.observed_placement.CopyFrom(result)
|
||||
self._record(definition, succeeded)
|
||||
return [started, succeeded]
|
||||
|
||||
def _execute_step(
|
||||
self, definition: job_pb2.JobDefinition, step: int
|
||||
):
|
||||
if step == job_pb2.JOB_STEP_KIND_SOURCE_STAGE:
|
||||
self._source_stage(definition)
|
||||
return None
|
||||
if step == job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER:
|
||||
self._wait_for_syncthing(definition)
|
||||
return None
|
||||
if step == job_pb2.JOB_STEP_KIND_TARGET_MATERIALIZE:
|
||||
self._target_materialize(definition)
|
||||
return None
|
||||
if step == job_pb2.JOB_STEP_KIND_QB_VERIFY:
|
||||
return self._qb_verify(definition)
|
||||
if step == job_pb2.JOB_STEP_KIND_STAGING_CLEANUP:
|
||||
cleanup_transfer(self._job_directory(definition))
|
||||
temporary_metainfo = self._metainfo_path(definition.job_id)
|
||||
if temporary_metainfo.exists():
|
||||
temporary_metainfo.unlink()
|
||||
for directory in (
|
||||
temporary_metainfo.parent,
|
||||
temporary_metainfo.parent.parent,
|
||||
):
|
||||
try:
|
||||
directory.rmdir()
|
||||
except OSError:
|
||||
break
|
||||
return None
|
||||
raise JobExecutionError("job step is unsupported")
|
||||
|
||||
def _source_stage(self, definition: job_pb2.JobDefinition) -> None:
|
||||
if definition.transfer.source_client_id != self.client_id:
|
||||
raise JobExecutionError("source stage was sent to the wrong client")
|
||||
resource = self._resource(definition)
|
||||
indices = _selection_indices(definition.transfer.transfer_delta_files)
|
||||
by_index = {item.file_index: item for item in resource.files}
|
||||
if not indices or any(index not in by_index for index in indices):
|
||||
raise JobExecutionError("transfer delta file selection is invalid")
|
||||
if any(
|
||||
not by_index[index].selected
|
||||
or by_index[index].completed_bytes != by_index[index].logical_bytes
|
||||
for index in indices
|
||||
):
|
||||
raise JobExecutionError("transfer delta contains unavailable files")
|
||||
manifest = transfer_pb2.TransferManifest(
|
||||
manifest_version=1,
|
||||
job_id=definition.job_id,
|
||||
source_client_id=definition.transfer.source_client_id,
|
||||
target_client_id=definition.transfer.target_client_id,
|
||||
route_id=definition.transfer.route_id,
|
||||
)
|
||||
manifest.resource_id.CopyFrom(definition.resource_id)
|
||||
manifest.requested_files.CopyFrom(definition.transfer.requested_files)
|
||||
manifest.target_baseline_files.CopyFrom(
|
||||
definition.transfer.target_baseline_files
|
||||
)
|
||||
manifest.transfer_delta_files.CopyFrom(
|
||||
definition.transfer.transfer_delta_files
|
||||
)
|
||||
manifest.source_fingerprint.CopyFrom(
|
||||
definition.transfer.source_fingerprint
|
||||
)
|
||||
if definition.transfer.HasField("target_baseline_fingerprint"):
|
||||
manifest.target_baseline_fingerprint.CopyFrom(
|
||||
definition.transfer.target_baseline_fingerprint
|
||||
)
|
||||
manifest.created_at.GetCurrentTime()
|
||||
for index in indices:
|
||||
item = by_index[index]
|
||||
manifest.files.add(
|
||||
file_index=index,
|
||||
payload_relative_path=f"payload/{item.canonical_path}",
|
||||
target_canonical_path=item.canonical_path,
|
||||
logical_bytes=item.logical_bytes,
|
||||
)
|
||||
metainfo_path = self._metainfo_path(definition.job_id)
|
||||
metainfo_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not metainfo_path.exists():
|
||||
metainfo_path.write_bytes(resource.metainfo_bytes)
|
||||
artifact = manifest.artifacts.add(
|
||||
kind=transfer_pb2.ARTIFACT_KIND_TORRENT_FILE,
|
||||
payload_relative_path="metainfo/source.torrent",
|
||||
logical_bytes=len(resource.metainfo_bytes),
|
||||
format="application/x-bittorrent",
|
||||
sha256_hex=hashlib.sha256(resource.metainfo_bytes).hexdigest(),
|
||||
)
|
||||
del artifact
|
||||
stage_transfer(
|
||||
manifest,
|
||||
source_root=self.qb_root,
|
||||
sync_root=self.route_path(definition.transfer.route_id),
|
||||
store=self.store,
|
||||
artifact_sources={"metainfo/source.torrent": metainfo_path},
|
||||
sparse_supported=self.sparse_supported,
|
||||
)
|
||||
self._observer(definition).rescan()
|
||||
|
||||
def _wait_for_syncthing(self, definition: job_pb2.JobDefinition) -> None:
|
||||
if definition.transfer.target_client_id != self.client_id:
|
||||
raise JobExecutionError(
|
||||
"Syncthing completion was sent to the wrong client"
|
||||
)
|
||||
observer = self._observer(definition)
|
||||
while True:
|
||||
try:
|
||||
if observer.status().complete:
|
||||
return
|
||||
except (FileNotFoundError, TransferIntegrityError):
|
||||
# Syncthing can expose the per-job directory before the
|
||||
# ready marker and manifest have arrived atomically as a set.
|
||||
pass
|
||||
time.sleep(self.poll_interval)
|
||||
|
||||
def _target_materialize(self, definition: job_pb2.JobDefinition) -> None:
|
||||
if definition.transfer.target_client_id != self.client_id:
|
||||
raise JobExecutionError(
|
||||
"target materialization was sent to the wrong client"
|
||||
)
|
||||
published = load_published_transfer(self._job_directory(definition))
|
||||
materialize_transfer(
|
||||
published,
|
||||
target_root=self.qb_root,
|
||||
store=self.store,
|
||||
sparse_supported=self.sparse_supported,
|
||||
)
|
||||
|
||||
def _qb_verify(self, definition: job_pb2.JobDefinition):
|
||||
if definition.transfer.target_client_id != self.client_id:
|
||||
raise JobExecutionError("qB verification was sent to the wrong client")
|
||||
published = load_published_transfer(self._job_directory(definition))
|
||||
metainfo_path = (
|
||||
published.job_directory / "metainfo" / "source.torrent"
|
||||
)
|
||||
info_hash = (
|
||||
definition.resource_id.info_hash_v1_hex
|
||||
or definition.resource_id.info_hash_v2_hex
|
||||
)
|
||||
resource = self.qbittorrent.get_resource(info_hash)
|
||||
should_start = (
|
||||
resource is None
|
||||
or resource.summary.runtime_state
|
||||
!= resource_pb2.TORRENT_RUNTIME_STATE_STOPPED
|
||||
)
|
||||
if resource is None:
|
||||
self.qbittorrent.add_stopped_with_retry(
|
||||
metainfo_path.read_bytes(),
|
||||
self.qb_api_root.as_posix(),
|
||||
info_hash,
|
||||
)
|
||||
target_union = _union_selection(
|
||||
definition.transfer.target_baseline_files,
|
||||
definition.transfer.transfer_delta_files,
|
||||
)
|
||||
selected = _selection_indices(target_union)
|
||||
total_files = len(
|
||||
decode_message_metainfo(metainfo_path).files
|
||||
)
|
||||
self.qbittorrent.set_selection(info_hash, selected, total_files)
|
||||
self.qbittorrent.recheck_and_wait(
|
||||
info_hash, selected, timeout=30 * 60, poll_interval=self.poll_interval
|
||||
)
|
||||
if should_start:
|
||||
self.qbittorrent.start(info_hash)
|
||||
verified = self.qbittorrent.get_resource(info_hash)
|
||||
if verified is None:
|
||||
raise JobExecutionError(
|
||||
"verified qBittorrent resource disappeared"
|
||||
)
|
||||
placement = resource_pb2.Placement(
|
||||
client_id=self.client_id,
|
||||
state=resource_pb2.PLACEMENT_STATE_PRESENT,
|
||||
generation=(
|
||||
definition.transfer.expected_target_placement_generation + 1
|
||||
if definition.transfer.HasField(
|
||||
"expected_target_placement_generation"
|
||||
)
|
||||
else 1
|
||||
),
|
||||
created_by_job_id=definition.job_id,
|
||||
)
|
||||
placement.resource_id.CopyFrom(definition.resource_id)
|
||||
placement.verified_files.CopyFrom(target_union)
|
||||
metainfo = decode_message_metainfo(metainfo_path)
|
||||
placement.verified_logical_bytes = sum(
|
||||
metainfo.files[index].logical_bytes for index in selected
|
||||
)
|
||||
placement.fingerprint.CopyFrom(
|
||||
_resource_fingerprint(verified, self.client_id)
|
||||
)
|
||||
placement.verified_at.GetCurrentTime()
|
||||
return placement
|
||||
|
||||
def _observer(
|
||||
self, definition: job_pb2.JobDefinition
|
||||
) -> SyncthingTransferObserver:
|
||||
return SyncthingTransferObserver(
|
||||
self.syncthing_transport,
|
||||
definition.transfer.route_id,
|
||||
f".archive-control/jobs/{definition.job_id}",
|
||||
self._job_directory(definition),
|
||||
)
|
||||
|
||||
def _job_directory(self, definition: job_pb2.JobDefinition) -> Path:
|
||||
return (
|
||||
self.route_path(definition.transfer.route_id)
|
||||
/ ".archive-control" / "jobs" / definition.job_id
|
||||
)
|
||||
|
||||
def _metainfo_path(self, job_id: str) -> Path:
|
||||
return self.store.database.parent / "artifacts" / job_id / "source.torrent"
|
||||
|
||||
def _resource(self, definition: job_pb2.JobDefinition) -> NormalizedResource:
|
||||
for info_hash in (
|
||||
definition.resource_id.info_hash_v1_hex,
|
||||
definition.resource_id.info_hash_v2_hex,
|
||||
):
|
||||
if info_hash:
|
||||
resource = self.qbittorrent.get_resource(info_hash)
|
||||
if resource is not None:
|
||||
return resource
|
||||
raise JobExecutionError("source resource is not present")
|
||||
|
||||
def _definition(self, job_id: str) -> job_pb2.JobDefinition:
|
||||
rows = self.store.job_snapshot_rows([job_id])
|
||||
if not rows:
|
||||
raise JobExecutionError("job is not assigned")
|
||||
return decode_message(
|
||||
str(rows[0]["definition_json"]), job_pb2.JobDefinition()
|
||||
)
|
||||
|
||||
def _committed(self, job_id: str) -> bool:
|
||||
rows = self.store.job_snapshot_rows([job_id])
|
||||
return bool(rows and rows[0]["committed"])
|
||||
|
||||
def _replay(
|
||||
self, job_id: str, after_sequence: int
|
||||
) -> list[control_pb2.JobEvent]:
|
||||
return [
|
||||
decode_message(str(row["event_json"]), control_pb2.JobEvent())
|
||||
for row in self.store.job_event_rows(job_id, after_sequence)
|
||||
]
|
||||
|
||||
def _record(
|
||||
self,
|
||||
definition: job_pb2.JobDefinition,
|
||||
event: control_pb2.JobEvent,
|
||||
) -> None:
|
||||
self.store.record_job_event(
|
||||
job_id=definition.job_id,
|
||||
definition_json=encode_message(definition),
|
||||
event_id=event.event_id,
|
||||
event_json=encode_message(event),
|
||||
state=job_pb2.JobState.Name(event.state),
|
||||
revision=event.job_revision,
|
||||
sequence=event.sequence,
|
||||
committed=event.committed,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _event(
|
||||
definition: job_pb2.JobDefinition,
|
||||
*,
|
||||
sequence: int,
|
||||
revision: int,
|
||||
event_type: int,
|
||||
state: int,
|
||||
committed: bool,
|
||||
step: int = job_pb2.JOB_STEP_KIND_UNSPECIFIED,
|
||||
step_state: int = job_pb2.STEP_STATE_UNSPECIFIED,
|
||||
) -> control_pb2.JobEvent:
|
||||
event = control_pb2.JobEvent(
|
||||
event_id=str(uuid.uuid4()),
|
||||
job_id=definition.job_id,
|
||||
sequence=sequence,
|
||||
job_revision=revision,
|
||||
type=event_type,
|
||||
state=state,
|
||||
committed=committed,
|
||||
)
|
||||
event.occurred_at.GetCurrentTime()
|
||||
if step != job_pb2.JOB_STEP_KIND_UNSPECIFIED:
|
||||
event.progress.step = step
|
||||
event.progress.state = step_state
|
||||
event.progress.fraction_complete = (
|
||||
1 if step_state == job_pb2.STEP_STATE_SUCCEEDED else 0
|
||||
)
|
||||
event.progress.overall_fraction_complete = _overall(step, step_state)
|
||||
event.progress.display_step_number = _step_number(step)
|
||||
event.progress.display_step_total = 5
|
||||
event.progress.last_progress_at.CopyFrom(event.occurred_at)
|
||||
return event
|
||||
|
||||
@staticmethod
|
||||
def _validate_definition(definition: job_pb2.JobDefinition) -> None:
|
||||
if (
|
||||
not definition.job_id
|
||||
or definition.operation not in {
|
||||
job_pb2.JOB_OPERATION_ARCHIVE,
|
||||
job_pb2.JOB_OPERATION_UNARCHIVE,
|
||||
}
|
||||
or definition.WhichOneof("spec") != "transfer"
|
||||
):
|
||||
raise JobExecutionError("transfer job definition is invalid")
|
||||
|
||||
|
||||
def _selection_indices(selection) -> list[int]:
|
||||
result: list[int] = []
|
||||
previous = -1
|
||||
for item in selection.ranges:
|
||||
if item.first > item.last or item.first <= previous:
|
||||
raise JobExecutionError("selection ranges are invalid")
|
||||
result.extend(range(item.first, item.last + 1))
|
||||
previous = item.last
|
||||
return result
|
||||
|
||||
|
||||
def _union_selection(*selections) -> resource_pb2.SelectionSet:
|
||||
indices = sorted({
|
||||
index
|
||||
for selection in selections
|
||||
for index in _selection_indices(selection)
|
||||
})
|
||||
result = resource_pb2.SelectionSet()
|
||||
if not indices:
|
||||
return result
|
||||
first = previous = indices[0]
|
||||
for index in indices[1:]:
|
||||
if index == previous + 1:
|
||||
previous = index
|
||||
continue
|
||||
result.ranges.add(first=first, last=previous)
|
||||
first = previous = index
|
||||
result.ranges.add(first=first, last=previous)
|
||||
return result
|
||||
|
||||
|
||||
def _step_number(step: int) -> int:
|
||||
return {
|
||||
job_pb2.JOB_STEP_KIND_SOURCE_STAGE: 1,
|
||||
job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER: 2,
|
||||
job_pb2.JOB_STEP_KIND_TARGET_MATERIALIZE: 3,
|
||||
job_pb2.JOB_STEP_KIND_QB_VERIFY: 4,
|
||||
job_pb2.JOB_STEP_KIND_STAGING_CLEANUP: 5,
|
||||
}.get(step, 0)
|
||||
|
||||
|
||||
def _overall(step: int, state: int) -> float:
|
||||
number = _step_number(step)
|
||||
return max(0, number - (0 if state == job_pb2.STEP_STATE_SUCCEEDED else 1)) / 5
|
||||
|
||||
|
||||
def decode_message_metainfo(path: Path):
|
||||
from archive_clients.bencode import decode_metainfo
|
||||
|
||||
return decode_metainfo(path.read_bytes())
|
||||
|
||||
|
||||
def _resource_fingerprint(
|
||||
resource: NormalizedResource, client_id: str
|
||||
) -> resource_pb2.ResourceStateFingerprint:
|
||||
summary = resource.summary
|
||||
fingerprint = resource_pb2.ResourceStateFingerprint(
|
||||
client_id=client_id,
|
||||
qb_torrent_id=summary.qb_torrent_id,
|
||||
content_revision=summary.content_revision,
|
||||
runtime_state=summary.runtime_state,
|
||||
)
|
||||
fingerprint.resource_id.CopyFrom(summary.resource_id)
|
||||
fingerprint.selected_files.CopyFrom(summary.selected_files)
|
||||
fingerprint.selected_complete_files.CopyFrom(
|
||||
summary.selected_complete_files
|
||||
)
|
||||
fingerprint.observed_at.CopyFrom(summary.observed_at)
|
||||
return fingerprint
|
||||
|
||||
|
||||
def _job_error_code(error: Exception) -> int:
|
||||
if isinstance(error, QBittorrentDownloadAttempt):
|
||||
return common_pb2.ERROR_CODE_MANUAL_INTERVENTION_REQUIRED
|
||||
if isinstance(error, TransferIntegrityError):
|
||||
return common_pb2.ERROR_CODE_INTEGRITY_CHECK_FAILED
|
||||
if isinstance(error, PermissionError):
|
||||
return common_pb2.ERROR_CODE_PERMISSION_DENIED
|
||||
if isinstance(error, JobExecutionError):
|
||||
return common_pb2.ERROR_CODE_PRECONDITION_FAILED
|
||||
if isinstance(error, TransferError):
|
||||
return common_pb2.ERROR_CODE_PATH_CONFLICT
|
||||
return common_pb2.ERROR_CODE_INTERNAL
|
||||
Reference in New Issue
Block a user