238 lines
8.0 KiB
Python
238 lines
8.0 KiB
Python
"""Thread-safe reference ledger for job/run and callback correlation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Dict, Optional, Protocol, Tuple
|
|
|
|
from arr_processing.contracts import ProcessingRequest, ProcessingResult
|
|
from arr_processing.errors import ProcessingError
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProcessingRecord:
|
|
request: ProcessingRequest
|
|
request_sha256: str
|
|
idempotency_key: str
|
|
remote_run_id: Optional[str]
|
|
remote_status: str
|
|
result: Optional[ProcessingResult] = None
|
|
failure_code: Optional[str] = None
|
|
|
|
|
|
class ProcessingLedger(Protocol):
|
|
def reserve(
|
|
self,
|
|
request: ProcessingRequest,
|
|
request_sha256: str,
|
|
idempotency_key: str,
|
|
) -> ProcessingRecord:
|
|
...
|
|
|
|
def get(self, job_id: str, attempt_no: int) -> Optional[ProcessingRecord]:
|
|
...
|
|
|
|
def bind_run(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
remote_run_id: str,
|
|
remote_status: str,
|
|
) -> ProcessingRecord:
|
|
...
|
|
|
|
def update_status(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
remote_status: str,
|
|
) -> ProcessingRecord:
|
|
...
|
|
|
|
def fail_start(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
failure_code: str,
|
|
) -> ProcessingRecord:
|
|
...
|
|
|
|
def attach_result(self, result: ProcessingResult) -> Tuple[ProcessingRecord, bool]:
|
|
...
|
|
|
|
def accept_callback(self, delivery_id: str, signed_sha256: str) -> bool:
|
|
...
|
|
|
|
|
|
class InMemoryProcessingLedger:
|
|
"""Reference implementation; production should map this protocol to PostgreSQL."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.RLock()
|
|
self._records: Dict[Tuple[str, int], ProcessingRecord] = {}
|
|
self._runs: Dict[str, Tuple[str, int]] = {}
|
|
self._callbacks: Dict[str, str] = {}
|
|
|
|
def reserve(
|
|
self,
|
|
request: ProcessingRequest,
|
|
request_sha256: str,
|
|
idempotency_key: str,
|
|
) -> ProcessingRecord:
|
|
key = (request.job_id, request.attempt_no)
|
|
with self._lock:
|
|
existing = self._records.get(key)
|
|
if existing is not None:
|
|
if (
|
|
existing.request != request
|
|
or existing.request_sha256 != request_sha256
|
|
or existing.idempotency_key != idempotency_key
|
|
):
|
|
raise ProcessingError(
|
|
"PROCESSING_JOB_CONFLICT", "processing attempt identity conflicts"
|
|
)
|
|
return copy.deepcopy(existing)
|
|
record = ProcessingRecord(
|
|
request=request,
|
|
request_sha256=request_sha256,
|
|
idempotency_key=idempotency_key,
|
|
remote_run_id=None,
|
|
remote_status="reserved",
|
|
)
|
|
self._records[key] = record
|
|
return copy.deepcopy(record)
|
|
|
|
def get(self, job_id: str, attempt_no: int) -> Optional[ProcessingRecord]:
|
|
with self._lock:
|
|
record = self._records.get((job_id, attempt_no))
|
|
return copy.deepcopy(record) if record is not None else None
|
|
|
|
def bind_run(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
remote_run_id: str,
|
|
remote_status: str,
|
|
) -> ProcessingRecord:
|
|
key = (job_id, attempt_no)
|
|
with self._lock:
|
|
record = self._required(key)
|
|
run_owner = self._runs.get(remote_run_id)
|
|
if run_owner is not None and run_owner != key:
|
|
raise ProcessingError(
|
|
"PROCESSING_RUN_CONFLICT", "remote run is already correlated"
|
|
)
|
|
if record.remote_run_id is not None and record.remote_run_id != remote_run_id:
|
|
raise ProcessingError(
|
|
"PROCESSING_RUN_CONFLICT", "processing attempt has a different remote run"
|
|
)
|
|
updated = ProcessingRecord(
|
|
request=record.request,
|
|
request_sha256=record.request_sha256,
|
|
idempotency_key=record.idempotency_key,
|
|
remote_run_id=remote_run_id,
|
|
remote_status=remote_status,
|
|
result=record.result,
|
|
failure_code=record.failure_code,
|
|
)
|
|
self._records[key] = updated
|
|
self._runs[remote_run_id] = key
|
|
return copy.deepcopy(updated)
|
|
|
|
def update_status(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
remote_status: str,
|
|
) -> ProcessingRecord:
|
|
key = (job_id, attempt_no)
|
|
with self._lock:
|
|
record = self._required(key)
|
|
updated = ProcessingRecord(
|
|
request=record.request,
|
|
request_sha256=record.request_sha256,
|
|
idempotency_key=record.idempotency_key,
|
|
remote_run_id=record.remote_run_id,
|
|
remote_status=remote_status,
|
|
result=record.result,
|
|
failure_code=record.failure_code,
|
|
)
|
|
self._records[key] = updated
|
|
return copy.deepcopy(updated)
|
|
|
|
def fail_start(
|
|
self,
|
|
job_id: str,
|
|
attempt_no: int,
|
|
failure_code: str,
|
|
) -> ProcessingRecord:
|
|
key = (job_id, attempt_no)
|
|
with self._lock:
|
|
record = self._required(key)
|
|
if record.remote_run_id is not None:
|
|
raise ProcessingError(
|
|
"PROCESSING_RUN_CONFLICT",
|
|
"processing attempt already has a remote run",
|
|
)
|
|
if record.remote_status == "failed":
|
|
return copy.deepcopy(record)
|
|
updated = ProcessingRecord(
|
|
request=record.request,
|
|
request_sha256=record.request_sha256,
|
|
idempotency_key=record.idempotency_key,
|
|
remote_run_id=None,
|
|
remote_status="failed",
|
|
result=record.result,
|
|
failure_code=failure_code,
|
|
)
|
|
self._records[key] = updated
|
|
return copy.deepcopy(updated)
|
|
|
|
def attach_result(self, result: ProcessingResult) -> Tuple[ProcessingRecord, bool]:
|
|
key = (result.job_id, result.attempt_no)
|
|
with self._lock:
|
|
record = self._required(key)
|
|
if record.remote_run_id != result.remote_run_id:
|
|
raise ProcessingError(
|
|
"PROCESSING_RESULT_MISMATCH", "result remote run does not match its attempt"
|
|
)
|
|
if record.result is not None:
|
|
if record.result != result:
|
|
raise ProcessingError(
|
|
"PROCESSING_RESULT_CONFLICT", "processing attempt has a different result"
|
|
)
|
|
return copy.deepcopy(record), True
|
|
updated = ProcessingRecord(
|
|
request=record.request,
|
|
request_sha256=record.request_sha256,
|
|
idempotency_key=record.idempotency_key,
|
|
remote_run_id=record.remote_run_id,
|
|
remote_status="delivered",
|
|
result=result,
|
|
failure_code=None,
|
|
)
|
|
self._records[key] = updated
|
|
return copy.deepcopy(updated), False
|
|
|
|
def accept_callback(self, delivery_id: str, signed_sha256: str) -> bool:
|
|
with self._lock:
|
|
existing = self._callbacks.get(delivery_id)
|
|
if existing is None:
|
|
self._callbacks[delivery_id] = signed_sha256
|
|
return False
|
|
if existing != signed_sha256:
|
|
raise ProcessingError(
|
|
"PROCESSING_CALLBACK_CONFLICT", "delivery identifier conflicts"
|
|
)
|
|
return True
|
|
|
|
def _required(self, key: Tuple[str, int]) -> ProcessingRecord:
|
|
try:
|
|
return self._records[key]
|
|
except KeyError:
|
|
raise ProcessingError(
|
|
"PROCESSING_JOB_NOT_FOUND", "processing attempt is not registered"
|
|
) from None
|