diff --git a/.coveragerc b/.coveragerc index c9001ad..a588790 100644 --- a/.coveragerc +++ b/.coveragerc @@ -1,4 +1,12 @@ [run] source = deadpool branch = True -#concurrency=multiprocessing \ No newline at end of file +relative_files = True +parallel = True +concurrency = + multiprocessing + thread +patch = + fork + subprocess + _exit diff --git a/README.rst b/README.rst index e1a808b..fa0ef20 100644 --- a/README.rst +++ b/README.rst @@ -763,6 +763,91 @@ when creating the instance: ): fut = exe.submit(...) +Remote executor (experimental) +============================== + +``deadpool.remote`` lets independent processes share one server-owned +``Deadpool`` over a Unix-domain or IPv4 TCP socket. The client implements the +standard ``Executor`` interface, including ordered ``map()`` results: + +.. code-block:: python + + from functools import partial + + import deadpool + from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + UnixAddress, + UnixListener, + ) + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=8), + listeners=[UnixListener("/run/example/deadpool.sock")], + ) + server.start() + + with DeadpoolClient(UnixAddress("/run/example/deadpool.sock")) as executor: + future = executor.submit(pow, 2, 10, deadpool_priority=5) + assert future.result() == 1024 + +A server may also expose registered operations. Registered mode rejects unknown +operation names before worker dispatch: + +.. code-block:: python + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=8), + listeners=[UnixListener("/run/example/deadpool.sock")], + task_registry={"math.pow": pow}, + ) + future = executor.submit_task("math.pow", 2, 10) + +TCP uses the same protocol. Plaintext is restricted to an explicit loopback +opt-in (``TcpListener("127.0.0.1", port, insecure=True)``); other deployments +must supply configured server and client ``ssl.SSLContext`` objects. Mutual TLS +is recommended. + +Client and server wire bounds for control data, frames, messages, metadata, and +chunk counts must match. The handshake rejects asymmetric wire limits before +accepting submissions. + +``RemoteFuture.add_done_callback()`` follows standard ``Future`` timing: +callbacks registered after terminal completion run synchronously before the +method returns, while callbacks registered earlier use the isolated callback +executor. ``callback_queue_size`` bounds callback batches admitted to that +executor, not application-owned registrations waiting for dispatch; callback +backpressure never blocks transport or Future state updates. + +Remote cancellation requires an authoritative broker decision. Unlike a local +``Future.cancel()``, cancelling submitted work may therefore block up to the +client's ``control_timeout`` or raise ``RemoteCancellationOutcomeUnknown`` when +the distributed outcome cannot be determined. Work still in the client's local +submission stage cancels immediately. + +.. warning:: + + The default ``PickleSerializer`` requires mutually trusted clients and + servers in both callable and registered-task modes. TLS authenticates and + encrypts peers but does not sandbox deserialization. Registered operation + authorization controls callable selection; it is not a deserialization + sandbox because workers still unpickle invocation arguments. Clients also + unpickle result and exception payloads, establishing reverse trust in the + server. + +The current framed wire layout is private and experimental because the remote +executor specification intentionally leaves stable frame octets and canonical +schemas to a later wire reference. Matching Deadpool releases support bounded +JSON control headers, opaque chunked payloads, request identity, typed failures, +ordinary and hard cancellation, health/statistics requests, fair broker +scheduling, and explicit shutdown behavior. The first implementation does not +negotiate session resumption, compression, IPv6, protocol-5 out-of-band buffers, +shared-memory/file result transports, durable state, progress events, or +streaming results. It also does not advertise stronger parent-death cleanup than +the local pool's existing process monitor. A lost unacknowledged submission is +reported as ambiguous and is never silently retried with a new request ID. + Developer Workflow ================== @@ -837,7 +922,7 @@ I currently do: 1. Ensure that ``main`` branch is fully up to date with all to be released, and all the tests succeed. -2. Change the ``__version__`` field in ``deadpool.py``. Flit +2. Change the ``__version__`` field in ``deadpool/__init__.py``. Flit uses this to stamp the version. 3. Verify that ``flit build`` succeeds. This will produce a wheel in the ``dist/`` directory. You can inspect this diff --git a/deadpool/__init__.py b/deadpool/__init__.py new file mode 100644 index 0000000..b49969b --- /dev/null +++ b/deadpool/__init__.py @@ -0,0 +1,36 @@ +"""A resilient process-pool executor and its optional remote interface.""" + +__version__ = "2026.6.1" + +from . import _pool +from ._pool import ( + CancelledError, + Deadpool, + Future, + PoolClosed, + ProcessError, + TimeoutError, + as_completed, +) + +__all__ = [ + "Deadpool", + "Future", + "CancelledError", + "TimeoutError", + "ProcessError", + "PoolClosed", + "as_completed", +] + + +def __getattr__(name: str): + """Preserve access to implementation helpers exposed by the old module.""" + try: + return getattr(_pool, name) + except AttributeError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(dir(_pool))) diff --git a/deadpool.py b/deadpool/_pool.py similarity index 94% rename from deadpool.py rename to deadpool/_pool.py index 0a783f9..b032f68 100644 --- a/deadpool.py +++ b/deadpool/_pool.py @@ -21,6 +21,7 @@ import concurrent.futures import ctypes +import itertools import logging import multiprocessing as mp import os @@ -32,7 +33,6 @@ import typing import weakref import atexit -import json from concurrent.futures import CancelledError, Executor, InvalidStateError, as_completed from dataclasses import dataclass, field from multiprocessing.connection import Connection @@ -44,7 +44,8 @@ import psutil from setproctitle import setproctitle -__version__ = "2026.6.1" +from . import __version__ as __version__ + __all__ = [ "Deadpool", "Future", @@ -117,6 +118,7 @@ def to_dict(self) -> dict[str, typing.Any]: class PrioritizedItem: priority: int item: typing.Any = field(compare=False) + sequence: int = 0 @dataclass(init=False) @@ -276,6 +278,11 @@ def pid(self, value): def add_pid_callback(self, fn): self.pid_callback = fn + if self.pid is not None: + try: + fn(self) + except Exception: # pragma: no cover + logger.exception("Error calling pid_callback") def cancel_and_kill_if_running(self, sig=signal.SIGKILL): self.cancel() @@ -376,6 +383,7 @@ def __init__( malloc_trim_rss_memory_threshold_bytes ) self._statistics = Statistics() + self._submission_sequence = itertools.count() # TODO: overcommit self.workers: SimpleQueue[WorkerProcess] = SimpleQueue() @@ -558,7 +566,20 @@ def run_task(self, fn, args, kwargs, timeout, fut: Future): retry_count -= 1 worker: WorkerProcess = self.get_process() try: - worker.submit_job((fn, args, kwargs, timeout)) + before_submit = getattr(fut, "_before_submit", None) + if before_submit is not None and not before_submit(worker.pid): + self.done_with_process(worker) + return + try: + worker.submit_job((fn, args, kwargs, timeout)) + except BaseException: + submit_failed = getattr(fut, "_submit_failed", None) + if submit_failed is not None: + submit_failed(worker.pid) + raise + after_submit = getattr(fut, "_after_submit", None) + if after_submit is not None: + after_submit(worker.pid) break except (pickle.PicklingError, AttributeError) as e: # If the user passed in a function or params that can't @@ -690,11 +711,43 @@ def submit( raise PoolClosed("The pool is closed. No more tasks can be submitted.") fut = Future() + return self._submit_future( + fut, + fn, + args, + kwargs, + deadpool_timeout, + deadpool_priority, + ) + + def _submit_future( + self, + fut: Future, + fn: Callable, + args: tuple, + kwargs: dict, + timeout, + priority, + *, + block: bool = True, + ) -> Future: + """Submit using a framework-owned Future. + + Remote brokerage uses this narrow private seam to arbitrate queued + cancellation against worker dispatch. Its broker lane may request + nonblocking admission so remote deadlines remain live when the local + backlog is full. Ordinary :meth:`submit` calls retain blocking + backpressure. + """ + if self.closed: + raise PoolClosed("The pool is closed. No more tasks can be submitted.") self.submitted_jobs.put( PrioritizedItem( - priority=deadpool_priority, - item=(fn, args, kwargs, deadpool_timeout, fut), - ) + priority=priority, + item=(fn, args, kwargs, timeout, fut), + sequence=next(self._submission_sequence), + ), + block=block, ) self._statistics.tasks_received.increment() return fut @@ -720,7 +773,11 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: try: self.submitted_jobs.put( - PrioritizedItem(priority=shutdown_priority, item=None), + PrioritizedItem( + priority=shutdown_priority, + item=None, + sequence=next(self._submission_sequence), + ), timeout=2.0, ) except TimeoutError: # pragma: no cover diff --git a/deadpool/remote/__init__.py b/deadpool/remote/__init__.py new file mode 100644 index 0000000..dd7525a --- /dev/null +++ b/deadpool/remote/__init__.py @@ -0,0 +1,95 @@ +"""Experimental socket-based access to a server-owned Deadpool. + +Security warning +---------------- +The default PickleSerializer requires mutually trusted clients and servers in +both callable and registered-task modes. TLS authenticates peers but does not +sandbox deserialization. Registered operation authorization is not a +serialization sandbox because workers still unpickle invocation arguments; +clients likewise unpickle result and exception payloads, creating reverse trust. + +The current wire layout is private and versioned but not a stable interoperability +contract. Session resumption, compression, IPv6, out-of-band pickle buffers, +and durable request state are intentionally not negotiated by this release. +""" + +from ._future import RemoteFuture, SubmissionState +from .client import DeadpoolClient +from .config import ( + Authenticator, + Authorizer, + PeerInfo, + Principal, + RemoteLimits, + TcpAddress, + TcpListener, + UnixAddress, + UnixListener, +) +from .errors import ( + AcceptanceCertainty, + ExecutionCertainty, + RemoteAuthenticationError, + RemoteCancellationOutcomeUnknown, + RemoteCancelledError, + RemoteCompatibilityError, + RemoteConnectionLost, + RemoteExecutorError, + RemoteExecutorUnavailable, + RemoteForkedProcessError, + RemoteProcessError, + RemoteProtocolError, + RemoteQueueFull, + RemoteQueueTimeout, + RemoteResultEncodingError, + RemoteResultLost, + RemoteResultTooLarge, + RemoteServerRestarted, + RemoteSubmissionTimeout, + RemoteTaskError, + SubmissionOutcomeUnknown, +) +from .serializer import PickleSerializer, SerializationLimitError, Serializer +from .server import DeadpoolServer, RequestState, ServerState + +__all__ = [ + "DeadpoolClient", + "DeadpoolServer", + "RemoteFuture", + "SubmissionState", + "RequestState", + "ServerState", + "UnixAddress", + "TcpAddress", + "UnixListener", + "TcpListener", + "RemoteLimits", + "PeerInfo", + "Principal", + "Authenticator", + "Authorizer", + "Serializer", + "PickleSerializer", + "SerializationLimitError", + "AcceptanceCertainty", + "ExecutionCertainty", + "RemoteExecutorError", + "RemoteExecutorUnavailable", + "RemoteAuthenticationError", + "RemoteCompatibilityError", + "RemoteQueueFull", + "RemoteSubmissionTimeout", + "SubmissionOutcomeUnknown", + "RemoteQueueTimeout", + "RemoteTaskError", + "RemoteResultEncodingError", + "RemoteResultTooLarge", + "RemoteCancellationOutcomeUnknown", + "RemoteForkedProcessError", + "RemoteProcessError", + "RemoteCancelledError", + "RemoteConnectionLost", + "RemoteServerRestarted", + "RemoteProtocolError", + "RemoteResultLost", +] diff --git a/deadpool/remote/_future.py b/deadpool/remote/_future.py new file mode 100644 index 0000000..a3dd2d5 --- /dev/null +++ b/deadpool/remote/_future.py @@ -0,0 +1,197 @@ +"""Process-local Future representing one remote request.""" + +from __future__ import annotations + +import concurrent.futures +import logging +import os +import threading +import weakref +from enum import Enum +from typing import TYPE_CHECKING, Callable + +from .errors import RemoteForkedProcessError + +if TYPE_CHECKING: + from .client import DeadpoolClient + +logger = logging.getLogger("deadpool.remote") + + +class SubmissionState(str, Enum): + LOCAL_PENDING = "LOCAL_PENDING" + SENT_UNACKNOWLEDGED = "SENT_UNACKNOWLEDGED" + ACCEPTED_QUEUED = "ACCEPTED_QUEUED" + RUNNING = "RUNNING" + SUCCEEDED = "SUCCEEDED" + TASK_FAILED = "TASK_FAILED" + CANCELLED = "CANCELLED" + FAILED = "FAILED" + + +class RemoteFuture(concurrent.futures.Future): + """A Future with request identity and authoritative remote state.""" + + def __init__(self, request_id: str, client: DeadpoolClient) -> None: + super().__init__() + self.request_id = request_id + self.submission_state = SubmissionState.LOCAL_PENDING + self.worker_id: str | None = None + self._pid: int | None = None + self._owner_pid = os.getpid() + self._client = weakref.ref(client) + self._state_lock = threading.RLock() + self._remote_callbacks: list[Callable[[RemoteFuture], object]] = [] + + @property + def pid(self) -> int | None: + self._check_process() + return self._pid + + def cancel(self) -> bool: + """Request an authoritative cancellation decision from the broker. + + Locally pending work cancels immediately. Once submitted, this call may + block up to the client's control timeout or raise when the distributed + cancellation outcome cannot be determined. + """ + self._check_process() + with self._state_lock: + if self.done(): + return False + if self.submission_state == SubmissionState.LOCAL_PENDING: + self.submission_state = SubmissionState.CANCELLED + cancelled = self._cancel_and_notify_waiters() + self._dispatch_callbacks() + return cancelled + client = self._client() + if client is None: + return False + return client._cancel(self, hard=False) + + def cancel_and_kill_if_running(self) -> bool: + self._check_process() + client = self._client() + if client is None or self.done(): + return False + return client._cancel(self, hard=True) + + def add_done_callback(self, fn: Callable[[RemoteFuture], object]) -> None: + """Register ``fn``, invoking it inline when already terminal. + + Callbacks registered before completion are dispatched on the client's + isolated callback executor. Matching the standard Future contract, + callbacks registered after completion run synchronously before this + method returns. + """ + self._check_process() + with self._state_lock: + if not self.done(): + self._remote_callbacks.append(fn) + return + self._invoke_callback(fn) + + def result(self, timeout: float | None = None): + self._check_process() + return super().result(timeout) + + def exception(self, timeout: float | None = None): + self._check_process() + return super().exception(timeout) + + def running(self) -> bool: + self._check_process() + with self._state_lock: + return ( + self.submission_state == SubmissionState.RUNNING and not super().done() + ) + + def done(self) -> bool: + self._check_process() + return super().done() + + def cancelled(self) -> bool: + self._check_process() + return super().cancelled() + + def _set_sent(self) -> bool: + with self._state_lock: + if self.done(): + return False + self.submission_state = SubmissionState.SENT_UNACKNOWLEDGED + return True + + def _set_accepted(self) -> None: + with self._state_lock: + if not self.done(): + self.submission_state = SubmissionState.ACCEPTED_QUEUED + + def _set_running(self, *, pid: int | None, worker_id: str | None) -> None: + # Keep the base Future pending. Remote RUNNING is authoritative broker + # metadata, while a later authorized hard cancellation still needs to + # make the local Future genuinely cancelled and terminal. + with self._state_lock: + if super().done(): + return + self._pid = pid + self.worker_id = worker_id + self.submission_state = SubmissionState.RUNNING + + def _set_result(self, value: object) -> None: + with self._state_lock: + if self.done(): + return + self.submission_state = SubmissionState.SUCCEEDED + super().set_result(value) + self._dispatch_callbacks() + + def _set_exception(self, error: BaseException, *, task: bool = False) -> None: + with self._state_lock: + if self.done(): + return + self.submission_state = ( + SubmissionState.TASK_FAILED if task else SubmissionState.FAILED + ) + super().set_exception(error) + self._dispatch_callbacks() + + def _set_cancelled(self) -> None: + with self._state_lock: + if self.done(): + return + self.submission_state = SubmissionState.CANCELLED + self._cancel_and_notify_waiters() + self._dispatch_callbacks() + + def _cancel_and_notify_waiters(self) -> bool: + # stdlib waiters require the executor-side cancellation notification. + cancelled = super().cancel() + if cancelled: + super().set_running_or_notify_cancel() + return cancelled + + def _dispatch_callbacks(self) -> None: + with self._state_lock: + callbacks, self._remote_callbacks = self._remote_callbacks, [] + if not callbacks: + return + client = self._client() + if client is None: + for callback in callbacks: + self._invoke_callback(callback) + return + client._schedule_callbacks(callbacks, self) + + def _invoke_callback(self, callback: Callable[[RemoteFuture], object]) -> None: + """Invoke one callback through the shared exception boundary.""" + try: + callback(self) + except BaseException: + logger.exception("exception calling RemoteFuture callback") + + def _check_process(self) -> None: + if os.getpid() != self._owner_pid: + raise RemoteForkedProcessError( + "RemoteFuture was inherited across fork and is not valid in the child", + request_id=self.request_id, + ) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py new file mode 100644 index 0000000..8e65353 --- /dev/null +++ b/deadpool/remote/_protocol.py @@ -0,0 +1,496 @@ +"""Experimental Deadpool remote wire protocol. + +The behavioral specification deliberately does not freeze wire octets. This +private version is therefore suitable only between matching Deadpool releases; +it must not be advertised as a stable interoperability protocol. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import select +import socket +import ssl +import struct +import time +import uuid +from dataclasses import dataclass +from enum import IntEnum + +from .config import RemoteLimits +from .errors import RemoteProtocolError + +MAGIC = b"DPR1" +MAJOR = 1 +MINOR = 0 +_PREFIX = struct.Struct("!4sBBBBII") +_FLAG_CHUNKED = 1 +_MAX_JSON_DEPTH = 12 +_WIRE_LIMIT_FIELDS = ( + "max_control_bytes", + "max_frame_payload_bytes", + "max_message_bytes", + "max_metadata_bytes", + "max_chunks", +) + + +def _wire_limits(limits: RemoteLimits) -> dict[str, int]: + return {name: getattr(limits, name) for name in _WIRE_LIMIT_FIELDS} + + +def _validate_wire_limits(value: object) -> dict[str, int]: + if not isinstance(value, dict) or set(value) != set(_WIRE_LIMIT_FIELDS): + raise RemoteProtocolError("wire limits have invalid fields") + for name, limit in value.items(): + if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0: + raise RemoteProtocolError(f"wire limit {name} must be a positive integer") + return dict(value) + + +class MessageType(IntEnum): + HELLO = 1 + WELCOME = 2 + HANDSHAKE_REJECTED = 3 + PING = 4 + PONG = 5 + GOAWAY = 6 + PROTOCOL_ERROR = 7 + SUBMIT = 10 + ACCEPTED = 11 + REJECTED = 12 + RUNNING = 13 + RESULT = 20 + TASK_ERROR = 21 + TIMED_OUT = 22 + CANCELLED = 23 + WORKER_LOST = 24 + RESULT_ENCODING_FAILED = 25 + RESULT_TOO_LARGE = 26 + RESULT_ACK = 27 + QUEUE_TIMED_OUT = 28 + CANCEL_REQUEST = 30 + CANCEL_RESPONSE = 31 + STATUS_REQUEST = 32 + STATUS_RESPONSE = 33 + STATS_REQUEST = 34 + STATS_RESPONSE = 35 + CANCEL_GROUP = 36 + CANCEL_GROUP_RESPONSE = 37 + CLOSE_SESSION = 38 + CLOSE_SESSION_RESPONSE = 39 + + +@dataclass(frozen=True, slots=True) +class Message: + kind: MessageType + control: dict + payload: bytes = b"" + + +@dataclass(slots=True) +class _Chunks: + kind: MessageType + control: dict + count: int + total: int + digest: str + next_index: int + parts: list[bytes] + received: int = 0 + + +def _validate_json_nesting(text: str) -> None: + """Reject excessive container nesting without recursing in the JSON decoder.""" + depth = 0 + in_string = False + escaped = False + for character in text: + if in_string: + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character == '"': + in_string = False + elif character == '"': + in_string = True + elif character in "[{": + depth += 1 + # Lexical depth includes the root container, whose semantic depth is zero. + if depth > _MAX_JSON_DEPTH + 1: + raise RemoteProtocolError("control JSON is too deeply nested") + elif character in "]}": + depth -= 1 + + +def _json_loads(data: bytes, limits: RemoteLimits) -> dict: + try: + text = data.decode("utf-8") + _validate_json_nesting(text) + value = json.loads( + text, + object_pairs_hook=_unique_object, + parse_constant=lambda value: (_ for _ in ()).throw( + ValueError(f"non-finite JSON number {value}") + ), + ) + except ( + UnicodeDecodeError, + json.JSONDecodeError, + RecursionError, + ValueError, + ) as error: + raise RemoteProtocolError(f"invalid control JSON: {error}") from error + _validate_json(value, limits) + if not isinstance(value, dict): + raise RemoteProtocolError("control header must be a JSON object") + return value + + +def _unique_object(items: list[tuple[str, object]]) -> dict: + value = {} + for key, item in items: + if key in value: + raise ValueError(f"duplicate JSON key {key!r}") + value[key] = item + return value + + +def _validate_json(value: object, limits: RemoteLimits, depth: int = 0) -> None: + if depth > _MAX_JSON_DEPTH: + raise RemoteProtocolError("control JSON is too deeply nested") + if isinstance(value, str): + if len(value.encode("utf-8")) > limits.max_metadata_bytes: + raise RemoteProtocolError("control string exceeds metadata limit") + elif isinstance(value, bool) or value is None: + return + elif isinstance(value, int): + if not -(1 << 63) <= value < (1 << 63): + raise RemoteProtocolError("control integer exceeds signed 64-bit range") + elif isinstance(value, float): + if not math.isfinite(value): + raise RemoteProtocolError("control number must be finite") + elif isinstance(value, list): + if len(value) > 1024: + raise RemoteProtocolError("control list is too large") + for item in value: + _validate_json(item, limits, depth + 1) + elif isinstance(value, dict): + if len(value) > 128: + raise RemoteProtocolError("control object has too many keys") + for key, item in value.items(): + if not isinstance(key, str): + raise RemoteProtocolError("control object keys must be strings") + _validate_json(key, limits, depth + 1) + _validate_json(item, limits, depth + 1) + else: + raise RemoteProtocolError(f"unsupported control value {type(value).__name__}") + + +def _json_dumps(value: dict, limits: RemoteLimits) -> bytes: + _validate_json(value, limits) + payload = json.dumps( + value, ensure_ascii=True, allow_nan=False, separators=(",", ":"), sort_keys=True + ).encode("ascii") + if len(payload) > limits.max_control_bytes: + raise RemoteProtocolError("control header exceeds configured limit") + return payload + + +def validate_control(control: dict, limits: RemoteLimits) -> None: + """Validate a logical control object without serializing application data.""" + _json_dumps(control, limits) + + +def validate_message_control( + control: dict, payload_size: int, limits: RemoteLimits +) -> None: + """Validate control data inside its complete first-frame envelope.""" + count = _message_chunk_count(payload_size, limits) + _encode_frame_control( + "0" * 32, + index=0, + count=count, + total=payload_size, + digest="0" * 64, + control=control, + limits=limits, + ) + + +def _message_chunk_count(payload_size: int, limits: RemoteLimits) -> int: + if payload_size < 0 or payload_size > limits.max_message_bytes: + raise RemoteProtocolError("message payload exceeds configured limit") + count = max( + 1, + (payload_size + limits.max_frame_payload_bytes - 1) + // limits.max_frame_payload_bytes, + ) + if count > limits.max_chunks: + raise RemoteProtocolError("message requires too many chunks") + return count + + +def _encode_frame_control( + message_id: str, + *, + index: int, + count: int, + total: int, + digest: str, + control: dict | None, + limits: RemoteLimits, +) -> bytes: + return _json_dumps( + { + "message_id": message_id, + "index": index, + "count": count, + "total": total, + "digest": digest, + "control": control, + }, + limits, + ) + + +def send_message( + sock: socket.socket, + message: Message, + limits: RemoteLimits, +) -> None: + """Send one logical message as one or more independently bounded frames.""" + payload = memoryview(message.payload) + frame_size = limits.max_frame_payload_bytes + count = _message_chunk_count(len(payload), limits) + message_id = uuid.uuid4().hex + digest = hashlib.sha256(payload).hexdigest() + for index in range(count): + start = index * frame_size + chunk = payload[start : start + frame_size] + control = _encode_frame_control( + message_id, + index=index, + count=count, + total=len(payload), + digest=digest, + control=message.control if index == 0 else None, + limits=limits, + ) + prefix = _PREFIX.pack( + MAGIC, + MAJOR, + MINOR, + int(message.kind), + _FLAG_CHUNKED if count > 1 else 0, + len(control), + len(chunk), + ) + _send_exact(sock, prefix, limits.partial_frame_timeout) + _send_exact(sock, control, limits.partial_frame_timeout) + if chunk: + _send_exact(sock, chunk, limits.partial_frame_timeout) + + +class MessageReader: + """Incrementally reassemble bounded, potentially interleaved messages.""" + + def __init__(self, limits: RemoteLimits) -> None: + self.limits = limits + self._messages: dict[str, _Chunks] = {} + self._incomplete_bytes = 0 + + def receive( + self, + sock: socket.socket, + *, + deadline: float | None = None, + ) -> Message: + while True: + prefix = _recv_exact( + sock, + _PREFIX.size, + timeout=_bounded_timeout( + self.limits.partial_frame_timeout, + deadline, + ), + allow_idle=True, + ) + magic, major, minor, raw_kind, flags, control_len, payload_len = ( + _PREFIX.unpack(prefix) + ) + if magic != MAGIC: + raise RemoteProtocolError("invalid protocol magic") + if major != MAJOR or minor > MINOR: + raise RemoteProtocolError( + f"unsupported protocol version {major}.{minor}" + ) + if flags & ~_FLAG_CHUNKED: + raise RemoteProtocolError("unknown mandatory frame flags") + if control_len > self.limits.max_control_bytes: + raise RemoteProtocolError("declared control header is too large") + if payload_len > self.limits.max_frame_payload_bytes: + raise RemoteProtocolError("declared frame payload is too large") + try: + kind = MessageType(raw_kind) + except ValueError as error: + raise RemoteProtocolError(f"unknown message type {raw_kind}") from error + frame_control = _json_loads( + _recv_exact( + sock, + control_len, + timeout=_bounded_timeout( + self.limits.partial_frame_timeout, + deadline, + ), + ), + self.limits, + ) + payload = _recv_exact( + sock, + payload_len, + timeout=_bounded_timeout( + self.limits.partial_frame_timeout, + deadline, + ), + ) + complete = self._accept(kind, frame_control, payload) + if complete is not None: + return complete + + def _accept( + self, kind: MessageType, header: dict, payload: bytes + ) -> Message | None: + required = {"message_id", "index", "count", "total", "digest", "control"} + if set(header) != required: + raise RemoteProtocolError("invalid chunk control fields") + message_id = header["message_id"] + index = header["index"] + count = header["count"] + total = header["total"] + digest = header["digest"] + if not isinstance(message_id, str) or len(message_id) > 64: + raise RemoteProtocolError("invalid message ID") + if not all( + isinstance(item, int) and not isinstance(item, bool) + for item in (index, count, total) + ): + raise RemoteProtocolError("chunk indices and lengths must be integers") + if count < 1 or count > self.limits.max_chunks or not 0 <= index < count: + raise RemoteProtocolError("invalid chunk index/count") + if total < 0 or total > self.limits.max_message_bytes: + raise RemoteProtocolError("declared message payload is too large") + if not isinstance(digest, str) or len(digest) != 64: + raise RemoteProtocolError("invalid payload digest") + + state = self._messages.get(message_id) + if state is None: + if index != 0 or not isinstance(header["control"], dict): + raise RemoteProtocolError("chunk sequence must start at index zero") + if len(self._messages) >= self.limits.max_incomplete_messages: + raise RemoteProtocolError("too many incomplete messages") + state = _Chunks(kind, header["control"], count, total, digest, 0, []) + self._messages[message_id] = state + elif ( + state.kind != kind + or state.count != count + or state.total != total + or state.digest != digest + or header["control"] is not None + ): + raise RemoteProtocolError("conflicting chunk metadata") + if index != state.next_index: + raise RemoteProtocolError("duplicate or out-of-order chunk") + if self._incomplete_bytes + len(payload) > self.limits.max_message_bytes: + raise RemoteProtocolError("aggregate incomplete payload is too large") + state.parts.append(payload) + state.received += len(payload) + self._incomplete_bytes += len(payload) + state.next_index += 1 + if state.received > state.total: + raise RemoteProtocolError("chunk sequence exceeds declared total") + if index + 1 != count: + return None + del self._messages[message_id] + self._incomplete_bytes -= state.received + if state.received != state.total: + raise RemoteProtocolError( + "chunk sequence length does not match declared total" + ) + complete = b"".join(state.parts) + if hashlib.sha256(complete).hexdigest() != state.digest: + raise RemoteProtocolError("message payload digest mismatch") + return Message(kind, state.control, complete) + + +def _bounded_timeout(timeout: float, deadline: float | None) -> float: + if deadline is None: + return timeout + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("remote handshake timed out") + return min(timeout, remaining) + + +def _send_exact(sock: socket.socket, data: bytes | memoryview, timeout: float) -> None: + """Write bytes within one monotonic deadline instead of unbounded sendall.""" + view = memoryview(data) + sent = 0 + deadline = time.monotonic() + timeout + while sent < len(view): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("remote frame write timed out") + try: + _, writable, _ = select.select([], [sock], [], remaining) + if not writable: + raise TimeoutError("remote frame write timed out") + count = sock.send(view[sent:]) + except ssl.SSLWantReadError: + readable, _, _ = select.select([sock], [], [], remaining) + if not readable: + raise TimeoutError("remote TLS frame write timed out") + continue + except ssl.SSLWantWriteError: + continue + if count == 0: + raise EOFError("connection closed during frame write") + sent += count + + +def _recv_exact( + sock: socket.socket, + size: int, + *, + timeout: float, + allow_idle: bool = False, +) -> bytes: + data = bytearray(size) + view = memoryview(data) + received = 0 + deadline = None + while received < size: + if not (allow_idle and received == 0): + if deadline is None: + deadline = time.monotonic() + timeout + remaining = deadline - time.monotonic() + if remaining <= 0 or not _socket_readable(sock, remaining): + raise TimeoutError("partial remote frame timed out") + count = sock.recv_into(view[received:]) + if count == 0: + raise EOFError("connection closed during frame") + received += count + if deadline is None: + deadline = time.monotonic() + timeout + return bytes(data) + + +def _socket_readable(sock: socket.socket, timeout: float) -> bool: + pending = getattr(sock, "pending", None) + if pending is not None and pending(): + return True + readable, _, _ = select.select([sock], [], [], timeout) + return bool(readable) diff --git a/deadpool/remote/_scheduler.py b/deadpool/remote/_scheduler.py new file mode 100644 index 0000000..fcc00f7 --- /dev/null +++ b/deadpool/remote/_scheduler.py @@ -0,0 +1,98 @@ +"""Strict-priority, principal-fair broker scheduler.""" + +from __future__ import annotations + +import heapq +from collections import defaultdict, deque +from dataclasses import dataclass, field +from typing import Generic, TypeVar + +T = TypeVar("T") + + +@dataclass(slots=True) +class _Priority(Generic[T]): + principals: deque[str] = field(default_factory=deque) + queued: dict[str, deque[T]] = field(default_factory=lambda: defaultdict(deque)) + + +class FairScheduler(Generic[T]): + """Select lower numeric priorities first and round-robin principals within one.""" + + def __init__(self) -> None: + self._priorities: dict[int, _Priority[T]] = {} + self._heap: list[int] = [] + self._size = 0 + + def put(self, item: T, *, priority: int, principal: str) -> None: + self._put(item, priority=priority, principal=principal, front=False) + + def put_front(self, item: T, *, priority: int, principal: str) -> None: + """Restore a popped item ahead of later work in its FIFO stream.""" + self._put(item, priority=priority, principal=principal, front=True) + + def _put( + self, + item: T, + *, + priority: int, + principal: str, + front: bool, + ) -> None: + bucket = self._priorities.get(priority) + if bucket is None: + bucket = self._priorities[priority] = _Priority() + heapq.heappush(self._heap, priority) + queue = bucket.queued[principal] + if not queue: + bucket.principals.append(principal) + if front: + queue.appendleft(item) + else: + queue.append(item) + self._size += 1 + + def pop(self) -> T: + if not self._size: + raise IndexError("scheduler is empty") + priority = self._heap[0] + bucket = self._priorities[priority] + principal = bucket.principals.popleft() + queue = bucket.queued[principal] + item = queue.popleft() + self._size -= 1 + if queue: + bucket.principals.append(principal) + else: + del bucket.queued[principal] + if not bucket.principals: + heapq.heappop(self._heap) + del self._priorities[priority] + return item + + def remove(self, item: T) -> bool: + for priority, bucket in list(self._priorities.items()): + for principal, queue in list(bucket.queued.items()): + try: + queue.remove(item) + except ValueError: + continue + self._size -= 1 + if not queue: + del bucket.queued[principal] + try: + bucket.principals.remove(principal) + except ValueError: + pass + if not bucket.principals: + del self._priorities[priority] + self._heap.remove(priority) + heapq.heapify(self._heap) + return True + return False + + def __bool__(self) -> bool: + return bool(self._size) + + def __len__(self) -> int: + return self._size diff --git a/deadpool/remote/_transport.py b/deadpool/remote/_transport.py new file mode 100644 index 0000000..f80140f --- /dev/null +++ b/deadpool/remote/_transport.py @@ -0,0 +1,226 @@ +"""Unix and TCP socket lifecycle shared by client and server.""" + +from __future__ import annotations + +import errno +import os +import socket +import stat +import struct +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path + +from .config import PeerInfo, TcpAddress, TcpListener, UnixAddress, UnixListener + + +@dataclass(slots=True) +class BoundListener: + config: UnixListener | TcpListener + socket: socket.socket + address: object + unix_identity: tuple[int, int] | None = None + + def close(self) -> None: + try: + self.socket.close() + finally: + if isinstance(self.config, UnixListener) and self.unix_identity is not None: + path = self.config.path + try: + current = path.lstat() + except FileNotFoundError: + current = None + if current is not None: + _unlink_if_identity(path, current, self.unix_identity) + + +def bind_listener(config: UnixListener | TcpListener) -> BoundListener: + if isinstance(config, UnixListener): + return _bind_unix(config) + return _bind_tcp(config) + + +def connect_address(address: UnixAddress | TcpAddress) -> socket.socket: + if isinstance(address, UnixAddress): + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.settimeout(address.connect_timeout) + try: + sock.connect(str(address.path)) + except BaseException: + sock.close() + raise + return sock + sock = socket.create_connection( + (address.host, address.port), timeout=address.connect_timeout + ) + if address.ssl_context is not None: + hostname = address.server_hostname or address.host + try: + sock = address.ssl_context.wrap_socket(sock, server_hostname=hostname) + except BaseException: + sock.close() + raise + return sock + + +def peer_info(sock: socket.socket, transport: str) -> PeerInfo: + uid = gid = pid = None + if transport == "unix" and hasattr(socket, "SO_PEERCRED"): + raw = sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12) + pid, uid, gid = struct.unpack("3i", raw) + certificate = None + if hasattr(sock, "getpeercert"): + try: + certificate = sock.getpeercert() or None + except (ValueError, OSError): + pass + try: + address = sock.getpeername() + except OSError: + address = None + return PeerInfo(transport, address, uid, gid, pid, certificate) + + +def _bind_unix(config: UnixListener) -> BoundListener: + path = config.path + path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + _validate_unix_directory(path.parent) + if path.exists() or path.is_symlink(): + info = path.lstat() + if stat.S_ISLNK(info.st_mode) or not stat.S_ISSOCK(info.st_mode): + raise ValueError(f"refusing unexpected Unix socket path type: {path}") + if config.stale_policy != "force_unlink": + raise FileExistsError(path) + if _unix_socket_is_live(path): + raise OSError(f"Unix socket is already live: {path}") + + # A socket may replace the stale path after its first failed probe. + # Confirm it remains unreachable before removing the pathname. + if _unix_socket_is_live(path): + raise _address_in_use(path) + + try: + current = path.lstat() + except FileNotFoundError: + pass + else: + if not _unlink_if_identity(path, current, (info.st_dev, info.st_ino)): + raise _address_in_use(path) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + identity: tuple[int, int] | None = None + try: + with _umask(0o077): + sock.bind(str(path)) + bound = path.lstat() + identity = (bound.st_dev, bound.st_ino) + if config.owner_uid is not None or config.owner_gid is not None: + os.chown( + path, + config.owner_uid if config.owner_uid is not None else -1, + config.owner_gid if config.owner_gid is not None else -1, + ) + os.chmod(path, config.mode) + sock.listen(config.backlog) + return BoundListener(config, sock, path, identity) + except BaseException: + sock.close() + if identity is not None: + try: + current = path.lstat() + except FileNotFoundError: + pass + else: + _unlink_if_identity(path, current, identity) + raise + + +def _unix_socket_is_live(path: Path) -> bool: + """Return whether a Unix socket pathname currently accepts connections.""" + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + probe.settimeout(0.2) + probe.connect(str(path)) + except (ConnectionRefusedError, FileNotFoundError, socket.timeout): + return False + finally: + probe.close() + return True + + +def _address_in_use(path: Path) -> OSError: + return OSError(errno.EADDRINUSE, os.strerror(errno.EADDRINUSE), str(path)) + + +def _unlink_if_identity( + path: Path, current: os.stat_result, identity: tuple[int, int] +) -> bool: + """Remove a Unix listener path only when it retains ``identity``. + + Returns whether the path was removed. This protects replacements visible + to the final ``lstat``; a portable ``lstat``/``unlink`` still has an + unavoidable TOCTOU window. + """ + if (current.st_dev, current.st_ino) != identity: + return False + path.unlink() + return True + + +def _validate_unix_directory(path) -> None: + info = path.lstat() + if not stat.S_ISDIR(info.st_mode) or stat.S_ISLNK(info.st_mode): + raise ValueError(f"Unix socket directory must be a real directory: {path}") + if info.st_mode & 0o022: + raise PermissionError( + f"Unix socket directory must not be group/world writable: {path}" + ) + if info.st_uid not in {0, os.getuid()}: + raise PermissionError(f"Unix socket directory has an untrusted owner: {path}") + + +def _bind_tcp(config: TcpListener) -> BoundListener: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if config.keepalive: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + sock.bind((config.host, config.port)) + sock.listen(config.backlog) + return BoundListener(config, sock, sock.getsockname()) + except BaseException: + sock.close() + raise + + +def accept_socket(bound: BoundListener) -> socket.socket: + """Accept only the raw socket; potentially slow TLS runs in a client thread.""" + sock, _ = bound.socket.accept() + return sock + + +def prepare_accepted_socket( + bound: BoundListener, + sock: socket.socket, + *, + handshake_timeout: float, +) -> socket.socket: + """Complete transport authentication under a finite deadline.""" + config = bound.config + sock.settimeout(handshake_timeout) + if isinstance(config, TcpListener) and config.ssl_context is not None: + try: + sock = config.ssl_context.wrap_socket(sock, server_side=True) + except BaseException: + sock.close() + raise + return sock + + +@contextmanager +def _umask(mask: int): + previous = os.umask(mask) + try: + yield + finally: + os.umask(previous) diff --git a/deadpool/remote/_worker.py b/deadpool/remote/_worker.py new file mode 100644 index 0000000..34838ac --- /dev/null +++ b/deadpool/remote/_worker.py @@ -0,0 +1,75 @@ +"""Opaque bytes-in/bytes-out worker entry point.""" + +from __future__ import annotations + +import traceback +from dataclasses import dataclass +from typing import Callable + +from .serializer import SerializationLimitError, Serializer + + +@dataclass(frozen=True, slots=True) +class WorkerOutcome: + kind: str + payload: bytes = b"" + descriptor: dict | None = None + + +def execute_opaque( + serializer: Serializer, + mode: str, + invocation: bytes, + operation: str | None, + registry: dict[str, Callable] | None, + max_result_bytes: int, +) -> WorkerOutcome: + """Materialize application objects only inside the selected worker.""" + try: + if mode == "callable": + fn, args, kwargs = serializer.loads(invocation) + elif mode == "registered": + if registry is None or operation not in registry: + raise LookupError(f"unknown registered operation {operation!r}") + fn = registry[operation] + args, kwargs = serializer.loads(invocation) + else: + raise ValueError(f"unknown invocation mode {mode!r}") + result = fn(*args, **kwargs) + except BaseException as error: + descriptor = _describe_exception(error) + try: + payload = serializer.dumps(error, limit=max_result_bytes) + except BaseException as serialization_error: + descriptor["serialization_error"] = _safe_repr(serialization_error) + payload = b"" + return WorkerOutcome("task_error", payload, descriptor) + + try: + payload = serializer.dumps(result, limit=max_result_bytes) + except SerializationLimitError as error: + return WorkerOutcome( + "result_too_large", descriptor={"message": _safe_repr(error)} + ) + except BaseException as error: + return WorkerOutcome( + "result_encoding_failed", + descriptor=_describe_exception(error), + ) + return WorkerOutcome("result", payload) + + +def _describe_exception(error: BaseException) -> dict: + return { + "module": type(error).__module__, + "type": type(error).__qualname__, + "message": _safe_repr(error), + "traceback": "".join(traceback.format_exception(error)), + } + + +def _safe_repr(value: object) -> str: + try: + return repr(value)[:4096] + except BaseException: + return f"<{type(value).__module__}.{type(value).__qualname__}>" diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py new file mode 100644 index 0000000..0fb31b6 --- /dev/null +++ b/deadpool/remote/client.py @@ -0,0 +1,921 @@ +"""Synchronous Executor-compatible client for a remote Deadpool server.""" + +from __future__ import annotations + +import concurrent.futures +import hashlib +import logging +import math +import os +import platform +import queue +import socket +import sys +import threading +import time +import uuid +from typing import Callable, Iterable + +from deadpool import TimeoutError as DeadpoolTimeoutError + +from ._future import RemoteFuture, SubmissionState +from ._protocol import ( + MAJOR, + MINOR, + Message, + MessageReader, + MessageType, + _validate_wire_limits, + _wire_limits, + send_message, + validate_control, +) +from ._transport import connect_address +from .config import RemoteLimits, TcpAddress, UnixAddress +from .errors import ( + AcceptanceCertainty, + ExecutionCertainty, + RemoteAuthenticationError, + RemoteCancellationOutcomeUnknown, + RemoteCompatibilityError, + RemoteConnectionLost, + RemoteExecutorError, + RemoteExecutorUnavailable, + RemoteProcessError, + RemoteProtocolError, + RemoteQueueFull, + RemoteQueueTimeout, + RemoteResultEncodingError, + RemoteResultTooLarge, + RemoteSubmissionTimeout, + RemoteTaskError, + SubmissionOutcomeUnknown, +) +from .serializer import PICKLE_TRUST_WARNING, PickleSerializer, Serializer + +logger = logging.getLogger("deadpool.remote") + + +class DeadpoolClient(concurrent.futures.Executor): + """Submit work to a server-owned pool across a mutually trusted boundary. + + The default :class:`PickleSerializer` is unsafe for untrusted peers in both + callable and registered-task modes. Results and exceptions are also pickle + payloads, so a client must trust the server just as the server trusts it. + """ + + def __init__( + self, + address: UnixAddress | TcpAddress, + *, + serializer: Serializer | None = None, + authenticator: Callable[[], dict] | None = None, + application_fingerprint: str | None = None, + registry_fingerprint: str | None = None, + limits: RemoteLimits | None = None, + submission_timeout: float = 5.0, + control_timeout: float = 5.0, + ) -> None: + self.address = address + if serializer is None: + logger.warning(PICKLE_TRUST_WARNING) + self.serializer = serializer or PickleSerializer() + self.authenticator = authenticator + self.application_fingerprint = application_fingerprint + self.registry_fingerprint = registry_fingerprint + self.limits = limits or RemoteLimits() + self.submission_timeout = _positive_timeout( + submission_timeout, "submission_timeout" + ) + self.control_timeout = _positive_timeout(control_timeout, "control_timeout") + self._owner_pid = os.getpid() + self._client_id = uuid.uuid4().hex + self._sequence = 0 + self._lock = threading.RLock() + self._socket: socket.socket | None = None + self._reader: MessageReader | None = None + self._outbound: queue.Queue[tuple[Message | None, RemoteFuture | None]] = ( + queue.Queue(self.limits.outbound_queue_size) + ) + self._futures: dict[str, RemoteFuture] = {} + self._terminal_received: set[str] = set() + self._rpc: dict[str, tuple[threading.Event, dict]] = {} + self._closed = False + self._transport_failed = False + # One ordered state lane preserves wire order. User callbacks have a + # separate pool and cannot block transport/state progression. + self._completion_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="deadpool.remote.completion", + ) + self._serialization_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.limits.completion_workers, + thread_name_prefix="deadpool.remote.serialization", + ) + self._callback_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.limits.completion_workers, + thread_name_prefix="deadpool.remote.callback", + ) + self._completion_slots = threading.BoundedSemaphore( + self.limits.inbound_queue_size + ) + self._callback_slots = threading.BoundedSemaphore( + self.limits.callback_queue_size + ) + self._serialization_slots = threading.BoundedSemaphore( + self.limits.outbound_queue_size + ) + self._callback_dispatcher: threading.Thread | None = None + self.server_id: str | None = None + self.server_epoch: str | None = None + self.session_id: str | None = None + self.negotiated_limits: dict = {} + self._connect() + self._start_callback_dispatcher() + + def _connect(self) -> None: + try: + sock = connect_address(self.address) + sock.settimeout(self.limits.handshake_timeout) + hello = { + "versions": [[MAJOR, MINOR]], + "features": ["chunking"], + "wire": "experimental-deadpool-private-v1", + "client_instance_id": self._client_id, + "serializer": self.serializer.name, + "serializer_protocol": self.serializer.protocol_version, + "python": [sys.version_info.major, sys.version_info.minor], + "implementation": platform.python_implementation(), + "deadpool_version": _deadpool_version(), + "capabilities": ["callable", "registered"], + "application_fingerprint": self.application_fingerprint, + "registry_fingerprint": self.registry_fingerprint, + "authentication": self.authenticator() if self.authenticator else None, + "max_result_bytes": self.limits.max_result_bytes, + "wire_limits": _wire_limits(self.limits), + } + send_message(sock, Message(MessageType.HELLO, hello), self.limits) + reader = MessageReader(self.limits) + response = reader.receive( + sock, + deadline=time.monotonic() + self.limits.handshake_timeout, + ) + if response.kind == MessageType.HANDSHAKE_REJECTED: + reason = response.control.get("reason") + error_type = ( + RemoteAuthenticationError + if reason == "authentication" + else RemoteCompatibilityError + ) + raise error_type(f"remote handshake rejected: {reason}") + if response.kind != MessageType.WELCOME: + raise RemoteProtocolError("server did not send WELCOME") + if response.control.get("wire") != "experimental-deadpool-private-v1": + raise RemoteCompatibilityError( + "server selected an incompatible wire protocol" + ) + server_wire_limits = _validate_wire_limits( + response.control.get("wire_limits") + ) + if server_wire_limits != _wire_limits(self.limits): + raise RemoteCompatibilityError( + "server selected incompatible wire limits" + ) + sock.settimeout(None) + except RemoteExecutorError: + try: + sock.close() + except (UnboundLocalError, OSError): + pass + raise + except BaseException as error: + try: + sock.close() + except (UnboundLocalError, OSError): + pass + raise RemoteExecutorUnavailable(str(error)) from error + with self._lock: + self._socket = sock + self._reader = reader + self.server_id = response.control.get("server_id") + self.server_epoch = response.control.get("epoch") + self.session_id = response.control.get("session_id") + self.negotiated_limits = { + **dict(response.control.get("limits") or {}), + **server_wire_limits, + } + self._transport_failed = False + threading.Thread( + target=self._sender_loop, name="deadpool.remote.sender", daemon=True + ).start() + threading.Thread( + target=self._receiver_loop, name="deadpool.remote.receiver", daemon=True + ).start() + + def submit(self, fn: Callable, /, *args, **kwargs) -> RemoteFuture: + return self._submit("callable", fn, args, kwargs) + + def submit_task(self, operation: str, /, *args, **kwargs) -> RemoteFuture: + if not isinstance(operation, str) or not operation: + raise ValueError("operation must be a non-empty string") + return self._submit("registered", operation, args, kwargs) + + def _submit( + self, mode: str, target: object, args: tuple, kwargs: dict + ) -> RemoteFuture: + self._ensure_process() + with self._lock: + if self._closed: + raise RuntimeError("cannot schedule new futures after shutdown") + options = _submission_options(kwargs, self.submission_timeout) + priority = options["priority"] + if not isinstance(priority, int) or isinstance(priority, bool) or priority < 0: + raise ValueError("deadpool_priority must be a non-negative integer") + deadline = time.monotonic() + options["submission_timeout"] + invocation_value = ( + (target, args, kwargs) if mode == "callable" else (args, kwargs) + ) + limit = min( + self.limits.max_invocation_bytes, + int( + self.negotiated_limits.get( + "max_invocation_bytes", self.limits.max_invocation_bytes + ) + ), + ) + payload = self._serialize_invocation(invocation_value, limit, deadline) + with self._lock: + # Serialization deliberately runs outside the state lock. Recheck + # both shutdown and transport state when linearizing registration, + # because either may have changed while serialization was running. + if self._closed: + raise RuntimeError("cannot schedule new futures after shutdown") + if self._transport_failed or self._socket is None: + raise RemoteExecutorUnavailable("connection is unavailable") + self._sequence += 1 + request_id = f"{self._client_id}:{self._sequence}" + future = RemoteFuture(request_id, self) + self._futures[request_id] = future + control = { + "request_id": request_id, + "digest": hashlib.sha256(payload).hexdigest(), + "mode": mode, + "operation": target if mode == "registered" else None, + "priority": priority, + "execution_timeout": options["execution_timeout"], + "queue_timeout": options["queue_timeout"], + "group_id": options["group_id"], + "metadata": options["metadata"], + } + try: + validate_control(control, self.limits) + except BaseException: + with self._lock: + self._futures.pop(request_id, None) + raise + remaining = deadline - time.monotonic() + if remaining <= 0: + with self._lock: + self._futures.pop(request_id, None) + raise RemoteSubmissionTimeout( + "local submission work exceeded its timeout", request_id=request_id + ) + try: + self._outbound.put( + (Message(MessageType.SUBMIT, control, payload), future), + timeout=remaining, + ) + except queue.Full as error: + with self._lock: + self._futures.pop(request_id, None) + raise RemoteSubmissionTimeout( + "client outbound queue is full", request_id=request_id + ) from error + return future + + def _serialize_invocation( + self, + invocation: object, + limit: int, + deadline: float, + ) -> bytes: + remaining = deadline - time.monotonic() + if remaining <= 0 or not self._serialization_slots.acquire(timeout=remaining): + raise RemoteSubmissionTimeout("serializer capacity wait timed out") + try: + work = self._serialization_executor.submit( + self.serializer.dumps, + invocation, + limit=limit, + ) + except BaseException: + self._serialization_slots.release() + raise + work.add_done_callback(lambda _: self._serialization_slots.release()) + remaining = deadline - time.monotonic() + try: + return work.result(timeout=max(0.0, remaining)) + except concurrent.futures.TimeoutError as error: + work.cancel() + raise RemoteSubmissionTimeout( + "invocation serialization exceeded submission timeout" + ) from error + + def submit_many( + self, submissions: Iterable[tuple[Callable, tuple, dict]] + ) -> list[RemoteFuture]: + return [self.submit(fn, *args, **kwargs) for fn, args, kwargs in submissions] + + def _sender_loop(self) -> None: + while True: + message, future = self._outbound.get() + try: + if message is None: + return + if future is not None and not future._set_sent(): + self._forget(future) + continue + with self._lock: + sock = self._socket + if sock is None: + raise OSError("connection is closed") + send_message(sock, message, self.limits) + except BaseException as error: + self._connection_lost(error) + return + finally: + self._outbound.task_done() + + def _receiver_loop(self) -> None: + try: + while True: + with self._lock: + sock, reader = self._socket, self._reader + if sock is None or reader is None: + return + message = reader.receive(sock) + self._receive(message) + except BaseException as error: + self._connection_lost(error) + + def _receive(self, message: Message) -> None: + request_id = message.control.get("request_id") + with self._lock: + future = ( + self._futures.get(request_id) if isinstance(request_id, str) else None + ) + if message.kind == MessageType.ACCEPTED and future is not None: + self._schedule_completion(future._set_accepted) + elif message.kind == MessageType.RUNNING and future is not None: + self._schedule_completion( + future._set_running, + pid=message.control.get("pid"), + worker_id=message.control.get("worker_id"), + ) + elif ( + message.kind + in { + MessageType.RESULT, + MessageType.TASK_ERROR, + MessageType.TIMED_OUT, + MessageType.CANCELLED, + MessageType.WORKER_LOST, + MessageType.RESULT_ENCODING_FAILED, + MessageType.RESULT_TOO_LARGE, + MessageType.QUEUE_TIMED_OUT, + } + and future is not None + ): + with self._lock: + self._terminal_received.add(future.request_id) + self._schedule_completion(self._complete, future, message) + elif message.kind == MessageType.REJECTED and future is not None: + self._schedule_completion(self._reject, future, message.control) + elif message.kind in { + MessageType.CANCEL_RESPONSE, + MessageType.CANCEL_GROUP_RESPONSE, + MessageType.STATUS_RESPONSE, + MessageType.STATS_RESPONSE, + MessageType.PONG, + MessageType.CLOSE_SESSION_RESPONSE, + }: + token = message.control.get("token") or message.control.get("nonce") + if isinstance(token, str): + with self._lock: + rpc = self._rpc.pop(token, None) + if rpc is not None: + rpc[1].update(message.control) + if rpc is not None: + rpc[0].set() + elif message.kind == MessageType.GOAWAY: + self._connection_lost(RemoteConnectionLost("server closed the session")) + elif message.kind in { + MessageType.PROTOCOL_ERROR, + MessageType.HANDSHAKE_REJECTED, + }: + self._connection_lost(RemoteProtocolError(str(message.control))) + + def _complete(self, future: RemoteFuture, message: Message) -> None: + try: + if message.kind == MessageType.RESULT: + future._set_result(self.serializer.loads(message.payload)) + elif message.kind == MessageType.TASK_ERROR: + error = None + if message.payload: + try: + decoded = self.serializer.loads(message.payload) + if isinstance(decoded, BaseException): + error = decoded + except BaseException: + pass + if error is None: + error = RemoteTaskError( + str(message.control.get("message", "remote task failed")), + request_id=future.request_id, + remote_traceback=str(message.control.get("traceback", "")), + ) + future._set_exception(error, task=True) + elif message.kind == MessageType.TIMED_OUT: + future._set_exception( + DeadpoolTimeoutError( + message.control.get("message", "remote task timed out") + ) + ) + elif message.kind == MessageType.CANCELLED: + future._set_cancelled() + elif message.kind == MessageType.QUEUE_TIMED_OUT: + future._set_exception( + RemoteQueueTimeout( + "remote queue wait expired", request_id=future.request_id + ) + ) + elif message.kind == MessageType.WORKER_LOST: + future._set_exception( + RemoteProcessError( + str(message.control.get("message", "worker lost")), + request_id=future.request_id, + ) + ) + elif message.kind == MessageType.RESULT_ENCODING_FAILED: + future._set_exception( + RemoteResultEncodingError( + str(message.control.get("message", "result encoding failed")), + request_id=future.request_id, + ) + ) + elif message.kind == MessageType.RESULT_TOO_LARGE: + future._set_exception( + RemoteResultTooLarge( + str(message.control.get("message", "result too large")), + request_id=future.request_id, + ) + ) + except BaseException as error: + future._set_exception( + error + if isinstance(error, RemoteExecutorError) + else RemoteProtocolError(str(error), request_id=future.request_id) + ) + finally: + try: + self._enqueue_control( + MessageType.RESULT_ACK, {"request_id": future.request_id} + ) + except RemoteConnectionLost as error: + # An unqueued acknowledgement cannot release the server-owned + # reservation; closing the session makes teardown release it. + self._connection_lost(error) + finally: + self._forget(future) + with self._lock: + self._terminal_received.discard(future.request_id) + + def _reject(self, future: RemoteFuture, control: dict) -> None: + reason = control.get("reason", "rejected") + if reason == "queue_full": + error = RemoteQueueFull( + "server queue is full", request_id=future.request_id + ) + elif reason == "queue_timeout": + error = RemoteQueueTimeout( + "server queue wait expired", request_id=future.request_id + ) + else: + error = RemoteExecutorError( + f"remote submission rejected: {reason}", + request_id=future.request_id, + acceptance_certainty=AcceptanceCertainty.NOT_ACCEPTED, + execution_certainty=ExecutionCertainty.NOT_STARTED, + ) + future._set_exception(error) + self._forget(future) + + def _cancel(self, future: RemoteFuture, *, hard: bool) -> bool: + token = uuid.uuid4().hex + event, response = threading.Event(), {} + with self._lock: + if self._socket is None: + raise RemoteCancellationOutcomeUnknown( + "connection is unavailable", request_id=future.request_id + ) + self._rpc[token] = (event, response) + try: + self._enqueue_control( + MessageType.CANCEL_REQUEST, + {"request_id": future.request_id, "hard": hard, "token": token}, + ) + if not event.wait(self.control_timeout): + raise RemoteCancellationOutcomeUnknown( + "timed out waiting for cancellation decision", + request_id=future.request_id, + ) + decision = response.get("response") + if decision == "cancelled": + # The response is authoritative enough for cancel()'s return + # value, but retain the request until the terminal CANCELLED + # frame is received and acknowledged. + future._set_cancelled() + return True + if decision in {"running", "terminal"}: + return False + raise RemoteCancellationOutcomeUnknown( + f"cancellation decision is {decision!r}", request_id=future.request_id + ) + finally: + with self._lock: + self._rpc.pop(token, None) + + def cancel_group( + self, + group_id: str, + *, + hard: bool = False, + timeout: float | None = None, + ) -> dict[str, int]: + if not isinstance(group_id, str) or not group_id: + raise ValueError("group_id must be a non-empty string") + response = self._rpc_request( + MessageType.CANCEL_GROUP, + {"group_id": group_id, "hard": hard}, + timeout, + ) + return { + name: int(response.get(name, 0)) + for name in ("cancelled", "running", "terminal", "unknown") + } + + def check_health(self, timeout: float | None = None) -> bool: + response = self._rpc_request(MessageType.PING, {}, timeout) + return "nonce" in response + + def get_statistics(self, timeout: float | None = None) -> dict: + response = self._rpc_request(MessageType.STATS_REQUEST, {}, timeout) + if "error" in response: + raise RemoteExecutorError(response["error"]) + return dict(response.get("statistics") or {}) + + def get_status( + self, future_or_request_id: RemoteFuture | str, timeout: float | None = None + ) -> str: + request_id = ( + future_or_request_id.request_id + if isinstance(future_or_request_id, RemoteFuture) + else future_or_request_id + ) + response = self._rpc_request( + MessageType.STATUS_REQUEST, {"request_id": request_id}, timeout + ) + return str(response.get("state", "UNKNOWN")) + + def _rpc_request( + self, kind: MessageType, control: dict, timeout: float | None + ) -> dict: + self._ensure_process() + token = uuid.uuid4().hex + event, response = threading.Event(), {} + with self._lock: + if self._socket is None: + raise RemoteConnectionLost("connection is unavailable") + self._rpc[token] = (event, response) + key = "nonce" if kind == MessageType.PING else "token" + try: + self._enqueue_control(kind, {**control, key: token}) + if not event.wait(timeout or self.control_timeout): + raise TimeoutError("remote control request timed out") + transport_error = response.get("_transport_error") + if transport_error is not None: + raise RemoteConnectionLost(str(transport_error)) + return response + finally: + with self._lock: + self._rpc.pop(token, None) + + def _enqueue_control(self, kind: MessageType, control: dict) -> None: + try: + self._outbound.put( + (Message(kind, control), None), timeout=self.control_timeout + ) + except queue.Full as error: + raise RemoteConnectionLost( + "client outbound control queue is full" + ) from error + + def _connection_lost(self, error: BaseException) -> None: + with self._lock: + if self._transport_failed: + return + self._transport_failed = True + sock, self._socket = self._socket, None + futures = list(self._futures.values()) + rpc = list(self._rpc.values()) + self._rpc.clear() + if sock is not None: + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + sock.close() + except OSError: + pass + for event, response in rpc: + response["_transport_error"] = error + event.set() + for future in futures: + with self._lock: + terminal_received = future.request_id in self._terminal_received + if future.done() or terminal_received: + continue + state = future.submission_state + if state == SubmissionState.LOCAL_PENDING: + exception = RemoteExecutorUnavailable( + str(error), request_id=future.request_id + ) + elif state == SubmissionState.SENT_UNACKNOWLEDGED: + exception = SubmissionOutcomeUnknown( + str(error), request_id=future.request_id + ) + else: + exception = RemoteConnectionLost( + str(error), + request_id=future.request_id, + acceptance_certainty=AcceptanceCertainty.ACCEPTED, + execution_certainty=ExecutionCertainty.MAY_HAVE_RUN, + ) + self._schedule_completion(self._fail_and_forget, future, exception) + + def _fail_and_forget( + self, + future: RemoteFuture, + exception: BaseException, + ) -> None: + future._set_exception(exception) + self._forget(future) + + def _schedule_completion(self, function, *args, **kwargs) -> None: + self._completion_slots.acquire() + try: + self._completion_executor.submit( + _run_bounded, + self._completion_slots, + function, + args, + kwargs, + ) + except RuntimeError: + self._completion_slots.release() + function(*args, **kwargs) + + def _forget(self, future: RemoteFuture) -> None: + with self._lock: + self._futures.pop(future.request_id, None) + + def _start_callback_dispatcher(self) -> None: + self._callback_dispatch_queue = queue.SimpleQueue() + self._callback_dispatch_closed = False + self._callback_dispatcher = threading.Thread( + target=self._callback_dispatcher_loop, + name="deadpool.remote.callback-dispatcher", + daemon=True, + ) + self._callback_dispatcher.start() + + def _callback_dispatcher_loop(self) -> None: + while True: + item = self._callback_dispatch_queue.get() + if item is None: + return + callbacks, future = item + # Registration stays stdlib-compatible and nonblocking. Capacity + # is acquired here, outside transport and completion-state lanes. + self._callback_slots.acquire() + try: + self._callback_executor.submit( + _run_callback_batch, + self._callback_slots, + callbacks, + future, + ) + except RuntimeError: + self._callback_slots.release() + for callback in callbacks: + future._invoke_callback(callback) + + def _schedule_callbacks(self, callbacks: list, future: RemoteFuture) -> None: + with self._lock: + if not self._callback_dispatch_closed: + self._callback_dispatch_queue.put((callbacks, future)) + return + for callback in callbacks: + future._invoke_callback(callback) + + def _ensure_process(self) -> None: + if os.getpid() == self._owner_pid: + return + # Threads and locks do not survive fork coherently. Close only the + # child-side descriptor and build a wholly new process-local session. + try: + if self._socket is not None: + self._socket.close() + except OSError: + pass + self._owner_pid = os.getpid() + self._client_id = uuid.uuid4().hex + self._sequence = 0 + self._lock = threading.RLock() + self._socket = None + self._reader = None + self._outbound = queue.Queue(self.limits.outbound_queue_size) + self._futures = {} + self._terminal_received = set() + self._rpc = {} + self._closed = False + self._transport_failed = False + self._completion_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="deadpool.remote.completion", + ) + self._serialization_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.limits.completion_workers, + thread_name_prefix="deadpool.remote.serialization", + ) + self._callback_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.limits.completion_workers, + thread_name_prefix="deadpool.remote.callback", + ) + self._completion_slots = threading.BoundedSemaphore( + self.limits.inbound_queue_size + ) + self._callback_slots = threading.BoundedSemaphore( + self.limits.callback_queue_size + ) + self._serialization_slots = threading.BoundedSemaphore( + self.limits.outbound_queue_size + ) + self._connect() + self._start_callback_dispatcher() + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + self._ensure_process() + with self._lock: + if self._closed: + return + self._closed = True + futures = list(self._futures.values()) + if cancel_futures: + for future in futures: + if not future.done(): + try: + future.cancel() + except RemoteExecutorError: + pass + lifecycle = threading.Thread( + target=self._finish_shutdown, + args=(futures,), + name="deadpool.remote.client-shutdown", + daemon=False, + ) + lifecycle.start() + if wait: + lifecycle.join() + + def _finish_shutdown(self, futures: list[RemoteFuture]) -> None: + for future in futures: + try: + future.result() + except BaseException: + pass + with self._lock: + can_close_cleanly = self._socket is not None and not self._transport_failed + if can_close_cleanly: + try: + self._rpc_request( + MessageType.CLOSE_SESSION, + {}, + self.control_timeout, + ) + except BaseException: + pass + try: + self._outbound.put((None, None), timeout=self.control_timeout) + except queue.Full: + pass + with self._lock: + sock, self._socket = self._socket, None + if sock is not None: + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + _drain_semaphore( + self._completion_slots, + self.limits.inbound_queue_size, + ) + with self._lock: + self._callback_dispatch_closed = True + self._callback_dispatch_queue.put(None) + self._completion_executor.shutdown(wait=False) + self._serialization_executor.shutdown(wait=False) + self._callback_executor.shutdown(wait=False) + + def __enter__(self) -> "DeadpoolClient": + return self + + def __exit__(self, exc_type, exc, traceback) -> bool: + self.shutdown() + return False + + +def _submission_options(kwargs: dict, default_submission_timeout: float) -> dict: + values = { + "priority": kwargs.pop("deadpool_priority", 0), + "execution_timeout": kwargs.pop("deadpool_timeout", None), + "queue_timeout": kwargs.pop("deadpool_queue_timeout", None), + "submission_timeout": kwargs.pop( + "deadpool_submission_timeout", default_submission_timeout + ), + "group_id": kwargs.pop("deadpool_group_id", None), + "metadata": kwargs.pop("deadpool_metadata", {}), + } + for name in ("execution_timeout", "queue_timeout"): + value = values[name] + if value is not None and ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value < 0 + ): + raise ValueError(f"{name} must be finite and non-negative") + values["submission_timeout"] = _positive_timeout( + values["submission_timeout"], "submission_timeout" + ) + if values["group_id"] is not None: + if not isinstance(values["group_id"], str): + raise TypeError("deadpool_group_id must be a string or None") + if not values["group_id"]: + raise ValueError("deadpool_group_id must be non-empty") + if not isinstance(values["metadata"], dict): + raise TypeError("deadpool_metadata must be a dict") + return values + + +def _positive_timeout(value: object, name: str) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be finite and greater than zero") + return float(value) + + +def _drain_semaphore(semaphore: threading.BoundedSemaphore, size: int) -> None: + for _ in range(size): + semaphore.acquire() + for _ in range(size): + semaphore.release() + + +def _run_bounded(semaphore, function, args: tuple, kwargs: dict) -> None: + try: + function(*args, **kwargs) + finally: + semaphore.release() + + +def _run_callback_batch( + semaphore: threading.BoundedSemaphore, + callbacks: list, + future: RemoteFuture, +) -> None: + semaphore.release() + for callback in callbacks: + future._invoke_callback(callback) + + +def _deadpool_version() -> str: + from deadpool import __version__ + + return __version__ diff --git a/deadpool/remote/config.py b/deadpool/remote/config.py new file mode 100644 index 0000000..5a72eba --- /dev/null +++ b/deadpool/remote/config.py @@ -0,0 +1,227 @@ +"""Validated addressing and resource policy for the remote executor.""" + +from __future__ import annotations + +import math +import ssl +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Protocol, runtime_checkable + +_INSECURE_LOOPBACK_HOSTS = frozenset({"127.0.0.1"}) + + +@dataclass(frozen=True, slots=True) +class UnixAddress: + path: str | Path + connect_timeout: float = 5.0 + + def __post_init__(self) -> None: + object.__setattr__(self, "path", Path(self.path)) + _positive(self.connect_timeout, "connect_timeout") + + +@dataclass(frozen=True, slots=True) +class TcpAddress: + host: str + port: int + ssl_context: ssl.SSLContext | None = None + server_hostname: str | None = None + insecure: bool = False + connect_timeout: float = 5.0 + + def __post_init__(self) -> None: + if isinstance(self.port, bool) or not isinstance(self.port, int): + raise TypeError("port must be an integer") + if not 0 <= self.port <= 65535: + raise ValueError("port must be between 0 and 65535") + _positive(self.connect_timeout, "connect_timeout") + if self.ssl_context is None and not self.insecure: + raise ValueError("TCP requires ssl_context or explicit insecure=True") + if self.ssl_context is not None: + _validate_client_tls(self.ssl_context) + if self.insecure and self.host not in _INSECURE_LOOPBACK_HOSTS: + raise ValueError("insecure TCP addresses are restricted to loopback") + + +@dataclass(frozen=True, slots=True) +class UnixListener: + """Unix listener; ``force_unlink`` is explicit unsafe stale cleanup.""" + + path: str | Path + mode: int = 0o600 + owner_uid: int | None = None + owner_gid: int | None = None + backlog: int = 128 + stale_policy: str = "error" + optional: bool = False + + def __post_init__(self) -> None: + object.__setattr__(self, "path", Path(self.path)) + if self.stale_policy not in {"error", "force_unlink"}: + raise ValueError("stale_policy must be 'error' or 'force_unlink'") + if isinstance(self.mode, bool) or not isinstance(self.mode, int): + raise TypeError("mode must be an integer") + if not 0 <= self.mode <= 0o777: + raise ValueError("mode must be a permission mode") + _positive_int(self.backlog, "backlog") + for name, value in ( + ("owner_uid", self.owner_uid), + ("owner_gid", self.owner_gid), + ): + if value is not None and ( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + ): + raise ValueError(f"{name} must be a non-negative integer or None") + + +@dataclass(frozen=True, slots=True) +class TcpListener: + host: str + port: int + ssl_context: ssl.SSLContext | None = None + insecure: bool = False + backlog: int = 128 + keepalive: bool = True + optional: bool = False + + def __post_init__(self) -> None: + if isinstance(self.port, bool) or not isinstance(self.port, int): + raise TypeError("port must be an integer") + if not 0 <= self.port <= 65535: + raise ValueError("port must be between 0 and 65535") + _positive_int(self.backlog, "backlog") + if self.ssl_context is None and not self.insecure: + raise ValueError("TCP requires ssl_context or explicit insecure=True") + if self.ssl_context is not None: + _validate_server_tls(self.ssl_context) + if self.insecure and self.host not in _INSECURE_LOOPBACK_HOSTS: + raise ValueError("insecure TCP listeners are restricted to loopback") + + +@dataclass(frozen=True, slots=True) +class RemoteLimits: + """Finite defaults bound protocol/resource queues and byte buffers.""" + + max_control_bytes: int = 64 * 1024 + max_frame_payload_bytes: int = 1024 * 1024 + max_message_bytes: int = 64 * 1024 * 1024 + max_invocation_bytes: int = 32 * 1024 * 1024 + max_result_bytes: int = 32 * 1024 * 1024 + max_metadata_bytes: int = 16 * 1024 + max_chunks: int = 256 + max_incomplete_messages: int = 16 + max_unauthenticated_connections: int = 64 + max_connections_global: int = 256 + max_connections_per_principal: int = 64 + max_pending_per_session: int = 1024 + max_pending_per_principal: int = 2048 + max_pending_global: int = 4096 + max_retained_outcomes_per_session: int = 64 + max_retained_outcomes_global: int = 1024 + max_retained_outcome_bytes_per_session: int = 256 * 1024 * 1024 + max_retained_outcome_bytes_global: int = 1024 * 1024 * 1024 + max_staged_tasks: int = 128 + outbound_queue_size: int = 1024 + inbound_queue_size: int = 1024 + # Bounds callback batches waiting to start; registration itself is nonblocking. + callback_queue_size: int = 1024 + completion_workers: int = 4 + control_timeout: float = 5.0 + handshake_timeout: float = 5.0 + partial_frame_timeout: float = 10.0 + + def __post_init__(self) -> None: + duration_fields = { + "control_timeout", + "handshake_timeout", + "partial_frame_timeout", + } + for item in fields(self): + value = getattr(self, item.name) + if item.name in duration_fields: + _positive(value, item.name) + else: + _positive_int(value, item.name) + if self.max_frame_payload_bytes > self.max_message_bytes: + raise ValueError("max_frame_payload_bytes cannot exceed max_message_bytes") + if self.max_invocation_bytes > self.max_message_bytes: + raise ValueError("max_invocation_bytes cannot exceed max_message_bytes") + if self.max_result_bytes > self.max_message_bytes: + raise ValueError("max_result_bytes cannot exceed max_message_bytes") + if self.max_frame_payload_bytes * self.max_chunks < self.max_message_bytes: + raise ValueError("max_chunks cannot represent max_message_bytes") + if self.max_connections_per_principal > self.max_connections_global: + raise ValueError( + "max_connections_per_principal cannot exceed max_connections_global" + ) + if self.max_pending_per_principal > self.max_pending_global: + raise ValueError( + "max_pending_per_principal cannot exceed max_pending_global" + ) + if self.max_retained_outcomes_per_session > self.max_retained_outcomes_global: + raise ValueError( + "per-session retained outcome count cannot exceed the global count" + ) + if ( + self.max_retained_outcome_bytes_per_session + > self.max_retained_outcome_bytes_global + ): + raise ValueError( + "per-session retained outcome bytes cannot exceed the global bytes" + ) + + +@dataclass(frozen=True, slots=True) +class PeerInfo: + transport: str + address: object + uid: int | None = None + gid: int | None = None + pid: int | None = None + certificate: dict | None = None + + +@dataclass(frozen=True, slots=True) +class Principal: + name: str + + +@runtime_checkable +class Authenticator(Protocol): + def __call__(self, peer: PeerInfo, hello: dict) -> Principal: ... + + +@runtime_checkable +class Authorizer(Protocol): + def __call__( + self, principal: Principal, operation: str, metadata: dict + ) -> bool: ... + + +def _validate_client_tls(context: ssl.SSLContext) -> None: + if context.verify_mode != ssl.CERT_REQUIRED or not context.check_hostname: + raise ValueError("TCP client TLS must verify certificates and hostnames") + if context.minimum_version < ssl.TLSVersion.TLSv1_2: + raise ValueError("TCP client TLS must require TLS 1.2 or newer") + + +def _validate_server_tls(context: ssl.SSLContext) -> None: + if context.verify_mode != ssl.CERT_REQUIRED: + raise ValueError("TCP server TLS must require client certificates") + if context.minimum_version < ssl.TLSVersion.TLSv1_2: + raise ValueError("TCP server TLS must require TLS 1.2 or newer") + + +def _positive_int(value: int, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value <= 0: + raise ValueError(f"{name} must be greater than zero") + + +def _positive(value: int | float, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{name} must be numeric") + if value <= 0 or not math.isfinite(value): + raise ValueError(f"{name} must be finite and greater than zero") diff --git a/deadpool/remote/errors.py b/deadpool/remote/errors.py new file mode 100644 index 0000000..58bbf70 --- /dev/null +++ b/deadpool/remote/errors.py @@ -0,0 +1,121 @@ +"""Public failure taxonomy for remote execution.""" + +from __future__ import annotations + +from enum import Enum + + +class AcceptanceCertainty(str, Enum): + NOT_ACCEPTED = "not_accepted" + ACCEPTED = "accepted" + UNKNOWN = "unknown" + + +class ExecutionCertainty(str, Enum): + NOT_STARTED = "not_started" + MAY_HAVE_RUN = "may_have_run" + COMPLETED = "completed" + + +class RemoteExecutorError(Exception): + """Base class carrying the facts an application needs before retrying.""" + + default_acceptance = AcceptanceCertainty.UNKNOWN + default_execution = ExecutionCertainty.MAY_HAVE_RUN + + def __init__( + self, + message: str = "", + *, + request_id: str | None = None, + acceptance_certainty: AcceptanceCertainty | str | None = None, + execution_certainty: ExecutionCertainty | str | None = None, + ) -> None: + super().__init__(message) + self.request_id = request_id + self.acceptance_certainty = AcceptanceCertainty( + acceptance_certainty or self.default_acceptance + ) + self.execution_certainty = ExecutionCertainty( + execution_certainty or self.default_execution + ) + + +class RemoteExecutorUnavailable(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.NOT_ACCEPTED + default_execution = ExecutionCertainty.NOT_STARTED + + +class RemoteAuthenticationError(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.NOT_ACCEPTED + default_execution = ExecutionCertainty.NOT_STARTED + + +class RemoteCompatibilityError(RemoteAuthenticationError): ... + + +class RemoteQueueFull(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.NOT_ACCEPTED + default_execution = ExecutionCertainty.NOT_STARTED + + +class RemoteSubmissionTimeout(RemoteQueueFull): ... + + +class SubmissionOutcomeUnknown(RemoteExecutorError): ... + + +class RemoteQueueTimeout(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.ACCEPTED + default_execution = ExecutionCertainty.NOT_STARTED + + +class RemoteTaskError(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.ACCEPTED + default_execution = ExecutionCertainty.COMPLETED + + def __init__( + self, + message: str = "Remote task failed", + *, + remote_traceback: str = "", + **kwargs, + ) -> None: + super().__init__(message, **kwargs) + self.remote_traceback = remote_traceback + + +class RemoteResultEncodingError(RemoteTaskError): ... + + +class RemoteResultTooLarge(RemoteTaskError): ... + + +class RemoteCancellationOutcomeUnknown(RemoteExecutorError): ... + + +class RemoteForkedProcessError(RemoteExecutorError): ... + + +class RemoteProcessError(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.ACCEPTED + default_execution = ExecutionCertainty.MAY_HAVE_RUN + + +class RemoteCancelledError(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.ACCEPTED + default_execution = ExecutionCertainty.NOT_STARTED + + +class RemoteConnectionLost(RemoteExecutorError): ... + + +class RemoteServerRestarted(RemoteConnectionLost): ... + + +class RemoteProtocolError(RemoteExecutorError): ... + + +class RemoteResultLost(RemoteExecutorError): + default_acceptance = AcceptanceCertainty.ACCEPTED + default_execution = ExecutionCertainty.COMPLETED diff --git a/deadpool/remote/serializer.py b/deadpool/remote/serializer.py new file mode 100644 index 0000000..18efdd1 --- /dev/null +++ b/deadpool/remote/serializer.py @@ -0,0 +1,82 @@ +"""Payload serializers for a bidirectional remote execution protocol. + +The default :class:`PickleSerializer` requires mutually trusted clients and +servers in both callable and registered-task modes. Registered operation +selection is authorization, not a deserialization sandbox: invocation arguments +are still unpickled by workers. Results and task exceptions are unpickled by +clients, establishing the reverse trust relationship as well. +""" + +from __future__ import annotations + +import pickle +from typing import Protocol, runtime_checkable + +PICKLE_TRUST_WARNING = ( + "The default PickleSerializer requires mutually trusted clients and servers " + "in both callable and registered-task modes: registered operation " + "authorization is not a deserialization sandbox, and result/exception " + "payloads require clients to trust servers in reverse" +) + + +class SerializationLimitError(ValueError): ... + + +@runtime_checkable +class Serializer(Protocol): + """Encode invocations and outcomes across the client/server boundary. + + Implementations define their own trust model. The bundled pickle + implementation is unsafe for untrusted input in either direction. + """ + + name: str + protocol_version: str + + def dumps(self, value: object, *, limit: int) -> bytes: ... + + def loads(self, payload: bytes) -> object: ... + + +class PickleSerializer: + """Serialize protocol payloads with pickle for mutually trusted peers. + + This serializer unpickles callable or registered-task invocation data on + the server and result or exception data on the client. Registered task + authorization limits callable selection but does not make unpickling safe. + """ + + name = "pickle" + protocol_version = str(pickle.HIGHEST_PROTOCOL) + + def __init__(self, protocol: int = pickle.HIGHEST_PROTOCOL) -> None: + self.protocol = protocol + self.protocol_version = str(protocol) + + def dumps(self, value: object, *, limit: int) -> bytes: + writer = _LimitedWriter(limit) + pickle.Pickler(writer, protocol=self.protocol).dump(value) + return bytes(writer.buffer) + + def loads(self, payload: bytes) -> object: + return pickle.loads(payload) + + +class _LimitedWriter: + """Pickle target which aborts before retaining bytes above the limit.""" + + def __init__(self, limit: int) -> None: + if limit < 0: + raise ValueError("serialization limit must be non-negative") + self.limit = limit + self.buffer = bytearray() + + def write(self, chunk: bytes) -> int: + size = len(chunk) + if len(self.buffer) + size > self.limit: + raise SerializationLimitError( + f"serialized payload exceeds the {self.limit}-byte limit" + ) + self.buffer.extend(chunk) + return size diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py new file mode 100644 index 0000000..5f4571c --- /dev/null +++ b/deadpool/remote/server.py @@ -0,0 +1,1569 @@ +"""Server-side broker owning one local :class:`deadpool.Deadpool`.""" + +from __future__ import annotations + +import hashlib +import itertools +import logging +import math +import os +import pickle +import queue +import socket +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from enum import Enum +from typing import Callable + +from deadpool import Future as LocalFuture +from deadpool import ProcessError, TimeoutError + +from ._protocol import ( + MAJOR, + MINOR, + Message, + MessageReader, + MessageType, + _validate_wire_limits, + _wire_limits, + send_message, + validate_message_control, +) +from ._scheduler import FairScheduler +from ._transport import ( + BoundListener, + accept_socket, + bind_listener, + peer_info, + prepare_accepted_socket, +) +from ._worker import WorkerOutcome, execute_opaque +from .config import ( + Authenticator, + Authorizer, + Principal, + RemoteLimits, + TcpListener, + UnixListener, +) +from .errors import RemoteProtocolError +from .serializer import PICKLE_TRUST_WARNING, PickleSerializer, Serializer + +logger = logging.getLogger("deadpool.remote") + + +class ServerState(str, Enum): + CREATED = "CREATED" + STARTING = "STARTING" + RUNNING = "RUNNING" + DRAINING = "DRAINING" + STOPPING = "STOPPING" + STOPPED = "STOPPED" + + +class RequestState(str, Enum): + ACCEPTED_QUEUED = "ACCEPTED_QUEUED" + RUNNING = "RUNNING" + SUCCEEDED = "SUCCEEDED" + TASK_FAILED = "TASK_FAILED" + EXECUTION_TIMED_OUT = "EXECUTION_TIMED_OUT" + CANCELLED = "CANCELLED" + QUEUE_TIMED_OUT = "QUEUE_TIMED_OUT" + WORKER_LOST = "WORKER_LOST" + RESULT_ENCODING_FAILED = "RESULT_ENCODING_FAILED" + RESULT_TOO_LARGE = "RESULT_TOO_LARGE" + + +_TERMINAL = { + RequestState.SUCCEEDED, + RequestState.TASK_FAILED, + RequestState.EXECUTION_TIMED_OUT, + RequestState.CANCELLED, + RequestState.QUEUE_TIMED_OUT, + RequestState.WORKER_LOST, + RequestState.RESULT_ENCODING_FAILED, + RequestState.RESULT_TOO_LARGE, +} + + +@dataclass(slots=True, eq=False) +class _Request: + request_id: str + digest: str + principal: str + session_id: str + connection: "_ServerConnection" + mode: str + operation: str | None + payload: bytes + priority: int + execution_timeout: float | None + queue_deadline: float | None + group_id: str | None + max_result_bytes: int + reserved_result_bytes: int + state: RequestState = RequestState.ACCEPTED_QUEUED + local_future: LocalFuture | None = None + terminal_kind: MessageType | None = None + terminal_payload: bytes = b"" + terminal_control: dict = field(default_factory=dict) + reservation_released: bool = False + + +@dataclass(slots=True) +class _Session: + session_id: str + client_id: str + principal: str + connection: "_ServerConnection" + max_result_bytes: int + watermark: int = 0 + reserved_outcomes: int = 0 + reserved_outcome_bytes: int = 0 + requests: dict[str, _Request] = field(default_factory=dict) + + +class _PoolFuture(LocalFuture): + def __init__(self, server: "DeadpoolServer", record: _Request) -> None: + super().__init__() + self._before_submit = lambda pid: server._claim_dispatch(record, pid) + + +@dataclass(slots=True) +class _WriterBarrier: + reached: threading.Event = field(default_factory=threading.Event) + + +class _ServerConnection: + def __init__( + self, server: "DeadpoolServer", sock: socket.socket, transport: str + ) -> None: + self.server = server + self.sock = sock + self.transport = transport + self.reader = MessageReader(server.limits) + self.session: _Session | None = None + self.closed = False + self.clean_close = False + self._close_lock = threading.Lock() + self._sequence = itertools.count() + self._outbound: queue.PriorityQueue[ + tuple[int, int, Message | _WriterBarrier | None] + ] = queue.PriorityQueue(server.limits.outbound_queue_size) + threading.Thread( + target=self._writer_loop, + name="deadpool.remote.writer", + daemon=True, + ).start() + + def send(self, kind: MessageType, control: dict, payload: bytes = b"") -> bool: + if self.closed: + return False + try: + # FIFO is intentional: state transitions and terminal outcomes may + # never be overtaken by GOAWAY or their control acknowledgements. + self._outbound.put_nowait( + (0, next(self._sequence), Message(kind, control, payload)) + ) + return True + except queue.Full: + self.close() + return False + + def _writer_loop(self) -> None: + while True: + _, _, item = self._outbound.get() + try: + if item is None: + return + if isinstance(item, _WriterBarrier): + item.reached.set() + continue + send_message(self.sock, item, self.server.limits) + except (OSError, EOFError, TimeoutError, RemoteProtocolError): + self.close() + return + finally: + self._outbound.task_done() + + def flush(self, timeout: float) -> bool: + if self.closed: + return False + barrier = _WriterBarrier() + try: + self._outbound.put_nowait((0, next(self._sequence), barrier)) + except queue.Full: + self.close() + return False + return barrier.reached.wait(timeout) + + def close(self) -> None: + with self._close_lock: + if self.closed: + return + self.closed = True + try: + self._outbound.put_nowait((0, next(self._sequence), None)) + except queue.Full: + pass + try: + self.sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + self.sock.close() + + +class DeadpoolServer: + """Own a Deadpool and expose it over bounded Unix/TCP connections. + + The wire format is private and experimental until a separate stable wire + reference is published. The default :class:`PickleSerializer` requires + mutually trusted clients and servers in callable and registered-task modes; + operation authorization is not a deserialization sandbox, and serialized + results and exceptions make trust bidirectional. + """ + + def __init__( + self, + pool_factory: Callable[[], object], + *, + listeners: list[UnixListener | TcpListener], + serializer: Serializer | None = None, + task_registry: dict[str, Callable] | None = None, + registry_fingerprint: str | None = None, + application_fingerprint: str | None = None, + authenticator: Authenticator | None = None, + authorizer: Authorizer | None = None, + limits: RemoteLimits | None = None, + scheduler=None, + disconnect_policy: str = "cancel_queued", + ) -> None: + if not listeners: + raise ValueError("at least one listener is required") + if disconnect_policy not in {"cancel_queued", "continue", "terminate"}: + raise ValueError("invalid disconnect policy") + self.pool_factory = pool_factory + self.listeners = list(listeners) + if serializer is None: + logger.warning(PICKLE_TRUST_WARNING) + self.serializer = serializer or PickleSerializer() + self.task_registry = dict(task_registry or {}) + self.registry_fingerprint = registry_fingerprint + self.application_fingerprint = application_fingerprint + self.authenticator = authenticator + self.authorizer = authorizer + self.limits = limits or RemoteLimits() + self.disconnect_policy = disconnect_policy + self.server_id = uuid.uuid4().hex + self.epoch = uuid.uuid4().hex + self.state = ServerState.CREATED + self.ready = threading.Event() + self._stopped = threading.Event() + self._lock = threading.RLock() + self._condition = threading.Condition(self._lock) + self._bound: list[BoundListener] = [] + self._connections: set[_ServerConnection] = set() + self._raw_sockets: set[socket.socket] = set() + self._sessions: dict[str, _Session] = {} + self._client_sessions: dict[str, _Session] = {} + self._requests: dict[str, _Request] = {} + self._reserved_outcomes = 0 + self._reserved_outcome_bytes = 0 + self._scheduler: FairScheduler[_Request] = scheduler or FairScheduler() + self._staged = 0 + self._pool = None + self._serve_thread: threading.Thread | None = None + self._shutdown_thread: threading.Thread | None = None + self._shutdown_cancel_futures = False + self._shutdown_deadline_at: float | None = None + self._startup_error: BaseException | None = None + self._statistics = { + "remote_tasks_received": 0, + "remote_tasks_accepted": 0, + "remote_tasks_rejected": 0, + "remote_tasks_running": 0, + "remote_tasks_terminal": 0, + "remote_connections": 0, + } + + @property + def bound_addresses(self) -> tuple[object, ...]: + return tuple(item.address for item in self._bound) + + def start(self) -> "DeadpoolServer": + with self._lock: + if self._serve_thread is None: + if self.state != ServerState.CREATED: + raise RuntimeError("a stopped remote server cannot be started") + self.state = ServerState.STARTING + self._serve_thread = threading.Thread( + target=self.serve_forever, + name="deadpool.remote.server", + daemon=False, + ) + try: + self._serve_thread.start() + except BaseException as error: + self._startup_error = error + self.state = ServerState.STOPPED + self.ready.set() + self._stopped.set() + raise + # Pool construction and listener binding are not handshake operations. + # Every caller waits for their shared readiness outcome. + if not self.wait_ready(): + raise RuntimeError("remote server stopped before becoming ready") + return self + + def wait_ready(self, timeout: float | None = None) -> bool: + if not self.ready.wait(timeout): + raise TimeoutError("remote server did not become ready") + if self._startup_error is not None: + raise self._startup_error + return self._pool is not None + + def serve_forever(self) -> None: + with self._condition: + current_thread = threading.current_thread() + background_start = self._serve_thread is not None + if self._serve_thread is None: + if self.state != ServerState.CREATED: + raise RuntimeError("a stopped remote server cannot be served") + self._serve_thread = current_thread + self.state = ServerState.STARTING + elif self._serve_thread is not current_thread: + raise RuntimeError("remote server is already being served") + try: + self._initialize() + except BaseException as error: + with self._lock: + self._startup_error = error + self.state = ServerState.STOPPED + self._close_bound() + self.ready.set() + self._stopped.set() + if not background_start: + raise + return + self._stopped.wait() + + def _initialize(self) -> None: + with self._condition: + if self.state != ServerState.STARTING: + if self.state == ServerState.STOPPING: + self.state = ServerState.STOPPED + self.ready.set() + self._stopped.set() + self._condition.notify_all() + return + pool = self.pool_factory() + bound: list[BoundListener] = [] + try: + for listener in self.listeners: + try: + bound.append(bind_listener(listener)) + except BaseException: + if not listener.optional: + raise + if not bound: + raise OSError("no listeners could be bound") + except BaseException: + self._discard_startup_resources(pool, bound) + raise + with self._condition: + if self.state != ServerState.STARTING: + startup_cancelled = True + else: + startup_cancelled = False + self._pool = pool + self._bound = bound + self.state = ServerState.RUNNING + threading.Thread( + target=self._dispatch_loop, + name="deadpool.remote.broker", + daemon=True, + ).start() + for item in bound: + threading.Thread( + target=self._accept_loop, + args=(item,), + name="deadpool.remote.accept", + daemon=True, + ).start() + self.ready.set() + if startup_cancelled: + self._discard_startup_resources(pool, bound) + with self._condition: + self.state = ServerState.STOPPED + self.ready.set() + self._stopped.set() + self._condition.notify_all() + + def _discard_startup_resources( + self, pool: object, bound: list[BoundListener] + ) -> None: + for item in bound: + try: + item.close() + except BaseException: + logger.exception("failed to close an unpublished remote listener") + try: + pool.shutdown(wait=True, cancel_futures=True) + except BaseException: + logger.exception("failed to shut down an unpublished remote pool") + + def _accept_loop(self, bound: BoundListener) -> None: + while self.state in {ServerState.RUNNING, ServerState.DRAINING}: + try: + sock = accept_socket(bound) + except OSError: + if self.state in {ServerState.STOPPING, ServerState.STOPPED}: + return + continue + with self._lock: + unauthenticated = len(self._raw_sockets) + sum( + connection.session is None for connection in self._connections + ) + total = len(self._raw_sockets) + len(self._connections) + if ( + unauthenticated >= self.limits.max_unauthenticated_connections + or total >= self.limits.max_connections_global + ): + sock.close() + continue + self._raw_sockets.add(sock) + threading.Thread( + target=self._prepare_connection, + args=(bound, sock), + name="deadpool.remote.handshake", + daemon=True, + ).start() + + def _prepare_connection(self, bound: BoundListener, raw: socket.socket) -> None: + connection = None + try: + sock = prepare_accepted_socket( + bound, + raw, + handshake_timeout=self.limits.handshake_timeout, + ) + transport = "unix" if isinstance(bound.config, UnixListener) else "tcp" + connection = _ServerConnection(self, sock, transport) + with self._lock: + self._raw_sockets.discard(raw) + if self.state not in {ServerState.RUNNING, ServerState.DRAINING}: + connection.close() + return + self._connections.add(connection) + self._statistics["remote_connections"] = len(self._connections) + self._connection_loop(connection) + except (OSError, TimeoutError): + try: + raw.close() + except OSError: + pass + finally: + with self._lock: + self._raw_sockets.discard(raw) + if connection is not None and not connection.closed: + self._disconnect(connection) + + def _connection_loop(self, connection: _ServerConnection) -> None: + connection.sock.settimeout(self.limits.handshake_timeout) + try: + hello = connection.reader.receive( + connection.sock, + deadline=time.monotonic() + self.limits.handshake_timeout, + ) + if hello.kind != MessageType.HELLO or hello.payload: + raise RemoteProtocolError("HELLO must be the first message") + self._handshake(connection, hello.control) + connection.sock.settimeout(None) + while not connection.closed: + message = connection.reader.receive(connection.sock) + self._handle_message(connection, message) + except (EOFError, OSError, TimeoutError): + pass + except Exception as error: + protocol_error = ( + error + if isinstance(error, RemoteProtocolError) + else RemoteProtocolError(f"malformed peer message: {error}") + ) + connection.send( + MessageType.PROTOCOL_ERROR, + {"message": str(protocol_error)[:1024]}, + ) + connection.flush(min(self.limits.control_timeout, 1.0)) + finally: + self._disconnect(connection) + + def _handshake(self, connection: _ServerConnection, hello: dict) -> None: + _validate_hello(hello) + if hello["wire_limits"] != _wire_limits(self.limits): + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "wire_limits"}, + ) + raise RemoteProtocolError("wire limit compatibility rejected") + if [MAJOR, MINOR] not in hello["versions"]: + connection.send(MessageType.HANDSHAKE_REJECTED, {"reason": "protocol"}) + raise RemoteProtocolError("no compatible protocol version") + if hello["serializer"] != self.serializer.name or str( + hello["serializer_protocol"] + ) != str(self.serializer.protocol_version): + connection.send(MessageType.HANDSHAKE_REJECTED, {"reason": "serializer"}) + raise RemoteProtocolError("serializer compatibility rejected") + if "callable" in hello["capabilities"] and hello["python"] != [ + sys.version_info.major, + sys.version_info.minor, + ]: + connection.send(MessageType.HANDSHAKE_REJECTED, {"reason": "python"}) + raise RemoteProtocolError("callable mode requires matching Python") + if ( + self.registry_fingerprint is not None + and hello.get("registry_fingerprint") != self.registry_fingerprint + ): + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "registry_fingerprint"}, + ) + raise RemoteProtocolError("task registry compatibility rejected") + if ( + self.application_fingerprint is not None + and hello.get("application_fingerprint") != self.application_fingerprint + ): + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "application_fingerprint"}, + ) + raise RemoteProtocolError("application compatibility rejected") + peer = peer_info(connection.sock, connection.transport) + try: + if self.authenticator is not None: + principal = self.authenticator(peer, hello) + elif connection.transport == "unix": + if peer.uid is not None and peer.uid != os.getuid(): + raise PermissionError("Unix peer UID is not authorized") + principal = Principal( + f"uid:{peer.uid if peer.uid is not None else os.getuid()}" + ) + elif peer.certificate: + identity = hashlib.sha256( + repr(peer.certificate).encode("utf-8") + ).hexdigest() + principal = Principal(f"tls:{identity}") + else: + # Plaintext is restricted by configuration to loopback. Its + # identity deliberately excludes the ephemeral source port so + # opening more sessions cannot increase scheduler share. + principal = Principal("tcp:loopback") + if not isinstance(principal, Principal) or not principal.name: + raise TypeError("authenticator must return a non-empty Principal") + except Exception as error: + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "authentication"}, + ) + raise RemoteProtocolError("peer authentication rejected") from error + offered_result_limit = hello.get("max_result_bytes") + if ( + isinstance(offered_result_limit, bool) + or not isinstance(offered_result_limit, int) + or offered_result_limit <= 0 + ): + raise RemoteProtocolError("client result limit must be a positive integer") + effective_result_limit = min( + offered_result_limit, + self.limits.max_result_bytes, + self.limits.max_retained_outcome_bytes_per_session, + self.limits.max_retained_outcome_bytes_global, + ) + client_id = hello["client_instance_id"] + session_id = uuid.uuid4().hex + session = _Session( + session_id, + client_id, + principal.name, + connection, + effective_result_limit, + ) + with self._condition: + if client_id in self._client_sessions: + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "duplicate_client_instance"}, + ) + raise RemoteProtocolError("client instance already has a live session") + principal_connections = sum( + item.principal == principal.name for item in self._sessions.values() + ) + if principal_connections >= self.limits.max_connections_per_principal: + connection.send( + MessageType.HANDSHAKE_REJECTED, + {"reason": "connection_limit"}, + ) + raise RemoteProtocolError("principal connection limit exceeded") + connection.session = session + self._sessions[session_id] = session + self._client_sessions[client_id] = session + connection.send( + MessageType.WELCOME, + { + "version": [MAJOR, MINOR], + "features": ["callable", "registered", "chunking"], + "wire": "experimental-deadpool-private-v1", + "wire_limits": _wire_limits(self.limits), + "server_id": self.server_id, + "epoch": self.epoch, + "session_id": session_id, + "serializer": self.serializer.name, + "serializer_protocol": self.serializer.protocol_version, + "registry_fingerprint": self.registry_fingerprint, + "application_fingerprint": self.application_fingerprint, + "limits": { + "max_invocation_bytes": self.limits.max_invocation_bytes, + "max_result_bytes": effective_result_limit, + "max_pending_per_session": self.limits.max_pending_per_session, + "max_pending_per_principal": self.limits.max_pending_per_principal, + "max_pending_global": self.limits.max_pending_global, + }, + "resumption": False, + }, + ) + + def _handle_message(self, connection: _ServerConnection, message: Message) -> None: + if message.kind == MessageType.PING: + connection.send(MessageType.PONG, {"nonce": message.control.get("nonce")}) + elif message.kind == MessageType.SUBMIT: + self._submit(connection, message) + elif message.kind == MessageType.CANCEL_REQUEST: + self._cancel(connection, message.control) + elif message.kind == MessageType.CANCEL_GROUP: + self._cancel_group(connection, message.control) + elif message.kind == MessageType.STATUS_REQUEST: + self._status(connection, message.control) + elif message.kind == MessageType.STATS_REQUEST: + if self._authorized(connection, "statistics", message.control): + connection.send( + MessageType.STATS_RESPONSE, + { + "statistics": self.get_statistics(), + "token": message.control.get("token"), + }, + ) + else: + connection.send( + MessageType.STATS_RESPONSE, + {"error": "unauthorized", "token": message.control.get("token")}, + ) + elif message.kind == MessageType.RESULT_ACK: + with self._condition: + request = self._owned_request( + connection, message.control.get("request_id") + ) + if request is not None and request.state in _TERMINAL: + self._remove_request_locked(request) + self._condition.notify_all() + elif message.kind == MessageType.CLOSE_SESSION: + session = connection.session + with self._condition: + if session is None or any( + record.state not in _TERMINAL + for record in session.requests.values() + ): + raise RemoteProtocolError( + "a session may close cleanly only after all requests are terminal" + ) + connection.clean_close = True + connection.send( + MessageType.CLOSE_SESSION_RESPONSE, + {"token": message.control.get("token")}, + ) + else: + raise RemoteProtocolError( + f"message {message.kind.name} is invalid in this state" + ) + + def _submit(self, connection: _ServerConnection, message: Message) -> None: + session = connection.session + if session is None: + raise RemoteProtocolError("submission before handshake") + control = message.control + request_id = control.get("request_id") + self._statistics["remote_tasks_received"] += 1 + if not isinstance(request_id, str) or not request_id.startswith( + session.client_id + ":" + ): + raise RemoteProtocolError( + "request ID does not belong to the client instance" + ) + try: + sequence = int(request_id.rsplit(":", 1)[1]) + except (ValueError, IndexError) as error: + raise RemoteProtocolError("invalid request sequence") from error + digest = hashlib.sha256(message.payload).hexdigest() + if control.get("digest") != digest: + raise RemoteProtocolError("invocation digest mismatch") + with self._condition: + existing = session.requests.get(request_id) + if existing is not None: + if existing.digest != digest: + raise RemoteProtocolError( + "request ID reused with different payload" + ) + self._send_record(existing) + return + if sequence <= session.watermark: + self._reject(connection, request_id, "stale_request") + return + session.watermark = sequence + if self.state != ServerState.RUNNING: + self._reject(connection, request_id, "server_draining") + return + # Active and unacknowledged-terminal records both consume finite + # bookkeeping. Reserve worst-case result retention before ACCEPTED + # so an accepted execution can always commit an explicit outcome. + principal_records = sum( + record.principal == session.principal + for record in self._requests.values() + ) + if ( + len(session.requests) >= self.limits.max_pending_per_session + or principal_records >= self.limits.max_pending_per_principal + or len(self._requests) >= self.limits.max_pending_global + or session.reserved_outcomes + >= self.limits.max_retained_outcomes_per_session + or self._reserved_outcomes >= self.limits.max_retained_outcomes_global + or session.reserved_outcome_bytes + session.max_result_bytes + > self.limits.max_retained_outcome_bytes_per_session + or self._reserved_outcome_bytes + session.max_result_bytes + > self.limits.max_retained_outcome_bytes_global + ): + self._reject(connection, request_id, "queue_full") + return + if len(message.payload) > self.limits.max_invocation_bytes: + self._reject(connection, request_id, "invocation_too_large") + return + mode = control.get("mode") + operation = control.get("operation") + if mode not in {"callable", "registered"}: + self._reject(connection, request_id, "invalid_mode") + return + if mode == "registered" and operation not in self.task_registry: + self._reject(connection, request_id, "unknown_operation") + return + action = ( + "submit_callable" if mode == "callable" else f"submit_task:{operation}" + ) + if not self._authorized(connection, action, control): + self._reject(connection, request_id, "unauthorized") + return + priority = _finite_nonnegative(control.get("priority", 0), "priority") + execution_timeout = _optional_finite_nonnegative( + control.get("execution_timeout"), "execution_timeout" + ) + queue_timeout = _optional_finite_nonnegative( + control.get("queue_timeout"), "queue_timeout" + ) + record = _Request( + request_id, + digest, + session.principal, + session.session_id, + connection, + mode, + operation, + message.payload, + int(priority), + execution_timeout, + time.monotonic() + queue_timeout if queue_timeout is not None else None, + control.get("group_id"), + session.max_result_bytes, + session.max_result_bytes, + ) + session.requests[request_id] = record + self._requests[request_id] = record + session.reserved_outcomes += 1 + session.reserved_outcome_bytes += record.reserved_result_bytes + self._reserved_outcomes += 1 + self._reserved_outcome_bytes += record.reserved_result_bytes + self._scheduler.put( + record, priority=record.priority, principal=record.principal + ) + self._statistics["remote_tasks_accepted"] += 1 + connection.send( + MessageType.ACCEPTED, {"request_id": request_id, "state": record.state} + ) + self._condition.notify_all() + + def _reject( + self, connection: _ServerConnection, request_id: str, reason: str + ) -> None: + self._statistics["remote_tasks_rejected"] += 1 + connection.send( + MessageType.REJECTED, {"request_id": request_id, "reason": reason} + ) + + def _dispatch_loop(self) -> None: + while True: + with self._condition: + self._expire_queued_locked() + while self.state in {ServerState.RUNNING, ServerState.DRAINING} and ( + not self._scheduler or self._staged >= self.limits.max_staged_tasks + ): + self._condition.wait(0.05) + self._expire_queued_locked() + if self.state in {ServerState.STOPPING, ServerState.STOPPED}: + return + try: + record = self._scheduler.pop() + except IndexError: + continue + if record.state != RequestState.ACCEPTED_QUEUED: + continue + local = _PoolFuture(self, record) + record.local_future = local + self._staged += 1 + try: + self._pool._submit_future( + local, + execute_opaque, + ( + self.serializer, + record.mode, + record.payload, + record.operation, + ( + {record.operation: self.task_registry[record.operation]} + if record.mode == "registered" + else None + ), + record.max_result_bytes, + ), + {}, + record.execution_timeout, + record.priority, + block=False, + ) + except queue.Full: + # The broker owns deadline processing, so it must never wait + # inside the local pool's backpressure path. Restore the record + # ahead of later work in its FIFO stream without changing the + # principal's round-robin position, then yield the lane. + with self._condition: + record.local_future = None + self._staged -= 1 + if record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.put_front( + record, + priority=record.priority, + principal=record.principal, + ) + self._condition.wait(0.01) + self._condition.notify_all() + except BaseException as error: + self._complete_local(record, error) + else: + # The local queue owns the bytes through its job tuple now; the + # broker record no longer retains a duplicate invocation. + record.payload = b"" + local.add_done_callback( + lambda future, record=record: self._pool_done(record, future) + ) + + def _expire_queued_locked(self) -> None: + now = time.monotonic() + for record in list(self._requests.values()): + if ( + record.state == RequestState.ACCEPTED_QUEUED + and record.queue_deadline is not None + and now >= record.queue_deadline + ): + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.QUEUE_TIMED_OUT, + MessageType.QUEUE_TIMED_OUT, + {"reason": "queue_timeout"}, + ) + if record.local_future is not None: + record.local_future.cancel() + + def _claim_dispatch(self, record: _Request, pid: int) -> bool: + """Atomically cross the queue/worker boundary before pipe I/O. + + Marking the request running makes cancellation and queue expiry lose the + arbitration race without holding the broker lock while a potentially + large invocation is written to the worker pipe. A broken worker retry + may claim the same already-running request until it becomes terminal. + """ + with self._condition: + if record.state == RequestState.ACCEPTED_QUEUED: + if ( + record.queue_deadline is not None + and time.monotonic() >= record.queue_deadline + ): + self._terminal_locked( + record, + RequestState.QUEUE_TIMED_OUT, + MessageType.QUEUE_TIMED_OUT, + {"reason": "queue_timeout"}, + ) + if record.local_future is not None: + record.local_future.cancel() + return False + record.state = RequestState.RUNNING + self._statistics["remote_tasks_running"] += 1 + elif record.state != RequestState.RUNNING: + return False + # Publish the selected PID before pipe I/O so a concurrent hard + # cancellation can terminate a writer blocked in submit_job(). + if record.local_future is not None: + record.local_future.pid = pid + # Broken-pipe retries publish the replacement worker identity + # without counting the request as running twice. + record.connection.send( + MessageType.RUNNING, + { + "request_id": record.request_id, + "pid": pid, + "worker_id": f"worker:{pid}", + }, + ) + return True + + def _pool_done(self, record: _Request, future: LocalFuture) -> None: + try: + outcome = future.result() + except BaseException as error: + self._complete_local(record, error) + else: + self._complete_local(record, outcome) + + def _complete_local(self, record: _Request, outcome: object) -> None: + with self._condition: + if record.state in _TERMINAL: + if self._staged: + self._staged -= 1 + if record.connection.closed: + session = self._sessions.get(record.session_id) + if session is not None: + self._purge_session_locked(session) + self._condition.notify_all() + return + if isinstance(outcome, WorkerOutcome): + mapping = { + "result": (RequestState.SUCCEEDED, MessageType.RESULT), + "task_error": (RequestState.TASK_FAILED, MessageType.TASK_ERROR), + "result_encoding_failed": ( + RequestState.RESULT_ENCODING_FAILED, + MessageType.RESULT_ENCODING_FAILED, + ), + "result_too_large": ( + RequestState.RESULT_TOO_LARGE, + MessageType.RESULT_TOO_LARGE, + ), + } + state, kind = mapping[outcome.kind] + descriptor = _fit_terminal_descriptor( + record.request_id, + outcome.descriptor or {}, + len(outcome.payload), + self.limits, + ) + self._terminal_locked(record, state, kind, descriptor, outcome.payload) + elif isinstance(outcome, TimeoutError): + self._terminal_locked( + record, + RequestState.EXECUTION_TIMED_OUT, + MessageType.TIMED_OUT, + {"message": str(outcome)}, + ) + elif isinstance(outcome, (ProcessError, pickle.PicklingError)): + self._terminal_locked( + record, + RequestState.WORKER_LOST, + MessageType.WORKER_LOST, + {"message": str(outcome)}, + ) + elif isinstance(outcome, BaseException): + self._terminal_locked( + record, + RequestState.WORKER_LOST, + MessageType.WORKER_LOST, + {"message": repr(outcome)}, + ) + else: + self._terminal_locked( + record, + RequestState.WORKER_LOST, + MessageType.WORKER_LOST, + {"message": "invalid worker outcome"}, + ) + if self._staged: + self._staged -= 1 + if record.connection.closed: + session = self._sessions.get(record.session_id) + if session is not None: + self._purge_session_locked(session) + self._condition.notify_all() + + def _terminal_locked( + self, + record: _Request, + state: RequestState, + kind: MessageType, + control: dict, + payload: bytes = b"", + ) -> None: + if record.state in _TERMINAL: + return + record.state = state + record.payload = b"" + record.terminal_kind = kind + record.terminal_control = dict(control) + record.terminal_payload = payload + self._statistics["remote_tasks_terminal"] += 1 + record.connection.send( + kind, {"request_id": record.request_id, **control}, payload + ) + + def _cancel(self, connection: _ServerConnection, control: dict) -> None: + request_id = control.get("request_id") + hard = bool(control.get("hard", False)) + token = control.get("token") + with self._condition: + record = self._owned_request(connection, request_id) + if record is None: + response = "unknown" + elif record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "cancelled"}, + ) + if record.local_future is not None: + record.local_future.cancel() + response = "cancelled" + elif ( + record.state == RequestState.RUNNING + and hard + and self._authorized(connection, "hard_cancel", control) + ): + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "terminated"}, + ) + if record.local_future is not None: + record.local_future.cancel_and_kill_if_running() + response = "cancelled" + elif record.state == RequestState.RUNNING: + response = "running" + else: + response = "terminal" + connection.send( + MessageType.CANCEL_RESPONSE, + {"request_id": request_id, "response": response, "token": token}, + ) + self._condition.notify_all() + + def _cancel_group(self, connection: _ServerConnection, control: dict) -> None: + group_id = control.get("group_id") + hard = bool(control.get("hard", False)) + counts = {"cancelled": 0, "running": 0, "terminal": 0, "unknown": 0} + if not isinstance(group_id, str): + counts["unknown"] = 1 + else: + with self._condition: + session = connection.session + records = ( + [ + record + for record in session.requests.values() + if record.group_id == group_id + ] + if session is not None + else [] + ) + if not records: + counts["unknown"] = 1 + for record in records: + if record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "group_cancelled"}, + ) + if record.local_future is not None: + record.local_future.cancel() + counts["cancelled"] += 1 + elif record.state == RequestState.RUNNING: + if hard and self._authorized( + connection, "hard_cancel", control + ): + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "group_terminated"}, + ) + if record.local_future is not None: + record.local_future.cancel_and_kill_if_running() + counts["cancelled"] += 1 + else: + counts["running"] += 1 + else: + counts["terminal"] += 1 + self._condition.notify_all() + connection.send( + MessageType.CANCEL_GROUP_RESPONSE, + {**counts, "token": control.get("token")}, + ) + + def _status(self, connection: _ServerConnection, control: dict) -> None: + record = self._owned_request(connection, control.get("request_id")) + connection.send( + MessageType.STATUS_RESPONSE, + { + "request_id": control.get("request_id"), + "state": record.state if record else "UNKNOWN", + "token": control.get("token"), + }, + ) + + def _owned_request( + self, connection: _ServerConnection, request_id: object + ) -> _Request | None: + session = connection.session + if session is None or not isinstance(request_id, str): + return None + return session.requests.get(request_id) + + def _send_record(self, record: _Request) -> None: + if record.state == RequestState.ACCEPTED_QUEUED: + record.connection.send( + MessageType.ACCEPTED, + {"request_id": record.request_id, "state": record.state}, + ) + elif record.state == RequestState.RUNNING: + record.connection.send( + MessageType.RUNNING, + { + "request_id": record.request_id, + "pid": record.local_future.pid if record.local_future else None, + "worker_id": None, + }, + ) + elif record.terminal_kind is not None: + record.connection.send( + record.terminal_kind, + {"request_id": record.request_id, **record.terminal_control}, + record.terminal_payload, + ) + + def _authorized( + self, connection: _ServerConnection, operation: str, metadata: dict + ) -> bool: + session = connection.session + if session is None: + return False + if self.authorizer is None: + # Filesystem permissions plus same-UID peer authentication form the + # explicit built-in local policy. TCP is default-deny even with TLS; + # operators must map certificate principals to allowed operations. + return ( + connection.transport == "unix" + and session.principal == f"uid:{os.getuid()}" + ) + return bool(self.authorizer(Principal(session.principal), operation, metadata)) + + def _disconnect(self, connection: _ServerConnection) -> None: + connection.close() + with self._condition: + self._connections.discard(connection) + self._statistics["remote_connections"] = len(self._connections) + session = connection.session + if ( + session is not None + and not connection.clean_close + and self.disconnect_policy in {"cancel_queued", "terminate"} + ): + # Hard cancellation can complete a worker concurrently and + # purge terminal records, so disconnect arbitration iterates a + # stable snapshot of the session's requests. + for record in list(session.requests.values()): + if record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "disconnected"}, + ) + if record.local_future is not None: + record.local_future.cancel() + elif ( + record.state == RequestState.RUNNING + and self.disconnect_policy == "terminate" + and record.local_future is not None + ): + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "disconnected_terminate"}, + ) + record.local_future.cancel_and_kill_if_running() + self._purge_session_locked(session) + elif session is not None: + # Resumption is not negotiated by this wire version. Retaining + # terminal outcomes after disconnect would make them + # unreachable and eventually exhaust the bounded result + # reservation. Active ``continue`` requests stay attached to + # the session and purge themselves when they become terminal. + self._purge_session_locked(session) + self._condition.notify_all() + + def _release_reservation_locked(self, session: _Session, record: _Request) -> None: + if record.reservation_released: + return + record.reservation_released = True + session.reserved_outcomes -= 1 + session.reserved_outcome_bytes -= record.reserved_result_bytes + self._reserved_outcomes -= 1 + self._reserved_outcome_bytes -= record.reserved_result_bytes + + def _remove_request_locked(self, record: _Request) -> None: + session = self._sessions.get(record.session_id) + if session is None: + return + self._release_reservation_locked(session, record) + record.payload = b"" + record.terminal_payload = b"" + session.requests.pop(record.request_id, None) + self._requests.pop(record.request_id, None) + + def _purge_session_locked(self, session: _Session) -> None: + for record in list(session.requests.values()): + if record.state in _TERMINAL: + self._remove_request_locked(record) + if not session.requests: + self._sessions.pop(session.session_id, None) + if self._client_sessions.get(session.client_id) is session: + self._client_sessions.pop(session.client_id, None) + + def get_statistics(self) -> dict: + with self._lock: + stats = dict(self._statistics) + stats["remote_sessions"] = len(self._sessions) + stats["remote_queued"] = len(self._scheduler) + stats["remote_staged"] = self._staged + stats["remote_retained_outcomes"] = self._reserved_outcomes + stats["remote_retained_outcome_bytes"] = self._reserved_outcome_bytes + if self._pool is not None: + stats.update(self._pool.get_statistics()) + return stats + + def shutdown( + self, + wait: bool = True, + *, + cancel_futures: bool = False, + deadline: float | None = None, + ) -> None: + validated_deadline = _validate_shutdown_deadline(deadline) + absolute_deadline = ( + time.monotonic() + validated_deadline + if validated_deadline is not None + else None + ) + startup_stopping = False + with self._condition: + if self.state == ServerState.STOPPED: + return + if self.state == ServerState.CREATED: + self.state = ServerState.STOPPED + self.ready.set() + self._stopped.set() + return + if self.state == ServerState.STARTING: + self.state = ServerState.STOPPING + startup_stopping = True + elif self.state == ServerState.STOPPING and not self.ready.is_set(): + startup_stopping = True + else: + if self.state == ServerState.RUNNING: + self.state = ServerState.DRAINING + self._shutdown_cancel_futures |= cancel_futures + if absolute_deadline is not None and ( + self._shutdown_deadline_at is None + or absolute_deadline < self._shutdown_deadline_at + ): + self._shutdown_deadline_at = absolute_deadline + if self._shutdown_thread is None: + self._shutdown_thread = threading.Thread( + target=self._finish_shutdown, + name="deadpool.remote.shutdown", + daemon=False, + ) + self._shutdown_thread.start() + self._condition.notify_all() + if startup_stopping: + if wait and self._serve_thread is not threading.current_thread(): + self._stopped.wait() + elif wait and self._shutdown_thread is not threading.current_thread(): + self._shutdown_thread.join() + + def _finish_shutdown(self) -> None: + with self._condition: + while any( + record.state not in _TERMINAL for record in self._requests.values() + ): + if self._shutdown_cancel_futures: + for record in list(self._requests.values()): + if record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "server_shutdown"}, + ) + if record.local_future is not None: + record.local_future.cancel() + deadline = self._shutdown_deadline_at + if deadline is not None and time.monotonic() >= deadline: + for record in list(self._requests.values()): + if record.state == RequestState.ACCEPTED_QUEUED: + self._scheduler.remove(record) + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "shutdown_deadline"}, + ) + if record.local_future is not None: + record.local_future.cancel() + elif ( + record.state == RequestState.RUNNING + and record.local_future is not None + ): + self._terminal_locked( + record, + RequestState.CANCELLED, + MessageType.CANCELLED, + {"reason": "shutdown_deadline"}, + ) + record.local_future.cancel_and_kill_if_running() + break + self._condition.wait(0.05) + self.state = ServerState.STOPPING + self._close_bound() + if self._pool is not None: + self._pool.shutdown(wait=True, cancel_futures=False) + flush_timeout = min(self.limits.control_timeout, 1.0) + for connection in list(self._connections): + # First drain every previously queued terminal outcome, then send + # GOAWAY and drain it. The explicit writer barriers cannot suffer + # the Event clear/enqueue lost-wakeup race. + connection.flush(flush_timeout) + connection.send(MessageType.GOAWAY, {"reason": "server_stopped"}) + connection.flush(flush_timeout) + connection.close() + with self._condition: + self.state = ServerState.STOPPED + self._stopped.set() + self._condition.notify_all() + + def _close_bound(self) -> None: + for item in self._bound: + item.close() + self._bound.clear() + with self._lock: + raw_sockets, self._raw_sockets = self._raw_sockets, set() + for sock in raw_sockets: + try: + sock.close() + except OSError: + pass + + def __enter__(self) -> "DeadpoolServer": + return self.start() + + def __exit__(self, exc_type, exc, traceback) -> bool: + self.shutdown() + return False + + +def _fit_terminal_descriptor( + request_id: str, + descriptor: dict, + payload_size: int, + limits: RemoteLimits, +) -> dict: + """Retain as much diagnostic text as the complete wire frame can carry.""" + fields = ("module", "type", "message", "serialization_error", "traceback") + bounded = { + name: _truncate_utf8(str(descriptor[name]), limits.max_metadata_bytes) + for name in fields + if name in descriptor + } + if _terminal_descriptor_fits(request_id, bounded, payload_size, limits): + return bounded + + # Remove least-essential/largest fields first. Once removal makes the + # frame valid, restore the largest prefix which still fits exactly. + for name in ("traceback", "serialization_error", "message", "type", "module"): + if name not in bounded: + continue + original = bounded[name] + bounded[name] = "" + if not _terminal_descriptor_fits(request_id, bounded, payload_size, limits): + continue + low = 0 + high = len(original.encode("utf-8")) + best = "" + while low <= high: + middle = (low + high) // 2 + candidate = _truncate_utf8(original, middle) + bounded[name] = candidate + if _terminal_descriptor_fits(request_id, bounded, payload_size, limits): + best = candidate + low = middle + 1 + else: + high = middle - 1 + bounded[name] = best + return bounded + + # A successful handshake is larger than this minimal terminal control, but + # omit optional descriptor keys defensively if a custom bound is extreme. + return {} + + +def _terminal_descriptor_fits( + request_id: str, + descriptor: dict, + payload_size: int, + limits: RemoteLimits, +) -> bool: + try: + validate_message_control( + {"request_id": request_id, **descriptor}, payload_size, limits + ) + except RemoteProtocolError: + return False + return True + + +def _truncate_utf8(value: str, limit: int) -> str: + encoded = value.encode("utf-8", errors="backslashreplace") + if len(encoded) <= limit: + return encoded.decode("utf-8") + marker = b"...[truncated]" + if limit <= len(marker): + return marker[:limit].decode("ascii") + prefix = encoded[: limit - len(marker)].decode("utf-8", errors="ignore") + return prefix + marker.decode("ascii") + + +def _validate_shutdown_deadline(value: object) -> float | None: + if value is None: + return None + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value < 0 + ): + raise ValueError("deadline must be finite and non-negative") + return float(value) + + +def _validate_hello(hello: dict) -> None: + required = { + "versions", + "client_instance_id", + "serializer", + "serializer_protocol", + "python", + "capabilities", + "max_result_bytes", + "wire", + "wire_limits", + } + if not required <= set(hello): + raise RemoteProtocolError("HELLO is missing required fields") + versions = hello["versions"] + python_version = hello["python"] + capabilities = hello["capabilities"] + if not ( + isinstance(versions, list) + and all( + isinstance(version, list) + and len(version) == 2 + and all( + isinstance(item, int) and not isinstance(item, bool) for item in version + ) + for version in versions + ) + ): + raise RemoteProtocolError("HELLO versions must be integer pairs") + if not ( + isinstance(python_version, list) + and len(python_version) == 2 + and all( + isinstance(item, int) and not isinstance(item, bool) + for item in python_version + ) + ): + raise RemoteProtocolError("HELLO python version must be an integer pair") + if not ( + isinstance(capabilities, list) + and all(isinstance(item, str) for item in capabilities) + ): + raise RemoteProtocolError("HELLO capabilities must be strings") + for name in ("client_instance_id", "serializer", "wire"): + if not isinstance(hello[name], str) or not hello[name]: + raise RemoteProtocolError(f"HELLO {name} must be a non-empty string") + if len(hello["client_instance_id"]) > 128: + raise RemoteProtocolError("HELLO client instance ID is too long") + serializer_protocol = hello["serializer_protocol"] + if isinstance(serializer_protocol, bool) or not isinstance( + serializer_protocol, (str, int) + ): + raise RemoteProtocolError("HELLO serializer protocol is invalid") + _validate_wire_limits(hello["wire_limits"]) + result_limit = hello["max_result_bytes"] + if ( + isinstance(result_limit, bool) + or not isinstance(result_limit, int) + or result_limit <= 0 + ): + raise RemoteProtocolError("client result limit must be a positive integer") + + +def _finite_nonnegative(value: object, name: str) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value < 0 + ): + raise RemoteProtocolError(f"{name} must be finite and non-negative") + return float(value) + + +def _optional_finite_nonnegative(value: object, name: str) -> float | None: + if value is None: + return None + return _finite_nonnegative(value, name) diff --git a/noxfile.py b/noxfile.py index 95ccb1b..a0b79c2 100644 --- a/noxfile.py +++ b/noxfile.py @@ -31,12 +31,12 @@ def test(session): ) def testcov(session): session.install(".") - session.install("pytest", "pytest-html", "coverage") + session.install("pytest", "pytest-html", "coverage>=7.10") EXTRA_PYTEST = " ".join(session.posargs) + session.run("coverage", "erase") session.run( *shlex.split( - f"coverage run --concurrency=multiprocessing,thread " - f"-m pytest {EXTRA_PYTEST}" + f"coverage run -m pytest {EXTRA_PYTEST}" f" --html=report.html --self-contained-html" ) ) diff --git a/pyproject.toml b/pyproject.toml index 370dc9d..9f16add 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ test = [ "pytest >= 7", "pytest-cov", "flake8", - "coverage[toml]" + "coverage[toml]>=7.10" ] [tool.black] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..96c7b9b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Test support package for importable multiprocessing fixtures.""" diff --git a/tests/remote/__init__.py b/tests/remote/__init__.py new file mode 100644 index 0000000..02be096 --- /dev/null +++ b/tests/remote/__init__.py @@ -0,0 +1 @@ +"""Remote executor integration tests.""" diff --git a/tests/remote/certs/README.txt b/tests/remote/certs/README.txt new file mode 100644 index 0000000..df8c2ef --- /dev/null +++ b/tests/remote/certs/README.txt @@ -0,0 +1 @@ +Static test-only TLS credentials. Never use outside the test suite. diff --git a/tests/remote/certs/ca.pem b/tests/remote/certs/ca.pem new file mode 100644 index 0000000..e4d024c --- /dev/null +++ b/tests/remote/certs/ca.pem @@ -0,0 +1,19 @@ +-----BEGIN CERTIFICATE----- +MIIDJzCCAg+gAwIBAgIUQ629E5usuJA9rFdH25Ik0xfS3LAwDQYJKoZIhvcNAQEL +BQAwGzEZMBcGA1UEAwwQRGVhZHBvb2wgVGVzdCBDQTAeFw0yNjA3MjkxNjA4MTBa +Fw0zNjA3MjYxNjA4MTBaMBsxGTAXBgNVBAMMEERlYWRwb29sIFRlc3QgQ0EwggEi +MA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQC10Rvk60vZTC6Q9HwstFev0WSr +BHBwI/L7L4436ZzYVhzMA4s2NXCqgVCKWQn2SfsaKr7qn+VR5suIeu1e/MWzDG3J +aFexvfNjrwNk/VeGFF3sucOWVHzeYJyjOQlgBchbJTeB0bHQf+TUCBUN8o7UysSy +md2EsHslMB4JvXgF0SjcUByVDWYbzkaH+Ow/GsswMAAcJ2/FIkCVKxScUawI51xn +XFWDRiN6g49CCVbcFH4alMfC0absaHXMz413p/wOcywXD5Fw7pob3YP1tnEfnF/s +EagXRA4wzsJ9/l8dCzFdH76naeES2HfEzeUsG0TsRjbnFfmklWVEWg4joNNTAgMB +AAGjYzBhMA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/BAQDAgEGMB0GA1UdDgQW +BBTNNofP9ttRTvTgsmIgQRJSZ6uT6DAfBgNVHSMEGDAWgBTNNofP9ttRTvTgsmIg +QRJSZ6uT6DANBgkqhkiG9w0BAQsFAAOCAQEAUxRP//5wJ6XqLqbM1DJLEU7rBA8i +xDhG/1jTvlw6VUrmeOGBgP4iz0QyYQg00fhty8VPhgIdhARO7YpS9o4QBs8Q7Xcf +7HxU+DQDsPrfxSxEkwGRFfP8RTnBEAvBWGpxRauuZrG67D89oMInJaGidZz0qCbv +8JImxzbJJSfYz2Tjv9bq3kkPWTJMNhb+KBjkSN6et5Fa3GTqpFtjimHBZcm+V3Zi +oyjIxIwgukQFyBuOEdSq9Q1VgZslZo/vvvZiBhdtE53udF4hU/l1+un0nHR07MV+ +jKDdV5GdWStSM+JlJ9ow9hrx8WLvpT0Fz+4Lz+KfqBAsiBMUAZAF4Upakg== +-----END CERTIFICATE----- diff --git a/tests/remote/certs/client-key.pem b/tests/remote/certs/client-key.pem new file mode 100644 index 0000000..2041520 --- /dev/null +++ b/tests/remote/certs/client-key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDODr+ti91IiQ1g +TSoSKuiYAsHfcfLURHTgMoSTqDRpUE31i8AVU1deywHzdfp0tVhGWUGtlXS9+TuF +E8f7Q5tWiylDjn45otjOssNBbJ65Uq+t1xae6GqT2e2XFOLksnDkhLbnURfZrp6O +D25Ucxv0eHKot6WFmGQgB7hPi+AnWyqAqofHYEem5mLaaCjyI0EcIrGAPvLFU1QE +BAoZUljs8gD2bIVgvhOwKIE+ukWCbcPioDZ/0YDXlCWmiAg21SCW77yhd5Xx3XjB ++oJMEUEUHKCl1iL29sYpiwwt8wLmpzj0DeVmjFfq1eDjdufA7K+WNFwfaLLL/9yZ +FZf+TyxBAgMBAAECggEAKdWZucjHHNZiZlZUP8sBXPdWWCCgzdK2LdxjyTyho23/ +3eXJEitQ9jheBYlC15YHvYSajnzk7jrfCCYbeFph+8CWGJah7t6uiMyHUs2ULgjp +DPC5sXp7RNg+YemGJStLO0HshmsS3gmTPEFKw495MfmQQcvDKavmTBB1LYU29ZYx +9oWQjUuA7ayx8Xd5z6GgvISAhvVp1P0GCHRPajPIuRHiKOtA7TDGtqAzzf8ysCh+ +Tjk7InHtp1Sg+eiEuzb2y27G9pGKv9UyP8oCz6VWKXA+rD/kGBpcf6hoYrod0bOv +5jPKJB6xDVxzlbmBSWWPKoU9HVfkiwDS0UmzJK3NCQKBgQDo3Uv04zslFufLpZYZ +QPhxZEf7azlhykqKirE0ZBtQln1cftLdL/XYVVcjhqbb64eaOGUaQSK02ft+qCZ+ +uMcasmckHLuDYbq0Z3+GXoYFFVWKD76n60KIKR1I7cOzxmIx+rhC+WCD/SbhXP+1 +CcDTF7EFMUuOk4rBbJ+UcqU6fQKBgQDih6TRdoEAo331UQTA+Aepg5ftg4ddUCLn +tdPI7qInU6nsdUOjHzFXvGLmaC5V4+n5Fsh/2QkqfknPUgIrp1SurdUYv150EKHd +rvMH+pX70vDWVXjvYakNVVPppshuWMI2H8XROGOXlxzGIOKsZt2okD6j1Nxo/xrm +7zqVEiTgFQKBgD3xR9gvbbcy/Zu9Q00abDv1efRWFGB/6A4sUHDoRB/OuDAXiE3Q +CSxvnwtTMSWE5IBQigxO0UWcSnrpjbvduRDP7gG8JuEO6RQ+B70dfbbyctuTzppq +STtg/Go+3PUAS3SSBdvC7DqP53zBOT/WBVXhknQYJcoaUymalGRvjqYdAoGAXSj7 +l0js9J5IBsMy+UkHnacIrB0TZkS2liQo3NGGzjWSBDSVhFgnqrVG1wqxQ5Fff5jn +C+zt4BPftus7CUjfgpbqtCq/ZWwRpWF1gSqE0/OEKCEugwPeyiT1RXnZo8fM1wVq +DQjrrEPxCWoszknfjpsDp8y6eFOxdJeduACmOaECgYAxDsgncB5lnOrkj8W3kAga +0C4eQoHrzaUouVoieiKy2MlgAvHtuh9rHZMqI0kfkssG1ikf8SD5wGOOKVN+dgLb +bShsB5rOpT27vKh35RW8TEB+rBEwTvyBONm21InDlTDS/yl1HwPpRoOYjyzdLs3J +ywYgG7d8Xj1Xk71p/1JrOw== +-----END PRIVATE KEY----- diff --git a/tests/remote/certs/client.pem b/tests/remote/certs/client.pem new file mode 100644 index 0000000..f171cfd --- /dev/null +++ b/tests/remote/certs/client.pem @@ -0,0 +1,20 @@ +-----BEGIN CERTIFICATE----- +MIIDPTCCAiWgAwIBAgIUa1KWjufkag+0cE9E5h90qMk3jqEwDQYJKoZIhvcNAQEL +BQAwGzEZMBcGA1UEAwwQRGVhZHBvb2wgVGVzdCBDQTAeFw0yNjA3MjkxNjA4MTBa +Fw0zNjA3MjYxNjA4MTBaMB8xHTAbBgNVBAMMFGRlYWRwb29sLXRlc3QtY2xpZW50 +MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAzg6/rYvdSIkNYE0qEiro +mALB33Hy1ER04DKEk6g0aVBN9YvAFVNXXssB83X6dLVYRllBrZV0vfk7hRPH+0Ob +VospQ45+OaLYzrLDQWyeuVKvrdcWnuhqk9ntlxTi5LJw5IS251EX2a6ejg9uVHMb +9HhyqLelhZhkIAe4T4vgJ1sqgKqHx2BHpuZi2mgo8iNBHCKxgD7yxVNUBAQKGVJY +7PIA9myFYL4TsCiBPrpFgm3D4qA2f9GA15QlpogINtUglu+8oXeV8d14wfqCTBFB +FBygpdYi9vbGKYsMLfMC5qc49A3lZoxX6tXg43bnwOyvljRcH2iyy//cmRWX/k8s +QQIDAQABo3UwczAMBgNVHRMBAf8EAjAAMA4GA1UdDwEB/wQEAwIFoDATBgNVHSUE +DDAKBggrBgEFBQcDAjAdBgNVHQ4EFgQUOnE4r+8srsQhb/cRpTvW4oON694wHwYD +VR0jBBgwFoAUzTaHz/bbUU704LJiIEESUmerk+gwDQYJKoZIhvcNAQELBQADggEB +ABf/SIqTzHds33kGCzvw5qn0sFyaZHpXMiqT8qWO7g/R20oHb9i9DWB1wS/XCuTv +MLDCPwW2WRvP+HHPPs1usvh6Mr+/dtHyGz5w4EPjPeu5ARp3bZkWP77gMB9UK72X +sUc45OjDXAgJ6tF4AXqME8Hwt2574pjiQfCa0Uh5Skl7rKWmYoL+QlPBa2sO2U3Y +XxYvrDWwFFRC7vYFMm/dUrlPQsSV8NoatyPVKCg9omyNIPS+A8ZmjF76HoLjhiGd +TjoB7q6Tub/li+t/SdKD3l4I1MgW9O4u2CzBXybKWLmzI22m98aD7oAWWmgcJ4A0 +PVf44VxHOtTzSiLRhJYEeyw= +-----END CERTIFICATE----- diff --git a/tests/remote/certs/rogue-client-key.pem b/tests/remote/certs/rogue-client-key.pem new file mode 100644 index 0000000..ddf8569 --- /dev/null +++ b/tests/remote/certs/rogue-client-key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCtSiwtPwXPa/2X +Ors6WQYCDImL73z5tpYhlbHGTqm0r9ezGY0Hrs4rHdXzpIRVrTXSbJSaGV9bqZIL +O+5JfuhO36Yw2pjQZrk3v4lYGiXCyXJQ80cGepqRN1eEnjwe76rh8qCuOA78Sttj +uyqYxd6pwN26Amh0tMdBJWbE6CLFSGZDU8biiHqOaJwsiZxWAZGQ/NaY9NlbVkr/ +gd++ujsYFimm/XKDcRLczLAUa1ovSoRuk8Bv8by1I4OhzbKBKVluAPiCMdJihWHn +JZc7nrst747CrG3PUJHsV/0sYAXmsHaUt41i8E2sqyBUAQyfyqUp3o7kWAXhPnnN +fO8Bze5tAgMBAAECggEACPDVbAA0bInwgOem1qaBM4SHaZDnAdCfLlxq2lsfURNB +gxwZCqxvjAnzX/68/9lP+EC3P2bPugr7CK6u51vaNeI+0NvLxp+GLuTS++4IGg9g +C+KBZ57WcBJK5wIYWABFm3gekjPnkgd1X6U92A6yzk9adp5pL/Gj1i331MZQBPMB +gdQdCeqJvdwTL2vSPZRtaoEFkAPeOEMYU0U9xE8YNDqv5kUS7UdQZur2MpBVjcVv +ExW0VIZU09NvAvMgL4xgPApXU9PpZII2uvTFxMdHAYlXwKxpx4/ue/gECxdqkWDo +hsIYzo7c2Bfss9JORxvvZJjwRllWh64AWjgc4f3UgQKBgQDmK1lrlN8wdTjA+hdl +RximY6wmPYmNVDNwtARNYmL9vYDYm59hwDFkDj9tOY+7tlIKtDbiYpi4SUzUar5K +VCb49TJxK6iEFGIXSEKuo+imnwlEnoWGZnYAMKXGeEY4+RREwacpcZbgtQL/c4YJ +f9EjTu2hB5Dl96OcsnVPooq28wKBgQDAvLNgTEezQSAHjPglYzMwtYJZdvhaDLCJ +dase9muxZHWCO9lU0wdUFsbRLLxRFtOQ6AhRkWLOwZ+OBnGxdlrSgoEmkXIYxZZT +t9uV5qxOCThcw7xSHJxVEmv0P1jEL8UI3ZDTveRFrWco3vDCNdA7jmYP1ZboHPJH +4+C5qSbdHwKBgEmvtRfo+C72SC4XoqfMxAp4vGMdrkytmS5Ko1n21oQvR/GQmMzd +j7JdkVaxZ9+LdeZxXWTKdeQBq3QRnEwFdtia/wQWGwP11pVnj9mDJfc50OjD6zFk +2gAjkt7gIHMa4q9EY/Szpb7YlFYdsNqXreek3BSCbJQC3MFMrJvLCKebAoGAcRkK +BVMdq+FmYyfkoUPR0R5hrA/08hqKKU9kJ9ogHcs+bTqjcHQY0849wpcGtmq3oUuX +Pg4bNgpGj3fWlXVHHEo6cSBeHI0thljYQOFIcM5WRZESW/iv//e8Y2ocs5r4exhR +GP9QCVcUttYD32Lmm1wXJemHEROTVH4y65+Hz68CgYATNM70HN8fPJJsMmzXcC3Q +ygcgjScbQ+FMeKo7+7+kQnq6G+89zoupOgsXWsSeLYoyayQ5PEAvdGIKjjaVP//m +d11UmCBA4o7DzFfV6wxbCSdS/y4OJ0NiPN1UzLztcBKTy5RMP2nyhoSrl/OuWqon +ScLnD92kCIduVxgd0PNRcw== +-----END PRIVATE KEY----- diff --git a/tests/remote/certs/rogue-client.pem b/tests/remote/certs/rogue-client.pem new file mode 100644 index 0000000..2636637 --- /dev/null +++ b/tests/remote/certs/rogue-client.pem @@ -0,0 +1,20 @@ +-----BEGIN CERTIFICATE----- +MIIDRDCCAiygAwIBAgIUU/jnvr6LLYeFsFuxk6b7AQN+daIwDQYJKoZIhvcNAQEL +BQAwITEfMB0GA1UEAwwWRGVhZHBvb2wgUm9ndWUgVGVzdCBDQTAeFw0yNjA3Mjkx +NjA4MTBaFw0zNjA3MjYxNjA4MTBaMCAxHjAcBgNVBAMMFWRlYWRwb29sLXJvZ3Vl +LWNsaWVudDCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAK1KLC0/Bc9r +/Zc6uzpZBgIMiYvvfPm2liGVscZOqbSv17MZjQeuzisd1fOkhFWtNdJslJoZX1up +kgs77kl+6E7fpjDamNBmuTe/iVgaJcLJclDzRwZ6mpE3V4SePB7vquHyoK44DvxK +22O7KpjF3qnA3boCaHS0x0ElZsToIsVIZkNTxuKIeo5onCyJnFYBkZD81pj02VtW +Sv+B3766OxgWKab9coNxEtzMsBRrWi9KhG6TwG/xvLUjg6HNsoEpWW4A+IIx0mKF +Yecllzueuy3vjsKsbc9QkexX/SxgBeawdpS3jWLwTayrIFQBDJ/KpSnejuRYBeE+ +ec187wHN7m0CAwEAAaN1MHMwDAYDVR0TAQH/BAIwADAOBgNVHQ8BAf8EBAMCBaAw +EwYDVR0lBAwwCgYIKwYBBQUHAwIwHQYDVR0OBBYEFPB/KWrl0dLJT38ou9NZcY71 +otDcMB8GA1UdIwQYMBaAFNN9G6qoEe/VAoR5Gdi9GKKhGRzHMA0GCSqGSIb3DQEB +CwUAA4IBAQBM34K7/JDy4HlIWqEtWZ1xAQPPpm/izcxqdt4OfOzLSQkkK16+ND5d +/yN5qqsvLrgK6TrXMacuFrEQ879h+gaUW0mOwCyZXSyRuqOKIPa0XflnfpQYqfab +bPFLdgI/tdsIAZSQNeQ48a8rE14+C+GdBqDzBpMqZdNZcTkT2JkMGc3+L2sjckMK +AgaPoK8HNEBxpZRAl3spHNQyBGV/qMvwNakRv0TdGuVO+Q2EnaPUg/gZPPvthudG +ey94jmAojkDuv7Y7gzv+vLgaVmzYleM2KS5acBBFu1Grb2WBEOfY8gX+JZtCipSR +vUXTWILHTltuzL91ccNUzjZG3jjWjtZr +-----END CERTIFICATE----- diff --git a/tests/remote/certs/server-key.pem b/tests/remote/certs/server-key.pem new file mode 100644 index 0000000..7113a2e --- /dev/null +++ b/tests/remote/certs/server-key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCX0MCIw+riNKEv +wqs4xAWt5P1skhiwxNqZsm5b/3ypEGJgs0UZOzawnFaGfAWUc80Pt/mP1rIx/ktJ +icCsV6rGNKdFlMtPCu0WDFWVHL3hv9bt8KY+A447+4KithuFPUx4hfOy8/RQYgHl +TU/jybqZBuWlUQTtrtUA0mTLC4i06SkBF7lq8QC/+LqkHL9LZmw6hcigRzFmADrZ +sRGW4DWszq/0mcAcFZXrceby6WPIpvWJB546ckCgjmZPVJ611iFpGpIalTREViW4 +mlGiC1myoy098iyoffn2Kzhsa0bQhUAi7Nsy2i+9zrpoavtjiHu4OSp8U/W0g6tu ++V4pl0etAgMBAAECggEAIXfz+jK25YeUwWK5PJp6ZgvFktnid5XuW1Pq4H1CjkL+ +bTMYWKelNupAaQWOdVvSKVT1v7AlOMBG+L9VqqjpDSkudziGWbae7e1qMKVVdx9D +0tDq1WETSkX4Vsau6oDJCwj84MbHZbE1mcehhSZD7NiIaBmkbIHOk4/OiLTSPo+R +ecCO9YQgcdq9dVadIkG6c5UVCWMwzPCAhYsvE8Qqdn8we2DKhtwJPnp6zwRJBVHf +Smc4hdrEzTcD1yU7s06oMYNu3mru12UVFDt0PhNx8IS/e7JAc2mBBqgIQd7+PJID +JDmT+3GjElIZSYWE7IecrqSszMaDlP1pRYKs0BxKAQKBgQDJJ1ZeF5vBRKFgDhjN +5Ct1bGN+FV6K3fINadHJoQSjcROGjgQplw8hMdWMsSoITic0t15csYraPmb6KEDj +JfcZoQhzQhgcBslPkMMFt1cKGuwvQbUmVVjN58I2jVyN3p0msaquWhmacT6+xM5i +617zOQnQ5bI6/RxU0EQTP+dhlQKBgQDBNZK1Xubcbzxu/x8deNPUH8pA7kV1M0QW +Xc3n32deKfutCmyoz8C2mBeCoQT0rnhM7wecBd5sMgiAmOIR4MS+ornyBRuek/Ks +4hc3npccNYFpI+XGb+LBLJWJNWV4q0+U8Ap2wEWSHEaiahUcQqRoolCZojJOyGVi +eaQ90o/3uQKBgBGPPoneM7mOXSv+bwniJi0M9i52fRdQQqYLG8YnKTl/UH4JfbGH +v0ldiU/L02iAMgn3C+S5lu5wThr+UpPBp610EkHFfkdMDFzGvU+NbqqyKPYDHWYo +QyVOWvTaKD82NK/BCK5JrTx/MEQE+CNIerRuwT2cH4/osoVw3NB1XqR1AoGAIiLP +/+F7knjIyaejnLZrXrAbMOuoM/PR1M3QAmazQazTPZh46D6Egv5OMRM35+8nbBhT +VNqomJ8iZ24biyMWP8RKbBtA7PygxLDRf08yzZYafjQMdcnw20aASRS5D9/cmKjs +gFyqLg124V0yy6Jun9oFuY8xqj+3wfVR8mZQb2kCgYEAq6IF4kNML176iSoaaCAP +qpMV3w7/ykviie9VSszrOpriGQOIjLGmj1TQm6kwrnmm69yY1Tqq8bWJC05K2edG +xeUtbksH/FgcvVVC1u4+e9Fc35XRNAMxTtM2elrsruton8ArWvnYlKa4RkRL/e3c +EU6xnNsle5bWY4J55Idr8Ow= +-----END PRIVATE KEY----- diff --git a/tests/remote/certs/server.pem b/tests/remote/certs/server.pem new file mode 100644 index 0000000..e816722 --- /dev/null +++ b/tests/remote/certs/server.pem @@ -0,0 +1,20 @@ +-----BEGIN CERTIFICATE----- +MIIDUDCCAjigAwIBAgIUa1KWjufkag+0cE9E5h90qMk3jqAwDQYJKoZIhvcNAQEL +BQAwGzEZMBcGA1UEAwwQRGVhZHBvb2wgVGVzdCBDQTAeFw0yNjA3MjkxNjA4MTBa +Fw0zNjA3MjYxNjA4MTBaMBQxEjAQBgNVBAMMCWxvY2FsaG9zdDCCASIwDQYJKoZI +hvcNAQEBBQADggEPADCCAQoCggEBAJfQwIjD6uI0oS/CqzjEBa3k/WySGLDE2pmy +blv/fKkQYmCzRRk7NrCcVoZ8BZRzzQ+3+Y/WsjH+S0mJwKxXqsY0p0WUy08K7RYM +VZUcveG/1u3wpj4Djjv7gqK2G4U9THiF87Lz9FBiAeVNT+PJupkG5aVRBO2u1QDS +ZMsLiLTpKQEXuWrxAL/4uqQcv0tmbDqFyKBHMWYAOtmxEZbgNazOr/SZwBwVletx +5vLpY8im9YkHnjpyQKCOZk9UnrXWIWkakhqVNERWJbiaUaILWbKjLT3yLKh9+fYr +OGxrRtCFQCLs2zLaL73Oumhq+2OIe7g5KnxT9bSDq275XimXR60CAwEAAaOBkjCB +jzAMBgNVHRMBAf8EAjAAMA4GA1UdDwEB/wQEAwIFoDATBgNVHSUEDDAKBggrBgEF +BQcDATAaBgNVHREEEzARgglsb2NhbGhvc3SHBH8AAAEwHQYDVR0OBBYEFKATDO9g +5j+dGts7cFRDLQi6Hc60MB8GA1UdIwQYMBaAFM02h8/221FO9OCyYiBBElJnq5Po +MA0GCSqGSIb3DQEBCwUAA4IBAQCUIG+nZ8X6/aevChi560gPhEfvP3LHW8RV1jVO +mmWFET+pzcubqwkAQEtrIcEuNfY6bju1igI+Uz9K4a/gCgfEcb6VjOzoRgXftem8 +WNYbsu8QVaUq20z3z9EckYOHfg8OXu6IcGGnRhYmPEWaH0Wx0hF1i2C4arX2EV+C +oQ4GeQU5JrGTBv2UA8n4vBELQVvPXx6NHL4TdsS0iFWsCVh1v43y0aMTRKHXmiVv +ueJl1EBiF8EhqkJeZ8E4LWD6Ywi74nnnSF61QpvNCXbNwkInHqliS3/CXVsVnASi +pTzoXH8NGp0foI9aPnn82MJznWMNv/HoQvFVMmvU+q44AZMz +-----END CERTIFICATE----- diff --git a/tests/remote/conftest.py b/tests/remote/conftest.py new file mode 100644 index 0000000..99e61f1 --- /dev/null +++ b/tests/remote/conftest.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from functools import partial +from pathlib import Path +import ssl + +import pytest + +import deadpool +from deadpool.remote import DeadpoolClient, DeadpoolServer, UnixAddress, UnixListener +from tests.remote.tasks import multiply + +_CERTS = Path(__file__).with_name("certs") + + +@dataclass(frozen=True, slots=True) +class TLSContexts: + """Mutually authenticated contexts backed by static test-only credentials.""" + + server: ssl.SSLContext + client: ssl.SSLContext + missing_client_certificate: ssl.SSLContext + untrusted_client: ssl.SSLContext + + +def _client_tls_context(*, certificate: str | None = "client") -> ssl.SSLContext: + context = ssl.create_default_context( + ssl.Purpose.SERVER_AUTH, + cafile=_CERTS / "ca.pem", + ) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + if certificate is not None: + context.load_cert_chain( + _CERTS / f"{certificate}.pem", + _CERTS / f"{certificate}-key.pem", + ) + return context + + +@pytest.fixture() +def tls_contexts() -> TLSContexts: + server = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server.minimum_version = ssl.TLSVersion.TLSv1_2 + server.verify_mode = ssl.CERT_REQUIRED + server.load_verify_locations(_CERTS / "ca.pem") + server.load_cert_chain(_CERTS / "server.pem", _CERTS / "server-key.pem") + return TLSContexts( + server=server, + client=_client_tls_context(), + missing_client_certificate=_client_tls_context(certificate=None), + untrusted_client=_client_tls_context(certificate="rogue-client"), + ) + + +@pytest.fixture() +def remote_pair(tmp_path): + socket_path = tmp_path / "deadpool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=2, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"multiply": multiply}, + ) + server.start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + yield server, client + finally: + client.shutdown(wait=True, cancel_futures=True) + server.shutdown(wait=True, cancel_futures=True, deadline=5) diff --git a/tests/remote/tasks.py b/tests/remote/tasks.py new file mode 100644 index 0000000..d55056d --- /dev/null +++ b/tests/remote/tasks.py @@ -0,0 +1,48 @@ +"""Importable registered tasks used by forkserver worker tests.""" + +import os +import time +from pathlib import Path + + +def multiply(left: int, right: int) -> int: + return left * right + + +def delayed(value: object, delay: float) -> object: + time.sleep(delay) + return value + + +def make_bytes(size: int) -> bytes: + return b"x" * size + + +def echo_bytes(value: bytes) -> bytes: + return value + + +def exit_abruptly(code: int = 17) -> None: + """Terminate the current worker without Python-level cleanup.""" + os._exit(code) + + +def mark(path: str | Path, value: str = "ran") -> str: + Path(path).write_text(value) + return value + + +def wait_then_mark( + release_path: str | Path, marker_path: str | Path, value: str +) -> str: + release = Path(release_path) + while not release.exists(): + time.sleep(0.005) + Path(marker_path).write_text(value) + return value + + +def append_line(path: str | Path, value: str) -> str: + with Path(path).open("a") as stream: + stream.write(f"{value}\n") + return value diff --git a/tests/remote/test_cancellation.py b/tests/remote/test_cancellation.py new file mode 100644 index 0000000..a45d3bf --- /dev/null +++ b/tests/remote/test_cancellation.py @@ -0,0 +1,357 @@ +import threading +import time +from collections.abc import Callable +from concurrent.futures import CancelledError, as_completed, wait +from functools import partial + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteExecutorError, + RemoteQueueTimeout, + ServerState, + SubmissionState, + UnixAddress, + UnixListener, +) +from tests.remote.tasks import mark, multiply, wait_then_mark + + +def marker(path): + path.write_text("ran") + + +def wait_until(predicate: Callable[[], bool], timeout: float = 3.0) -> None: + """Poll an observable state transition under a bounded deadline.""" + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError("condition did not become true before deadline") + time.sleep(0.005) + + +def assert_completes(call: Callable[[], object], timeout: float = 1.0) -> None: + """Require a lifecycle call to return without relying on elapsed timing.""" + completed = threading.Event() + errors: list[BaseException] = [] + + def invoke() -> None: + try: + call() + except BaseException as error: + errors.append(error) + finally: + completed.set() + + threading.Thread(target=invoke, daemon=True).start() + assert completed.wait(timeout), "call did not complete before the deadline" + if errors: + raise errors[0] + + +def test_queued_cancellation_prevents_execution(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + queued_marker = tmp_path / "queued" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + running = client.submit(wait_then_mark, release, running_marker, "done") + wait_until(running.running) + queued = client.submit(marker, queued_marker) + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) + assert queued.cancel() + assert queued.cancelled() + done, not_done = wait([queued], timeout=0.1) + assert done == {queued} + assert not not_done + assert list(as_completed([queued], timeout=0.1)) == [queued] + release.touch() + assert running.result(timeout=5) == "done" + assert running_marker.read_text() == "done" + assert not queued_marker.exists() + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_group_cancellation_is_scoped_and_reports_counts(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + queued_marker = tmp_path / "queued" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + running = client.submit( + wait_then_mark, + release, + running_marker, + "done", + deadpool_group_id="batch", + ) + wait_until(running.running) + queued = client.submit(mark, queued_marker, "late", deadpool_group_id="batch") + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) + assert client.cancel_group("batch") == { + "cancelled": 1, + "running": 1, + "terminal": 0, + "unknown": 0, + } + wait_until(queued.done) + assert queued.cancelled() + release.touch() + assert running.result(timeout=5) == "done" + assert not queued_marker.exists() + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_unknown_and_hard_running_group_cancellation(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + queued_marker = tmp_path / "queued" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + assert client.cancel_group("missing") == { + "cancelled": 0, + "running": 0, + "terminal": 0, + "unknown": 1, + } + + running = client.submit( + wait_then_mark, + release, + running_marker, + "running", + deadpool_group_id="batch", + ) + wait_until(running.running) + queued = client.submit(mark, queued_marker, "queued", deadpool_group_id="batch") + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) + + assert client.cancel_group("batch", hard=True) == { + "cancelled": 2, + "running": 0, + "terminal": 0, + "unknown": 0, + } + with pytest.raises(CancelledError): + running.result(2) + with pytest.raises(CancelledError): + queued.result(2) + assert not running_marker.exists() + assert not queued_marker.exists() + assert client.submit(multiply, 6, 7).result(5) == 42 + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_queue_timeout_is_distinct_from_execution_timeout(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + running = client.submit(wait_then_mark, release, running_marker, "done") + wait_until(running.running) + queued = client.submit( + mark, + tmp_path / "queued", + "late", + deadpool_queue_timeout=0.02, + ) + with pytest.raises(RemoteQueueTimeout): + queued.result(timeout=5) + release.touch() + assert running.result(timeout=5) == "done" + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_queue_timeout_survives_saturated_local_backlog(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + expired_marker = tmp_path / "expired" + server = DeadpoolServer( + partial( + deadpool.Deadpool, + max_workers=1, + max_backlog=1, + mp_context="forkserver", + ), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + pool = server._pool + assert pool is not None + try: + running = pool.submit( + wait_then_mark, release, running_marker, "unrelated-running" + ) + wait_until(lambda: running.pid is not None) + backlogged = pool.submit(multiply, 6, 7) + wait_until(pool.submitted_jobs.full) + + expiring = client.submit( + mark, + expired_marker, + "must-not-run", + deadpool_queue_timeout=0.02, + ) + with pytest.raises(RemoteQueueTimeout): + expiring.result(timeout=2) + + assert not running.done() + assert not expired_marker.exists() + assert pool.get_statistics()["tasks_received"] == 2 + wait_until(lambda: server.get_statistics()["remote_staged"] == 0) + + release.touch() + assert running.result(timeout=5) == "unrelated-running" + assert backlogged.result(timeout=5) == 42 + assert client.submit(multiply, 3, 4).result(timeout=5) == 12 + assert not expired_marker.exists() + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_ordinary_cancel_does_not_kill_running_task(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + marker_path = tmp_path / "marker" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + future = client.submit(wait_then_mark, release, marker_path, "done") + wait_until(future.running) + assert not future.cancel() + assert not future.done() + release.touch() + assert future.result(timeout=5) == "done" + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_hard_cancel_running_future_is_terminal_and_pool_recovers(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + marker_path = tmp_path / "marker" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + future = client.submit(wait_then_mark, release, marker_path, "never") + wait_until(future.running) + assert future.cancel_and_kill_if_running() + assert future.done() + assert future.cancelled() + with pytest.raises(CancelledError): + future.result() + assert not marker_path.exists() + assert client.submit(multiply, 6, 7).result(timeout=5) == 42 + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_server_shutdown_deadline_cancels_running_and_queued_work(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + running_marker = tmp_path / "running" + queued_marker = tmp_path / "queued" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + running = client.submit(wait_then_mark, release, running_marker, "running") + wait_until(running.running) + + queued = client.submit(mark, queued_marker, "queued") + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) + + server.shutdown(wait=False, cancel_futures=False, deadline=5) + server.shutdown(wait=False, cancel_futures=True, deadline=0.1) + wait_until(lambda: server.state is ServerState.STOPPED, timeout=5) + + with pytest.raises(CancelledError): + running.result(2) + with pytest.raises(CancelledError): + queued.result(2) + assert not running_marker.exists() + assert not queued_marker.exists() + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=0) + + +def test_nonblocking_server_shutdown_gracefully_drains_accepted_work(tmp_path): + socket_path = tmp_path / "pool.sock" + release = tmp_path / "release" + marker_path = tmp_path / "marker" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + accepted = client.submit(wait_then_mark, release, marker_path, "drained") + wait_until(accepted.running) + + assert_completes(lambda: server.shutdown(wait=False, deadline=5)) + assert server.state is ServerState.DRAINING + assert not accepted.done() + with pytest.raises(RemoteExecutorError, match="server_draining"): + client.submit(multiply, 2, 3).result(5) + + release.touch() + assert accepted.result(5) == "drained" + wait_until(lambda: server.state is ServerState.STOPPED, timeout=5) + assert marker_path.read_text() == "drained" + assert not socket_path.exists() + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=0) diff --git a/tests/remote/test_client_api.py b/tests/remote/test_client_api.py new file mode 100644 index 0000000..5b5b9f5 --- /dev/null +++ b/tests/remote/test_client_api.py @@ -0,0 +1,493 @@ +import concurrent.futures +import logging +import os +import queue +import threading +import time +from types import SimpleNamespace +from functools import partial +from pathlib import Path +from typing import Callable + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteConnectionLost, + RemoteExecutorError, + RemoteExecutorUnavailable, + RemoteLimits, + RemoteProcessError, + RemoteProtocolError, + RemoteQueueFull, + RemoteResultTooLarge, + ServerState, + SubmissionState, + UnixAddress, + UnixListener, +) +from deadpool.remote._protocol import Message, MessageType +from deadpool.remote.serializer import PickleSerializer +from tests.remote.tasks import ( + delayed, + exit_abruptly, + make_bytes, + multiply, + wait_then_mark, +) + + +def wait_until(predicate: Callable[[], bool], timeout: float = 3.0) -> None: + """Poll an observable condition under a bounded monotonic deadline.""" + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError("condition did not become true before deadline") + time.sleep(0.005) + + +def assert_completes(call: Callable[[], object], timeout: float = 1.0) -> None: + """Require a lifecycle call to return without relying on elapsed timing.""" + completed = threading.Event() + errors: list[BaseException] = [] + + def invoke() -> None: + try: + call() + except BaseException as error: + errors.append(error) + finally: + completed.set() + + threading.Thread(target=invoke, daemon=True).start() + assert completed.wait(timeout), "call did not complete before the deadline" + if errors: + raise errors[0] + + +def make_pair( + socket_path: Path, *, limits: RemoteLimits | None = None +) -> tuple[DeadpoolServer, DeadpoolClient]: + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"multiply": multiply}, + limits=limits, + ).start() + return server, DeadpoolClient(UnixAddress(socket_path), limits=limits) + + +class BlockingSerializer(PickleSerializer): + """Expose invocation serialization as a deterministic lifecycle barrier.""" + + def __init__(self) -> None: + super().__init__() + self.entered = threading.Event() + self.release = threading.Event() + + def dumps(self, value: object, *, limit: int) -> bytes: + self.entered.set() + if not self.release.wait(5): + raise TimeoutError("serializer test barrier was not released") + return super().dumps(value, limit=limit) + + +class RejectingResultSerializer(PickleSerializer): + """Model a result whose Python type is unavailable on the client.""" + + def loads(self, payload: bytes) -> object: + value = super().loads(payload) + if value == "unavailable-client-type": + raise ModuleNotFoundError("client cannot import result type") + return value + + +def test_result_decode_failure_acknowledges_and_forgets_request(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + ).start() + client = DeadpoolClient(UnixAddress(path), serializer=RejectingResultSerializer()) + try: + future = client.submit(delayed, "unavailable-client-type", 0) + with pytest.raises(RemoteProtocolError, match="cannot import result type"): + future.result(5) + + wait_until(lambda: server.get_statistics()["remote_retained_outcomes"] == 0) + assert future.request_id not in client._futures + assert future.request_id not in client._terminal_received + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_result_ack_enqueue_failure_closes_session_and_forgets_request( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + server, client = make_pair(tmp_path / "pool.sock") + release = tmp_path / "release" + marker = tmp_path / "marker" + future = client.submit(wait_then_mark, release, marker, "done") + original_enqueue = client._enqueue_control + + def reject_result_ack(kind: MessageType, control: dict) -> None: + if kind == MessageType.RESULT_ACK: + raise RemoteConnectionLost("client outbound control queue is full") + original_enqueue(kind, control) + + try: + wait_until(future.running) + monkeypatch.setattr(client, "_enqueue_control", reject_result_ack) + release.touch() + assert future.result(5) == "done" + + wait_until(lambda: client._socket is None) + wait_until(lambda: server.get_statistics()["remote_retained_outcomes"] == 0) + assert future.request_id not in client._futures + assert future.request_id not in client._terminal_received + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_submit_rechecks_shutdown_after_serialization(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + serializer = BlockingSerializer() + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + ).start() + client = DeadpoolClient(UnixAddress(path), serializer=serializer) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as calls: + submission = calls.submit(client.submit, multiply, 2, 3) + try: + assert serializer.entered.wait(2) + client.shutdown(wait=False) + serializer.release.set() + with pytest.raises(RuntimeError, match="after shutdown"): + submission.result(2) + finally: + serializer.release.set() + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_client_server_contexts_and_repeated_shutdown(tmp_path: Path) -> None: + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"multiply": multiply}, + ) + + with server: + assert server.state is ServerState.RUNNING + with DeadpoolClient(UnixAddress(socket_path)) as client: + assert client.submit_many([(multiply, (2, 3), {})])[0].result(5) == 6 + client.shutdown() + with pytest.raises(RuntimeError, match="after shutdown"): + client.submit(multiply, 1, 2) + + server.shutdown() + assert server.state is ServerState.STOPPED + assert not socket_path.exists() + + +def test_known_server_loss_rejects_submit_and_control_rpc(tmp_path: Path) -> None: + server, client = make_pair(tmp_path / "pool.sock") + try: + assert client.check_health() + server.shutdown(cancel_futures=True, deadline=0) + + def transport_is_lost() -> bool: + try: + client.check_health(timeout=0.05) + except RemoteConnectionLost: + return True + except TimeoutError: + return False + return False + + wait_until(transport_is_lost) + with pytest.raises(RemoteConnectionLost): + client.get_status("never-existed") + with pytest.raises(RemoteExecutorUnavailable, match="connection"): + client.submit(multiply, 2, 3) + + assert_completes(client.shutdown) + assert_completes(client.shutdown) + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=0) + + +def test_rpc_response_claims_token_before_following_disconnect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A complete response wins even if transport loss wakes immediately after it.""" + + class PausedWake: + def __init__(self) -> None: + self.signaled = threading.Event() + self.waiter_reached = threading.Event() + self.release_waiter = threading.Event() + + def set(self) -> None: + self.signaled.set() + + def wait(self, timeout: float | None = None) -> bool: + if not self.signaled.wait(timeout): + return False + self.waiter_reached.set() + return self.release_waiter.wait(timeout) + + class SocketStub: + def shutdown(self, how: int) -> None: + pass + + def close(self) -> None: + pass + + client = object.__new__(DeadpoolClient) + client._owner_pid = os.getpid() + client._lock = threading.RLock() + client._socket = SocketStub() + client._rpc = {} + client._transport_failed = False + client._futures = {} + client._terminal_received = set() + client.control_timeout = 2.0 + request_sent = threading.Event() + sent_control: dict[str, object] = {} + + def capture_request(kind: MessageType, control: dict) -> None: + sent_control.update(control) + request_sent.set() + + client._enqueue_control = capture_request + paused_wake = PausedWake() + client_module = __import__("deadpool.remote.client", fromlist=["client"]) + real_threading = client_module.threading + monkeypatch.setattr( + client_module, + "threading", + SimpleNamespace(Event=lambda: paused_wake), + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as calls: + rpc = calls.submit(client.check_health) + assert request_sent.wait(1) + token = sent_control["nonce"] + client._receive(Message(MessageType.PONG, {"nonce": token})) + assert paused_wake.waiter_reached.wait(1) + client._connection_lost(OSError("disconnect after response")) + paused_wake.release_waiter.set() + assert rpc.result(1) is True + + monkeypatch.setattr(client_module, "threading", real_threading) + assert client._rpc == {} + + +def test_connection_loss_during_rpc_is_not_a_normal_error_response( + tmp_path: Path, +) -> None: + path = tmp_path / "pool.sock" + entered = threading.Event() + release = threading.Event() + + def blocking_authorizer(principal: object, operation: str, metadata: dict) -> bool: + if operation == "statistics": + entered.set() + release.wait(5) + return True + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + authorizer=blocking_authorizer, + ).start() + client = DeadpoolClient(UnixAddress(path), control_timeout=1) + executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + rpc = executor.submit(client.get_statistics) + try: + assert entered.wait(2) + server.shutdown(cancel_futures=True, deadline=0) + with pytest.raises(RemoteConnectionLost): + rpc.result(2) + finally: + release.set() + executor.shutdown() + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=0) + + +def test_control_enqueue_failure_releases_rpc_tokens( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + server, client = make_pair(tmp_path / "pool.sock") + release = tmp_path / "release" + marker = tmp_path / "marker" + running = client.submit(wait_then_mark, release, marker, "done") + original_put = client._outbound.put + + def raise_full(*args: object, **kwargs: object) -> None: + raise queue.Full + + try: + wait_until(running.running) + monkeypatch.setattr(client._outbound, "put", raise_full) + with pytest.raises(RemoteConnectionLost, match="control queue is full"): + client.check_health() + assert client._rpc == {} + + with pytest.raises(RemoteConnectionLost, match="control queue is full"): + running.cancel() + assert client._rpc == {} + finally: + monkeypatch.setattr(client._outbound, "put", original_put) + release.touch(exist_ok=True) + running.result(5) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_worker_abrupt_death_is_typed_and_pool_recovers(tmp_path: Path) -> None: + server, client = make_pair(tmp_path / "pool.sock") + try: + with pytest.raises(RemoteProcessError): + client.submit(exit_abruptly).result(5) + assert client.submit(multiply, 6, 7).result(5) == 42 + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_unknown_task_and_oversized_result_do_not_poison_session( + tmp_path: Path, +) -> None: + limits = RemoteLimits(max_result_bytes=256) + server, client = make_pair(tmp_path / "pool.sock", limits=limits) + try: + with pytest.raises(RemoteExecutorError, match="unknown_operation"): + client.submit_task("missing.operation").result(5) + with pytest.raises(RemoteResultTooLarge): + client.submit(make_bytes, 2048).result(5) + assert client.submit_task("multiply", 6, 7).result(5) == 42 + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_status_transitions_cover_queued_running_terminal_and_unknown( + tmp_path: Path, +) -> None: + server, client = make_pair(tmp_path / "pool.sock") + release = tmp_path / "release" + marker = tmp_path / "marker" + try: + running = client.submit(wait_then_mark, release, marker, "first") + wait_until(running.running) + assert running.pid is not None + assert client.get_status(running) == "RUNNING" + + queued = client.submit(delayed, "second", 0.01) + wait_until( + lambda: client.get_status(queued) == "ACCEPTED_QUEUED" + and queued.submission_state is SubmissionState.ACCEPTED_QUEUED + ) + assert client.get_status("unknown-request") == "UNKNOWN" + + release.touch() + assert running.result(5) == "first" + assert queued.result(5) == "second" + wait_until(lambda: client.get_status(queued) == "UNKNOWN") + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_submission_rejects_empty_group_id_without_reaching_server( + tmp_path: Path, +) -> None: + server, client = make_pair(tmp_path / "pool.sock") + try: + before = server.get_statistics()["remote_tasks_received"] + with pytest.raises(ValueError, match="group_id must be non-empty"): + client.submit(multiply, 2, 3, deadpool_group_id="") + assert server.get_statistics()["remote_tasks_received"] == before + assert client.submit(multiply, 3, 4).result(5) == 12 + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_per_session_admission_recovers_after_completion(tmp_path: Path) -> None: + limits = RemoteLimits(max_pending_per_session=1) + server, client = make_pair(tmp_path / "pool.sock", limits=limits) + release = tmp_path / "release" + marker = tmp_path / "marker" + try: + first = client.submit(wait_then_mark, release, marker, "first") + wait_until(first.running) + with pytest.raises(RemoteQueueFull): + client.submit(delayed, "rejected", 0).result(5) + + release.touch() + assert first.result(5) == "first" + wait_until(lambda: server.get_statistics()["remote_retained_outcomes"] == 0) + assert client.submit(delayed, "recovered", 0).result(5) == "recovered" + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_public_future_exception_pid_and_callback_isolation( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + server, client = make_pair(tmp_path / "pool.sock") + release = tmp_path / "release" + marker = tmp_path / "marker" + callbacks: list[str] = [] + callbacks_done = threading.Event() + caplog.set_level(logging.ERROR, logger="deadpool.remote") + try: + running = client.submit(wait_then_mark, release, marker, "done") + wait_until(running.running) + assert running.pid is not None + release.touch() + assert running.result(5) == "done" + + failed = client.submit(int, "not-an-int") + error = failed.exception(5) + assert isinstance(error, ValueError) + + def broken_callback(future: object) -> None: + callbacks.append("broken") + raise RuntimeError("callback exploded") + + def healthy_callback(future: object) -> None: + callbacks.append("healthy") + callbacks_done.set() + + callback_future = client.submit(multiply, 6, 7) + callback_future.add_done_callback(broken_callback) + callback_future.add_done_callback(healthy_callback) + assert callback_future.result(5) == 42 + assert callbacks_done.wait(2) + callback_future.add_done_callback(lambda future: callbacks.append("late")) + wait_until(lambda: callbacks == ["broken", "healthy", "late"]) + assert "callback exploded" in caplog.text + assert client.submit(multiply, 3, 4).result(5) == 12 + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) diff --git a/tests/remote/test_fork_timeout.py b/tests/remote/test_fork_timeout.py new file mode 100644 index 0000000..f6c97c4 --- /dev/null +++ b/tests/remote/test_fork_timeout.py @@ -0,0 +1,45 @@ +import os +import time + +import pytest + +import deadpool +from deadpool.remote import RemoteForkedProcessError + + +def delayed(delay): + time.sleep(delay) + return delay + + +def test_execution_timeout_retires_worker_and_pool_continues(remote_pair): + _, client = remote_pair + with pytest.raises(deadpool.TimeoutError): + client.submit(delayed, 0.3, deadpool_timeout=0.05).result(timeout=5) + assert client.submit(delayed, 0.01).result(timeout=5) == 0.01 + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="requires fork") +def test_inherited_future_is_invalid_in_child(remote_pair): + _, client = remote_pair + future = client.submit(delayed, 0.15) + read_fd, write_fd = os.pipe() + pid = os.fork() + if pid == 0: + os.close(read_fd) + try: + future.done() + except RemoteForkedProcessError: + result = client.submit(delayed, 0.01).result(timeout=5) + os.write(write_fd, f"invalid:{result}".encode()) + else: + os.write(write_fd, b"shared") + finally: + os.close(write_fd) + os._exit(0) + os.close(write_fd) + result = os.read(read_fd, 16) + os.close(read_fd) + os.waitpid(pid, 0) + assert result == b"invalid:0.01" + assert future.result(timeout=5) == 0.15 diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py new file mode 100644 index 0000000..7989366 --- /dev/null +++ b/tests/remote/test_integration.py @@ -0,0 +1,578 @@ +import importlib +import socket +import threading +import time +from types import SimpleNamespace +from functools import partial + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteAuthenticationError, + RemoteCompatibilityError, + RemoteLimits, + RemoteProtocolError, + RemoteQueueFull, + RemoteSubmissionTimeout, + RemoteResultEncodingError, + RemoteTaskError, + UnixAddress, + UnixListener, +) +from deadpool.remote._protocol import ( + Message, + MessageReader, + MessageType, + _wire_limits, + send_message, +) +from deadpool.remote.serializer import PickleSerializer +from tests.remote.tasks import wait_then_mark + + +def add(left, right): + return left + right + + +def delayed(value, delay): + time.sleep(delay) + return value + + +def fail_value_error(): + raise ValueError("remote boom") + + +def fail_large_value_error(): + raise ValueError("task exploded: " + "\x00" * 20_000) + + +def large_encoding_failure_result(): + return "trigger-large-encoding-failure" + + +def unencodable_result(): + return lambda: None + + +class UnpicklableError(Exception): + def __init__(self): + super().__init__("cannot pickle me") + self.callback = lambda: None + + +def fail_unpicklable(): + raise UnpicklableError() + + +def test_callable_registered_stats_and_health(remote_pair): + server, client = remote_pair + + assert client.check_health() + assert client.submit(add, 2, 3).result(timeout=5) == 5 + assert client.submit_task("multiply", 6, 7).result(timeout=5) == 42 + stats = client.get_statistics() + assert stats["remote_tasks_terminal"] == 2 + assert stats["tasks_received"] == 2 + assert server.state == "RUNNING" + + +def test_independent_clients_share_server_without_owning_it(remote_pair): + server, first = remote_pair + second = DeadpoolClient(UnixAddress(server.bound_addresses[0])) + try: + assert first.submit(add, 1, 2).result(timeout=5) == 3 + assert second.submit(add, 3, 4).result(timeout=5) == 7 + second.shutdown() + assert first.submit(add, 5, 6).result(timeout=5) == 11 + assert server.state == "RUNNING" + finally: + second.shutdown() + + +def test_authentication_rejection_is_typed(tmp_path): + socket_path = tmp_path / "pool.sock" + + def reject(peer, hello): + raise PermissionError("no") + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + authenticator=reject, + ).start() + try: + with pytest.raises(RemoteAuthenticationError): + DeadpoolClient(UnixAddress(socket_path)) + finally: + server.shutdown(cancel_futures=True, deadline=5) + + +def test_duplicate_live_client_instance_is_rejected(remote_pair, monkeypatch): + server, first = remote_pair + client_module = importlib.import_module("deadpool.remote.client") + monkeypatch.setattr( + client_module.uuid, + "uuid4", + lambda: SimpleNamespace(hex=first._client_id), + ) + with pytest.raises(RemoteCompatibilityError): + DeadpoolClient(UnixAddress(server.bound_addresses[0])) + + +@pytest.mark.parametrize( + "server_limit_changes", + [ + {"max_control_bytes": 32 * 1024}, + {"max_frame_payload_bytes": 512 * 1024}, + {"max_message_bytes": 128 * 1024 * 1024}, + {"max_metadata_bytes": 8 * 1024}, + {"max_chunks": 128}, + ], +) +def test_handshake_rejects_asymmetric_wire_limits(tmp_path, server_limit_changes): + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + limits=RemoteLimits(**server_limit_changes), + ).start() + try: + with pytest.raises(RemoteCompatibilityError, match="wire_limits"): + DeadpoolClient(UnixAddress(socket_path)) + finally: + server.shutdown(cancel_futures=True, deadline=5) + + +def test_client_rejects_server_selected_asymmetric_wire_limits(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits() + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + listener.bind(str(socket_path)) + listener.listen() + listener.settimeout(5) + server_errors = [] + + def serve_incompatible_welcome() -> None: + try: + connection, _ = listener.accept() + with connection: + MessageReader(limits).receive(connection) + selected = _wire_limits(limits) + selected["max_chunks"] -= 1 + send_message( + connection, + Message( + MessageType.WELCOME, + { + "wire": "experimental-deadpool-private-v1", + "wire_limits": selected, + }, + ), + limits, + ) + except BaseException as error: + server_errors.append(error) + + server_thread = threading.Thread(target=serve_incompatible_welcome) + server_thread.start() + try: + with pytest.raises(RemoteCompatibilityError, match="wire limits"): + DeadpoolClient(UnixAddress(socket_path)) + finally: + listener.close() + server_thread.join(5) + socket_path.unlink(missing_ok=True) + + assert not server_thread.is_alive() + assert server_errors == [] + + +def test_callable_mode_never_pickles_registered_task_registry(tmp_path): + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"deliberately.unpicklable": lambda: None}, + ).start() + client = DeadpoolClient(UnixAddress(socket_path)) + try: + assert client.submit(add, 20, 22).result(timeout=5) == 42 + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_out_of_order_correlation_and_ordered_map(remote_pair): + _, client = remote_pair + slow = client.submit(delayed, "slow", 0.2) + fast = client.submit(delayed, "fast", 0.01) + + assert fast.result(timeout=5) == "fast" + assert slow.result(timeout=5) == "slow" + assert list(client.map(add, [1, 2, 3], [10, 20, 30])) == [11, 22, 33] + + +def test_invalid_control_metadata_fails_before_transport_use(remote_pair): + _, client = remote_pair + with pytest.raises(RemoteProtocolError): + client.submit(add, 1, 2, deadpool_metadata={"bad": object()}) + assert client.submit(add, 2, 3).result(timeout=5) == 5 + + +def test_task_exception_and_encoding_boundaries(remote_pair): + _, client = remote_pair + + with pytest.raises(ValueError, match="remote boom"): + client.submit(fail_value_error).result(timeout=5) + with pytest.raises(RemoteTaskError) as captured: + client.submit(fail_unpicklable).result(timeout=5) + assert "UnpicklableError" in captured.value.remote_traceback + with pytest.raises(RemoteResultEncodingError): + client.submit(unencodable_result).result(timeout=5) + + +class RejectResultSerializer(PickleSerializer): + def loads(self, payload): + raise ValueError("cannot decode outcome") + + +class SlowSerializer(PickleSerializer): + def dumps(self, value, *, limit): + time.sleep(0.1) + return super().dumps(value, limit=limit) + + +class LargeEncodingFailureSerializer(PickleSerializer): + def dumps(self, value, *, limit): + if value == "trigger-large-encoding-failure": + raise ValueError("encoding exploded: " + "\x00" * 20_000) + return super().dumps(value, limit=limit) + + +def test_large_error_descriptors_do_not_disconnect_session(tmp_path): + socket_path = tmp_path / "pool.sock" + serializer = LargeEncodingFailureSerializer() + limits = RemoteLimits(max_control_bytes=4096, max_metadata_bytes=1024) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=2, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + serializer=serializer, + limits=limits, + ).start() + client = DeadpoolClient( + UnixAddress(socket_path), serializer=serializer, limits=limits + ) + try: + task_failure = client.submit(fail_large_value_error) + encoding_failure = client.submit(large_encoding_failure_result) + unrelated = client.submit(delayed, "healthy", 0.05) + + with pytest.raises(ValueError, match="task exploded"): + task_failure.result(timeout=5) + with pytest.raises(RemoteResultEncodingError, match="encoding exploded"): + encoding_failure.result(timeout=5) + assert unrelated.result(timeout=5) == "healthy" + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_failed_result_decode_is_not_acknowledged(tmp_path): + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient( + UnixAddress(socket_path), + serializer=RejectResultSerializer(), + ) + try: + future = client.submit(add, 1, 2) + with pytest.raises(RemoteProtocolError, match="cannot decode outcome"): + future.result(timeout=5) + assert server.get_statistics()["remote_retained_outcomes"] == 1 + finally: + client.shutdown() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_submission_timeout_includes_serialization(tmp_path): + socket_path = tmp_path / "pool.sock" + serializer = SlowSerializer() + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + serializer=serializer, + ).start() + client = DeadpoolClient( + UnixAddress(socket_path), + serializer=serializer, + submission_timeout=0.02, + ) + try: + with pytest.raises(RemoteSubmissionTimeout): + client.submit(add, 1, 2) + assert server.get_statistics()["remote_tasks_received"] == 0 + finally: + client.shutdown() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_retained_outcome_capacity_is_reserved_and_released(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits( + max_retained_outcomes_per_session=1, + max_retained_outcomes_global=1, + ) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + limits=limits, + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + original_enqueue = client._enqueue_control + + def omit_result_ack(kind, control): + if kind != MessageType.RESULT_ACK: + original_enqueue(kind, control) + + client._enqueue_control = omit_result_ack + try: + first = client.submit(add, 1, 2) + assert first.result(timeout=5) == 3 + assert server.get_statistics()["remote_retained_outcomes"] == 1 + with pytest.raises(RemoteQueueFull): + client.submit(add, 2, 3).result(timeout=5) + original_enqueue(MessageType.RESULT_ACK, {"request_id": first.request_id}) + deadline = time.monotonic() + 2 + while ( + server.get_statistics()["remote_retained_outcomes"] + and time.monotonic() < deadline + ): + time.sleep(0.01) + assert server.get_statistics()["remote_retained_outcomes"] == 0 + assert client.submit(add, 3, 4).result(timeout=5) == 7 + finally: + client._enqueue_control = original_enqueue + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_live_session_releases_outcomes_across_idle_cycles(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits( + max_retained_outcomes_per_session=1, + max_retained_outcomes_global=1, + ) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + limits=limits, + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + try: + for args, expected in [((1, 2), 3), ((3, 4), 7)]: + assert client.submit(add, *args).result(timeout=5) == expected + deadline = time.monotonic() + 2 + while ( + server.get_statistics()["remote_retained_outcomes"] + and time.monotonic() < deadline + ): + time.sleep(0.01) + stats = server.get_statistics() + assert stats["remote_retained_outcomes"] == 0 + assert stats["remote_sessions"] == 1 + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_continue_disconnect_releases_unreachable_terminal_outcome(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits( + max_retained_outcomes_per_session=1, + max_retained_outcomes_global=1, + ) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + limits=limits, + disconnect_policy="continue", + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + original_enqueue = client._enqueue_control + + def omit_result_ack(kind, control): + if kind != MessageType.RESULT_ACK: + original_enqueue(kind, control) + + client._enqueue_control = omit_result_ack + try: + assert client.submit(add, 1, 2).result(timeout=5) == 3 + assert server.get_statistics()["remote_retained_outcomes"] == 1 + sock = client._socket + assert sock is not None + sock.shutdown(socket.SHUT_RDWR) + sock.close() + deadline = time.monotonic() + 2 + while ( + server.get_statistics()["remote_retained_outcomes"] + and time.monotonic() < deadline + ): + time.sleep(0.01) + assert server.get_statistics()["remote_retained_outcomes"] == 0 + assert server.get_statistics()["remote_sessions"] == 0 + finally: + client._enqueue_control = original_enqueue + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_callback_capacity_one_allows_reentrant_registration(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits(callback_queue_size=1) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + finished = threading.Event() + try: + future = client.submit(add, 2, 3) + + def first_callback(done): + done.add_done_callback(lambda _: finished.set()) + + future.add_done_callback(first_callback) + assert future.result(timeout=5) == 5 + assert finished.wait(2) + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_callback_registration_does_not_wait_for_dispatch_capacity(tmp_path): + socket_path = tmp_path / "pool.sock" + limits = RemoteLimits(callback_queue_size=1, completion_workers=1) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + release = tmp_path / "release" + marker = tmp_path / "marker" + registration_finished = threading.Event() + callback_started = threading.Event() + callback_release = threading.Event() + second_callback_finished = threading.Event() + third_callback_finished = threading.Event() + registration_errors = [] + callback_order = [] + registration_thread = None + try: + future = client.submit(wait_then_mark, release, marker, "done") + deadline = time.monotonic() + 2 + while not future.running() and time.monotonic() < deadline: + time.sleep(0.005) + assert future.running() + + def register_callbacks() -> None: + try: + + def blocking_callback(_) -> None: + callback_order.append(0) + callback_started.set() + callback_release.wait(5) + + future.add_done_callback(blocking_callback) + for index in (1, 2): + future.add_done_callback( + lambda _, index=index: callback_order.append(index) + ) + except BaseException as error: + registration_errors.append(error) + finally: + registration_finished.set() + + registration_thread = threading.Thread(target=register_callbacks) + registration_thread.start() + assert registration_finished.wait(1) + assert registration_errors == [] + + release.touch() + assert future.result(timeout=5) == "done" + assert callback_started.wait(2) + + second = client.submit(add, 2, 3) + second.add_done_callback(lambda _: second_callback_finished.set()) + assert second.result(timeout=5) == 5 + deadline = time.monotonic() + 2 + capacity_exhausted = False + while time.monotonic() < deadline: + if client._callback_slots.acquire(blocking=False): + client._callback_slots.release() + time.sleep(0.005) + else: + capacity_exhausted = True + break + assert capacity_exhausted + + third = client.submit(add, 3, 4) + third.add_done_callback(lambda _: third_callback_finished.set()) + assert third.result(timeout=5) == 7 + late_callbacks = [] + third.add_done_callback(lambda _: late_callbacks.append("late")) + assert late_callbacks == ["late"] + assert client.submit(add, 20, 22).result(timeout=5) == 42 + assert client.check_health() + + callback_release.set() + assert second_callback_finished.wait(2) + assert third_callback_finished.wait(2) + deadline = time.monotonic() + 2 + while len(callback_order) < 3 and time.monotonic() < deadline: + time.sleep(0.005) + assert callback_order == [0, 1, 2] + finally: + release.touch(exist_ok=True) + callback_release.set() + if registration_thread is not None: + registration_thread.join(5) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_slow_callback_does_not_block_results_or_health(remote_pair): + _, client = remote_pair + callback_started = threading.Event() + callback_release = threading.Event() + + first = client.submit(add, 1, 1) + + def block_callback(future): + callback_started.set() + callback_release.wait(5) + + first.add_done_callback(block_callback) + assert first.result(timeout=5) == 2 + assert callback_started.wait(5) + second = client.submit(add, 2, 2) + assert second.result(timeout=5) == 4 + assert client.check_health() + callback_release.set() + + callback_order = [] + ordered = client.submit(add, 3, 3) + ordered.add_done_callback(lambda future: callback_order.append(1)) + ordered.add_done_callback(lambda future: callback_order.append(2)) + assert ordered.result(timeout=5) == 6 + deadline = time.monotonic() + 2 + while len(callback_order) < 2 and time.monotonic() < deadline: + time.sleep(0.01) + assert callback_order == [1, 2] diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py new file mode 100644 index 0000000..466286c --- /dev/null +++ b/tests/remote/test_protocol.py @@ -0,0 +1,321 @@ +import hashlib +import socket +import struct +import threading + +import pytest + +from deadpool.remote import RemoteLimits, RemoteProtocolError +from deadpool.remote._protocol import ( + MAGIC, + MAJOR, + MINOR, + Message, + MessageReader, + MessageType, + _json_loads, + _send_exact, + send_message, + validate_control, +) + + +def small_limits(**changes): + values = { + "max_control_bytes": 4096, + "max_frame_payload_bytes": 16, + "max_message_bytes": 128, + "max_invocation_bytes": 128, + "max_result_bytes": 128, + "max_metadata_bytes": 1024, + "max_chunks": 16, + } + values.update(changes) + return RemoteLimits(**values) + + +def test_chunked_round_trip_preserves_opaque_payload(): + sender, receiver = socket.socketpair() + payload = bytes(range(80)) + thread = threading.Thread( + target=send_message, + args=( + sender, + Message(MessageType.RESULT, {"request_id": "r"}, payload), + small_limits(), + ), + ) + thread.start() + received = MessageReader(small_limits()).receive(receiver) + thread.join() + sender.close() + receiver.close() + + assert received.kind == MessageType.RESULT + assert received.control == {"request_id": "r"} + assert received.payload == payload + + +def test_partial_frame_has_a_bounded_deadline(): + sender, receiver = socket.socketpair() + sender.sendall(MAGIC[:1]) + with pytest.raises(TimeoutError, match="partial remote frame"): + MessageReader(small_limits(partial_frame_timeout=0.02)).receive(receiver) + sender.close() + receiver.close() + + +def test_partial_frame_write_has_one_total_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class PartialSendSocket: + def __init__(self) -> None: + self.sent = 0 + + def send(self, data: memoryview) -> int: + self.sent += 1 + return 1 + + now = [-0.03] + + def monotonic() -> float: + now[0] += 0.03 + return now[0] + + monkeypatch.setattr("deadpool.remote._protocol.time.monotonic", monotonic) + monkeypatch.setattr( + "deadpool.remote._protocol.select.select", + lambda readable, writable, exceptional, timeout: ([], writable, []), + ) + sock = PartialSendSocket() + + with pytest.raises(TimeoutError, match="remote frame write"): + _send_exact(sock, b"slow", timeout=0.05) + assert sock.sent < len(b"slow") + + +def test_declared_frame_limit_is_rejected_before_payload_read(): + sender, receiver = socket.socketpair() + prefix = struct.pack("!4sBBBBII", MAGIC, MAJOR, MINOR, MessageType.RESULT, 0, 2, 17) + sender.sendall(prefix + b"{}") + with pytest.raises(RemoteProtocolError, match="frame payload"): + MessageReader(small_limits()).receive(receiver) + sender.close() + receiver.close() + + +def test_remote_limits_reject_unbounded_values(): + with pytest.raises(ValueError): + RemoteLimits(max_message_bytes=0) + with pytest.raises(ValueError): + RemoteLimits(control_timeout=float("inf")) + with pytest.raises(TypeError): + RemoteLimits(max_chunks=1.5) + + +@pytest.mark.parametrize("value", [-(2**63), 2**63 - 1]) +def test_control_accepts_signed_64_bit_boundaries(value): + validate_control({"value": value}, small_limits()) + + +@pytest.mark.parametrize("value", [-(2**63) - 1, 2**63]) +def test_control_rejects_values_outside_signed_64_bit_range(value): + with pytest.raises(RemoteProtocolError, match="signed 64-bit"): + validate_control({"value": value}, small_limits()) + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + (b"\xff", "invalid control JSON"), + (b'{"key":1,"key":2}', "duplicate JSON key"), + (b'{"value":NaN}', "non-finite JSON number"), + (b"[]", "JSON object"), + ], +) +def test_hostile_control_json_is_rejected(payload, message): + with pytest.raises(RemoteProtocolError, match=message): + _json_loads(payload, small_limits()) + + +@pytest.mark.parametrize( + ("prefix", "message"), + [ + (struct.pack("!4sBBBBII", b"NOPE", MAJOR, MINOR, 20, 0, 0, 0), "magic"), + (struct.pack("!4sBBBBII", MAGIC, MAJOR + 1, MINOR, 20, 0, 0, 0), "version"), + (struct.pack("!4sBBBBII", MAGIC, MAJOR, MINOR, 20, 2, 0, 0), "flags"), + ( + struct.pack("!4sBBBBII", MAGIC, MAJOR, MINOR, 20, 0, 4097, 0), + "control header", + ), + (struct.pack("!4sBBBBII", MAGIC, MAJOR, MINOR, 255, 0, 0, 0), "message type"), + ], +) +def test_invalid_prefix_is_rejected_without_waiting_for_a_body(prefix, message): + sender, receiver = socket.socketpair() + try: + sender.sendall(prefix) + with pytest.raises(RemoteProtocolError, match=message): + MessageReader(small_limits(partial_frame_timeout=0.02)).receive(receiver) + finally: + sender.close() + receiver.close() + + +def chunk_header(**changes): + values = { + "message_id": "message", + "index": 0, + "count": 1, + "total": 0, + "digest": hashlib.sha256(b"").hexdigest(), + "control": {}, + } + values.update(changes) + return values + + +@pytest.mark.parametrize( + ("header", "payload", "message"), + [ + (chunk_header(index=True), b"", "must be integers"), + (chunk_header(index=1), b"", "index/count"), + (chunk_header(index=1, count=2), b"", "start at index zero"), + (chunk_header(total=0), b"x", "exceeds declared total"), + (chunk_header(total=1), b"x", "digest mismatch"), + ], +) +def test_chunk_state_machine_rejects_inconsistent_input(header, payload, message): + with pytest.raises(RemoteProtocolError, match=message): + MessageReader(small_limits())._accept(MessageType.RESULT, header, payload) + + +def test_deep_control_json_is_a_typed_protocol_error() -> None: + payload = b"[" * 2000 + b"0" + b"]" * 2000 + with pytest.raises(RemoteProtocolError, match="too deeply nested"): + _json_loads(payload, small_limits()) + + +def test_maximum_json_nesting_is_accepted() -> None: + payload = b'{"value":' + b"[" * 12 + b"]" * 12 + b"}" + assert _json_loads(payload, small_limits()) + + +def test_forbidden_json_nesting_is_rejected_before_decoding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_if_called(*args: object, **kwargs: object) -> object: + raise AssertionError("over-depth JSON reached the recursive decoder") + + monkeypatch.setattr("deadpool.remote._protocol.json.loads", fail_if_called) + payload = b'{"escaped":"\\\\","value":' + b"[" * 13 + b"]" * 13 + b"}" + with pytest.raises(RemoteProtocolError, match="too deeply nested"): + _json_loads(payload, small_limits()) + + +def test_json_nesting_check_ignores_delimiters_in_strings() -> None: + payload = b'{"value":"' + b'\\"' + b"[" * 20 + b'"}' + assert _json_loads(payload, small_limits()) == {"value": '"' + "[" * 20} + + +def test_json_decoder_recursion_error_is_normalized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def recurse(*args: object, **kwargs: object) -> object: + raise RecursionError("decoder recursion limit") + + monkeypatch.setattr("deadpool.remote._protocol.json.loads", recurse) + with pytest.raises(RemoteProtocolError, match="invalid control JSON"): + _json_loads(b"{}", small_limits()) + + +def test_chunk_state_bounds_incomplete_messages_and_conflicting_metadata(): + reader = MessageReader(small_limits(max_incomplete_messages=1)) + digest = hashlib.sha256(b"xx").hexdigest() + first = chunk_header(count=2, total=2, digest=digest) + assert reader._accept(MessageType.RESULT, first, b"x") is None + + with pytest.raises(RemoteProtocolError, match="too many incomplete"): + reader._accept( + MessageType.RESULT, + chunk_header(message_id="other", count=2, total=2, digest=digest), + b"x", + ) + with pytest.raises(RemoteProtocolError, match="conflicting chunk metadata"): + reader._accept( + MessageType.RESULT, + chunk_header(index=1, count=2, total=3, digest=digest, control=None), + b"x", + ) + + completed = reader._accept( + MessageType.RESULT, + chunk_header(index=1, count=2, total=2, digest=digest, control=None), + b"x", + ) + assert completed == Message(MessageType.RESULT, {}, b"xx") + + +def test_chunk_state_bounds_aggregate_incomplete_payload_bytes() -> None: + limits = small_limits( + max_frame_payload_bytes=2, + max_message_bytes=4, + max_invocation_bytes=4, + max_result_bytes=4, + max_chunks=2, + max_incomplete_messages=4, + ) + digest = hashlib.sha256(b"xxxx").hexdigest() + reader = MessageReader(limits) + for message_id in ("first", "second"): + assert ( + reader._accept( + MessageType.RESULT, + chunk_header(message_id=message_id, count=2, total=4, digest=digest), + b"xx", + ) + is None + ) + + with pytest.raises(RemoteProtocolError, match="aggregate incomplete payload"): + reader._accept( + MessageType.RESULT, + chunk_header(message_id="third", count=2, total=4, digest=digest), + b"x", + ) + + completing = MessageReader(limits) + assert ( + completing._accept( + MessageType.RESULT, + chunk_header(message_id="complete", count=2, total=4, digest=digest), + b"xx", + ) + is None + ) + assert completing._accept( + MessageType.RESULT, + chunk_header( + message_id="complete", + index=1, + count=2, + total=4, + digest=digest, + control=None, + ), + b"xx", + ) == Message(MessageType.RESULT, {}, b"xxxx") + assert completing._incomplete_bytes == 0 + + +def test_oversized_outbound_message_is_rejected_before_socket_io(): + class NoIoSocket: + def send(self, data): + raise AssertionError("send must not be called") + + with pytest.raises(RemoteProtocolError, match="message payload"): + send_message( + NoIoSocket(), + Message(MessageType.RESULT, {}, b"x" * 129), + small_limits(), + ) diff --git a/tests/remote/test_public_end_to_end.py b/tests/remote/test_public_end_to_end.py new file mode 100644 index 0000000..55e0990 --- /dev/null +++ b/tests/remote/test_public_end_to_end.py @@ -0,0 +1,122 @@ +"""End-to-end coverage for public remote client/server interfaces.""" + +from __future__ import annotations + +import multiprocessing +from functools import partial +from pathlib import Path + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteLimits, + UnixAddress, + UnixListener, +) +from tests.remote.tasks import echo_bytes, multiply + + +def _spawn_submitter( + socket_path: str, + index: int, + start: multiprocessing.synchronize.Event, + output: multiprocessing.queues.Queue, +) -> None: + """Connect after the shared gate and report only public client outcomes.""" + output.put(("ready", index)) + if not start.wait(10): + output.put(("error", index, "start gate timed out")) + return + + client = None + try: + client = DeadpoolClient(UnixAddress(socket_path)) + result = client.submit_task("multiply", index, index + 1).result(10) + output.put(("result", index, result)) + except BaseException as error: + output.put(("error", index, f"{type(error).__name__}: {error}")) + finally: + if client is not None: + client.shutdown(wait=True, cancel_futures=True) + + +def test_spawn_clients_submit_concurrently_to_one_server_pool(tmp_path: Path) -> None: + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=2, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"multiply": multiply}, + ).start() + context = multiprocessing.get_context("spawn") + start = context.Event() + output = context.Queue() + process_count = 4 + processes = [ + context.Process( + target=_spawn_submitter, + args=(str(socket_path), index, start, output), + ) + for index in range(1, process_count + 1) + ] + + try: + for process in processes: + process.start() + ready = {output.get(timeout=10) for _ in processes} + assert ready == {("ready", index) for index in range(1, process_count + 1)} + + start.set() + outcomes = [output.get(timeout=15) for _ in processes] + errors = [outcome for outcome in outcomes if outcome[0] == "error"] + assert errors == [] + assert sorted(outcomes) == [ + ("result", index, index * (index + 1)) + for index in range(1, process_count + 1) + ] + for process in processes: + process.join(10) + assert process.exitcode == 0 + finally: + start.set() + for process in processes: + if process.is_alive(): + process.terminate() + process.join(5) + if process.is_alive(): + process.kill() + process.join(5) + assert not process.is_alive() + output.close() + output.join_thread() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_chunked_invocation_and_result_round_trip_and_session_health( + tmp_path: Path, +) -> None: + frame_size = 8 * 1024 + limits = RemoteLimits( + max_frame_payload_bytes=frame_size, + max_message_bytes=128 * 1024, + max_invocation_bytes=64 * 1024, + max_result_bytes=64 * 1024, + max_chunks=16, + ) + socket_path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(socket_path)], + task_registry={"echo_bytes": echo_bytes, "multiply": multiply}, + limits=limits, + ).start() + client = DeadpoolClient(UnixAddress(socket_path), limits=limits) + payload = bytes(range(256)) * 40 + b"exact-tail" + assert len(payload) > frame_size + try: + assert client.submit_task("echo_bytes", payload).result(10) == payload + assert client.submit_task("multiply", 6, 7).result(5) == 42 + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) diff --git a/tests/remote/test_scheduling_resilience.py b/tests/remote/test_scheduling_resilience.py new file mode 100644 index 0000000..f5613d9 --- /dev/null +++ b/tests/remote/test_scheduling_resilience.py @@ -0,0 +1,322 @@ +import socket +import threading +import time +from functools import partial +from pathlib import Path +from typing import Callable + +import pytest + +import deadpool +from deadpool._pool import WorkerProcess +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + Principal, + RemoteCompatibilityError, + RemoteConnectionLost, + RemoteExecutorError, + RemoteLimits, + RemoteQueueFull, + SubmissionState, + UnixAddress, + UnixListener, +) +from tests.remote.tasks import append_line, mark, multiply, wait_then_mark + + +def wait_until(predicate: Callable[[], bool], timeout: float = 5.0) -> None: + """Poll a cross-thread/process observation without assuming scheduler speed.""" + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError("condition did not become true before deadline") + time.sleep(0.005) + + +def principal_from_hello(peer: object, hello: dict) -> Principal: + authentication = hello.get("authentication") or {} + return Principal(authentication.get("principal", "anonymous")) + + +def allow_all(principal: Principal, operation: str, metadata: dict) -> bool: + return True + + +def authenticated_client( + path: Path, principal: str, **kwargs: object +) -> DeadpoolClient: + return DeadpoolClient( + UnixAddress(path), + authenticator=lambda: {"principal": principal}, + **kwargs, + ) + + +def disconnect_abruptly(client: DeadpoolClient) -> None: + """Inject a real transport break without invoking clean session shutdown.""" + sock = client._socket + assert sock is not None + sock.shutdown(socket.SHUT_RDWR) + sock.close() + + +def test_fingerprint_compatibility_and_authorization_are_end_to_end( + tmp_path: Path, +) -> None: + path = tmp_path / "pool.sock" + + def authorize(principal: Principal, operation: str, metadata: dict) -> bool: + return operation != "submit_task:multiply" + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + task_registry={"multiply": multiply}, + registry_fingerprint="registry-v1", + application_fingerprint="app-v1", + authenticator=principal_from_hello, + authorizer=authorize, + ).start() + client = authenticated_client( + path, + "alice", + registry_fingerprint="registry-v1", + application_fingerprint="app-v1", + ) + try: + assert client.submit(multiply, 6, 7).result(5) == 42 + with pytest.raises(RemoteExecutorError, match="unauthorized"): + client.submit_task("multiply", 2, 3).result(5) + + with pytest.raises(RemoteCompatibilityError): + authenticated_client( + path, + "alice", + registry_fingerprint="registry-v2", + application_fingerprint="app-v1", + ) + with pytest.raises(RemoteCompatibilityError): + authenticated_client( + path, + "alice", + registry_fingerprint="registry-v1", + application_fingerprint="app-v2", + ) + assert client.check_health() + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_worker_pipe_write_does_not_hold_broker_lock( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + entered = threading.Event() + release = threading.Event() + original_submit_job = WorkerProcess.submit_job + + def blocked_submit_job(self: WorkerProcess, job: object) -> None: + entered.set() + if not release.wait(5): + raise TimeoutError("worker pipe test barrier was not released") + original_submit_job(self, job) + + monkeypatch.setattr(WorkerProcess, "submit_job", blocked_submit_job) + path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + ).start() + client = DeadpoolClient(UnixAddress(path)) + marker = tmp_path / "must-not-run" + future = client.submit(mark, marker) + try: + assert entered.wait(2) + queued = client.submit(multiply, 6, 7) + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) + assert future.cancel_and_kill_if_running() + release.set() + + assert future.cancelled() + assert queued.result(5) == 42 + assert not marker.exists() + finally: + release.set() + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_worker_pipe_broken_pipe_retry_preserves_running_state( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + attempted_pids: list[int] = [] + original_submit_job = WorkerProcess.submit_job + + def fail_first_submit(self: WorkerProcess, job: object) -> None: + attempted_pids.append(self.pid) + if len(attempted_pids) == 1: + raise BrokenPipeError("injected worker pipe failure") + original_submit_job(self, job) + + monkeypatch.setattr(WorkerProcess, "submit_job", fail_first_submit) + path = tmp_path / "pool.sock" + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + ).start() + client = DeadpoolClient(UnixAddress(path)) + try: + future = client.submit(multiply, 6, 7) + assert future.result(5) == 42 + assert len(attempted_pids) == 2 + assert attempted_pids[0] != attempted_pids[1] + assert future.pid == attempted_pids[-1] + assert server.get_statistics()["remote_tasks_running"] == 1 + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_global_admission_recovers_for_another_principal(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + limits = RemoteLimits( + max_pending_per_session=1, + max_pending_per_principal=1, + max_pending_global=1, + ) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + authenticator=principal_from_hello, + authorizer=allow_all, + limits=limits, + ).start() + first = authenticated_client(path, "alice", limits=limits) + second = authenticated_client(path, "bob", limits=limits) + release = tmp_path / "release" + marker_path = tmp_path / "first" + try: + active = first.submit(wait_then_mark, release, marker_path, "active") + wait_until(active.running) + with pytest.raises(RemoteQueueFull): + second.submit(multiply, 2, 3).result(5) + + release.touch() + assert active.result(5) == "active" + wait_until(lambda: server.get_statistics()["remote_retained_outcomes"] == 0) + assert second.submit(multiply, 6, 7).result(5) == 42 + finally: + release.touch(exist_ok=True) + first.shutdown(cancel_futures=True) + second.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_priority_and_principal_fairness_are_observable(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + limits = RemoteLimits(max_staged_tasks=1) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + authenticator=principal_from_hello, + authorizer=allow_all, + limits=limits, + ).start() + alice = authenticated_client(path, "alice", limits=limits) + bob = authenticated_client(path, "bob", limits=limits) + release = tmp_path / "release" + blocker_marker = tmp_path / "blocker" + order_path = tmp_path / "order" + try: + blocker = alice.submit( + wait_then_mark, release, blocker_marker, "blocker", deadpool_priority=5 + ) + wait_until(blocker.running) + alice_one = alice.submit( + append_line, order_path, "alice-1", deadpool_priority=2 + ) + wait_until(lambda: alice.get_status(alice_one) == "ACCEPTED_QUEUED") + alice_two = alice.submit( + append_line, order_path, "alice-2", deadpool_priority=2 + ) + wait_until(lambda: alice.get_status(alice_two) == "ACCEPTED_QUEUED") + bob_one = bob.submit(append_line, order_path, "bob-1", deadpool_priority=2) + wait_until(lambda: bob.get_status(bob_one) == "ACCEPTED_QUEUED") + urgent = bob.submit(append_line, order_path, "urgent", deadpool_priority=0) + wait_until(lambda: bob.get_status(urgent) == "ACCEPTED_QUEUED") + queued = (alice_one, alice_two, bob_one, urgent) + + release.touch() + assert blocker.result(5) == "blocker" + assert [future.result(5) for future in queued] == [ + "alice-1", + "alice-2", + "bob-1", + "urgent", + ] + assert order_path.read_text().splitlines() == [ + "urgent", + "alice-1", + "bob-1", + "alice-2", + ] + finally: + release.touch(exist_ok=True) + alice.shutdown(cancel_futures=True) + bob.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +@pytest.mark.parametrize( + ("policy", "first_runs", "queued_runs"), + [ + ("cancel_queued", True, False), + ("continue", True, True), + ("terminate", False, False), + ], +) +def test_disconnect_policies_have_observable_worker_side_effects( + tmp_path: Path, policy: str, first_runs: bool, queued_runs: bool +) -> None: + path = tmp_path / "pool.sock" + limits = RemoteLimits(max_staged_tasks=1) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[UnixListener(path)], + limits=limits, + disconnect_policy=policy, + ).start() + client = DeadpoolClient(UnixAddress(path), limits=limits) + release = tmp_path / "release" + first_marker = tmp_path / "first" + queued_marker = tmp_path / "queued" + try: + active = client.submit(wait_then_mark, release, first_marker, "first") + wait_until(active.running) + queued = client.submit(mark, queued_marker, "queued") + wait_until( + lambda: client.get_status(queued) == "ACCEPTED_QUEUED" + and queued.submission_state is SubmissionState.ACCEPTED_QUEUED + ) + + disconnect_abruptly(client) + wait_until(lambda: server.get_statistics()["remote_connections"] == 0) + with pytest.raises(RemoteConnectionLost): + active.result(2) + with pytest.raises(RemoteConnectionLost): + queued.result(2) + release.touch() + + recovery = DeadpoolClient(UnixAddress(path), limits=limits) + try: + assert recovery.submit(multiply, 6, 7).result(5) == 42 + finally: + recovery.shutdown(cancel_futures=True) + + assert first_marker.exists() is first_runs + assert queued_marker.exists() is queued_runs + finally: + release.touch(exist_ok=True) + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) diff --git a/tests/remote/test_tcp.py b/tests/remote/test_tcp.py new file mode 100644 index 0000000..8947aa9 --- /dev/null +++ b/tests/remote/test_tcp.py @@ -0,0 +1,244 @@ +from functools import partial +import socket +import ssl +import time +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteCompatibilityError, + RemoteExecutorError, + RemoteExecutorUnavailable, + Principal, + RemoteLimits, + TcpAddress, + TcpListener, +) +from tests.remote.tasks import multiply + + +def identity(value): + return value + + +def test_insecure_loopback_tcp_uses_same_protocol(): + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, insecure=True)], + authorizer=lambda principal, operation, metadata: True, + ).start() + host, port = server.bound_addresses[0] + client = DeadpoolClient(TcpAddress(host, port, insecure=True)) + try: + assert client.submit(identity, "tcp").result(timeout=5) == "tcp" + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_loopback_mutual_tls_authenticates_client_identity(tls_contexts): + authenticated_common_names = [] + authorized_calls = [] + + def authenticate(peer, hello): + subject = { + key: value + for distinguished_name in peer.certificate["subject"] + for key, value in distinguished_name + } + authenticated_common_names.append(subject["commonName"]) + return Principal(subject["commonName"]) + + def authorize(principal, operation, metadata): + authorized_calls.append((principal.name, operation)) + return True + + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, ssl_context=tls_contexts.server)], + task_registry={"multiply": multiply}, + authenticator=authenticate, + authorizer=authorize, + ).start() + host, port = server.bound_addresses[0] + client = None + try: + client = DeadpoolClient( + TcpAddress( + host, + port, + ssl_context=tls_contexts.client, + server_hostname="localhost", + ) + ) + assert client.submit_task("multiply", 6, 7).result(5) == 42 + assert client.check_health() + assert authenticated_common_names == ["deadpool-test-client"] + assert authorized_calls == [("deadpool-test-client", "submit_task:multiply")] + finally: + if client is not None: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +@pytest.mark.parametrize( + "context_name", + ["missing_client_certificate", "untrusted_client"], +) +def test_mutual_tls_rejects_missing_or_untrusted_client_certificate( + tls_contexts, context_name +): + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, ssl_context=tls_contexts.server)], + authorizer=lambda principal, operation, metadata: True, + ).start() + host, port = server.bound_addresses[0] + try: + with pytest.raises(RemoteExecutorUnavailable): + DeadpoolClient( + TcpAddress( + host, + port, + ssl_context=getattr(tls_contexts, context_name), + server_hostname="localhost", + ) + ) + finally: + server.shutdown(cancel_futures=True, deadline=5) + + +def test_tls_rejects_untrusted_server_hostname(tls_contexts): + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, ssl_context=tls_contexts.server)], + authorizer=lambda principal, operation, metadata: True, + ).start() + host, port = server.bound_addresses[0] + try: + with pytest.raises(RemoteExecutorUnavailable, match="(?i)hostname"): + DeadpoolClient( + TcpAddress( + host, + port, + ssl_context=tls_contexts.client, + server_hostname="untrusted.invalid", + ) + ) + finally: + server.shutdown(cancel_futures=True, deadline=5) + + +def test_tcp_without_authorizer_is_default_deny(): + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, insecure=True)], + ).start() + host, port = server.bound_addresses[0] + client = DeadpoolClient(TcpAddress(host, port, insecure=True)) + try: + with pytest.raises(RemoteExecutorError, match="unauthorized"): + client.submit(identity, "tcp").result(timeout=5) + with pytest.raises(RemoteExecutorError, match="unauthorized"): + client.get_statistics() + finally: + client.shutdown() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_tcp_principal_connection_limit_cannot_be_multiplied(): + limits = RemoteLimits(max_connections_per_principal=1) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, insecure=True)], + authorizer=lambda principal, operation, metadata: True, + limits=limits, + ).start() + host, port = server.bound_addresses[0] + first = DeadpoolClient(TcpAddress(host, port, insecure=True), limits=limits) + try: + with pytest.raises(RemoteCompatibilityError): + DeadpoolClient(TcpAddress(host, port, insecure=True), limits=limits) + finally: + first.shutdown() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_unauthenticated_connection_limit_is_bounded(): + limits = RemoteLimits( + max_unauthenticated_connections=1, + max_connections_global=2, + max_connections_per_principal=2, + ) + server = DeadpoolServer( + partial(deadpool.Deadpool, max_workers=1, mp_context="forkserver"), + listeners=[TcpListener("127.0.0.1", 0, insecure=True)], + authorizer=lambda principal, operation, metadata: True, + limits=limits, + ).start() + host, port = server.bound_addresses[0] + stalled = socket.create_connection((host, port)) + try: + deadline = time.monotonic() + 2 + while ( + server.get_statistics()["remote_connections"] < 1 + and time.monotonic() < deadline + ): + time.sleep(0.01) + with pytest.raises(RemoteExecutorUnavailable): + DeadpoolClient(TcpAddress(host, port, insecure=True), limits=limits) + finally: + stalled.close() + deadline = time.monotonic() + 2 + while server.get_statistics()["remote_connections"] and time.monotonic() < deadline: + time.sleep(0.01) + client = DeadpoolClient(TcpAddress(host, port, insecure=True), limits=limits) + client.shutdown() + server.shutdown(cancel_futures=True, deadline=5) + + +def test_tls_configuration_requires_mutual_verification(): + client_context = ssl.create_default_context() + client_context.minimum_version = ssl.TLSVersion.TLSv1_2 + TcpAddress("example.com", 443, ssl_context=client_context) + + insecure_client = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + insecure_client.check_hostname = False + insecure_client.verify_mode = ssl.CERT_NONE + with pytest.raises(ValueError, match="certificates and hostnames"): + TcpAddress("example.com", 443, ssl_context=insecure_client) + + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.minimum_version = ssl.TLSVersion.TLSv1_2 + with pytest.raises(ValueError, match="client certificates"): + TcpListener("127.0.0.1", 443, ssl_context=server_context) + + +def test_tls_configuration_rejects_versions_older_than_tls_1_2(): + client_context = ssl.create_default_context() + with pytest.warns(DeprecationWarning): + client_context.minimum_version = ssl.TLSVersion.TLSv1 + with pytest.raises(ValueError, match="client TLS must require TLS 1.2"): + TcpAddress("example.com", 443, ssl_context=client_context) + + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.verify_mode = ssl.CERT_REQUIRED + with pytest.warns(DeprecationWarning): + server_context.minimum_version = ssl.TLSVersion.TLSv1 + with pytest.raises(ValueError, match="server TLS must require TLS 1.2"): + TcpListener("127.0.0.1", 443, ssl_context=server_context) + + +def test_plaintext_tcp_requires_explicit_loopback_opt_in(): + with pytest.raises(ValueError): + TcpAddress("127.0.0.1", 1) + with pytest.raises(ValueError, match="addresses are restricted to loopback"): + TcpAddress("example.com", 443, insecure=True) + with pytest.raises(ValueError, match="addresses are restricted to loopback"): + TcpAddress("localhost", 443, insecure=True) + with pytest.raises(ValueError, match="listeners are restricted to loopback"): + TcpListener("0.0.0.0", 0, insecure=True) + with pytest.raises(ValueError, match="listeners are restricted to loopback"): + TcpListener("localhost", 0, insecure=True) diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py new file mode 100644 index 0000000..12a79f6 --- /dev/null +++ b/tests/remote/test_units.py @@ -0,0 +1,193 @@ +from concurrent.futures import as_completed, wait + +import pytest + +from deadpool.remote import ( + AcceptanceCertainty, + ExecutionCertainty, + PickleSerializer, + RemoteFuture, + RemoteLimits, + RemoteQueueFull, + SerializationLimitError, + SubmissionState, +) +from deadpool.remote._scheduler import FairScheduler +from deadpool.remote.serializer import _LimitedWriter + + +def test_pickle_serializer_enforces_limit_and_round_trips(): + serializer = PickleSerializer() + payload = serializer.dumps({"answer": 42}, limit=1000) + assert serializer.loads(payload) == {"answer": 42} + with pytest.raises(SerializationLimitError): + serializer.dumps(b"x" * 100, limit=10) + + +def test_infrastructure_error_exposes_retry_certainty(): + error = RemoteQueueFull("full", request_id="client:1") + assert error.request_id == "client:1" + assert error.acceptance_certainty is AcceptanceCertainty.NOT_ACCEPTED + assert error.execution_certainty is ExecutionCertainty.NOT_STARTED + + +def test_scheduler_is_strict_priority_and_principal_fair(): + scheduler = FairScheduler() + scheduler.put("a1", priority=2, principal="a") + scheduler.put("a2", priority=2, principal="a") + scheduler.put("b1", priority=2, principal="b") + scheduler.put("urgent", priority=1, principal="a") + + assert [scheduler.pop() for _ in range(4)] == ["urgent", "a1", "b1", "a2"] + + +def test_scheduler_front_requeue_preserves_fifo_priority_and_fairness(): + scheduler = FairScheduler() + scheduler.put("a1", priority=2, principal="a") + scheduler.put("a2", priority=2, principal="a") + scheduler.put("b1", priority=2, principal="b") + + rejected = scheduler.pop() + assert rejected == "a1" + scheduler.put("urgent", priority=1, principal="a") + scheduler.put_front(rejected, priority=2, principal="a") + + assert [scheduler.pop() for _ in range(4)] == ["urgent", "b1", "a1", "a2"] + + +def test_scheduler_removal_preserves_other_principals_and_priorities(): + scheduler = FairScheduler() + scheduler.put("a1", priority=2, principal="a") + scheduler.put("a2", priority=2, principal="a") + scheduler.put("b1", priority=2, principal="b") + scheduler.put("urgent", priority=1, principal="a") + + assert scheduler.remove("a1") + assert not scheduler.remove("missing") + assert scheduler.remove("urgent") + assert len(scheduler) == 2 + assert [scheduler.pop(), scheduler.pop()] == ["a2", "b1"] + with pytest.raises(IndexError, match="empty"): + scheduler.pop() + + +def test_limited_writer_rejects_atomically_and_validates_limit(): + with pytest.raises(ValueError, match="non-negative"): + _LimitedWriter(-1) + + writer = _LimitedWriter(3) + assert writer.write(b"abc") == 3 + before = bytes(writer.buffer) + with pytest.raises(SerializationLimitError): + writer.write(b"d") + assert bytes(writer.buffer) == before + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"max_frame_payload_bytes": 65 * 1024 * 1024}, "frame_payload"), + ({"max_invocation_bytes": 65 * 1024 * 1024}, "invocation"), + ({"max_result_bytes": 65 * 1024 * 1024}, "result"), + ({"max_frame_payload_bytes": 1, "max_chunks": 1}, "max_chunks"), + ( + {"max_connections_global": 1, "max_connections_per_principal": 2}, + "connections_per_principal", + ), + ( + {"max_pending_global": 1, "max_pending_per_principal": 2}, + "pending_per_principal", + ), + ( + { + "max_retained_outcomes_global": 1, + "max_retained_outcomes_per_session": 2, + }, + "retained outcome count", + ), + ( + { + "max_retained_outcome_bytes_global": 1, + "max_retained_outcome_bytes_per_session": 2, + }, + "retained outcome bytes", + ), + ], +) +def test_remote_limit_cross_field_relationships(changes, message): + with pytest.raises(ValueError, match=message): + RemoteLimits(**changes) + + +class CallbackClient: + def __init__(self): + self.cancel_calls = 0 + self.single_callbacks = [] + + def _schedule_callbacks(self, callbacks, future): + for callback in callbacks: + callback(future) + + def _schedule_callback(self, callback, future): + # Model the old asynchronous late-callback path without running it. + self.single_callbacks.append((callback, future)) + + def _cancel(self, future, *, hard): + self.cancel_calls += 1 + return False + + +def test_future_terminal_state_and_callbacks_are_idempotent(): + client = CallbackClient() + future = RemoteFuture("request:1", client) + callbacks = [] + future.add_done_callback(lambda done: callbacks.append(done.result())) + future._set_running(pid=123, worker_id="worker:123") + future._set_result("ok") + + assert not future._set_sent() + future._set_accepted() + future._set_running(pid=456, worker_id="worker:456") + future._set_exception(RuntimeError("late")) + future._set_cancelled() + + assert future.result() == "ok" + assert future.submission_state is SubmissionState.SUCCEEDED + assert future.pid == 123 + assert future.worker_id == "worker:123" + assert callbacks == ["ok"] + + +def test_future_late_callback_runs_inline_and_logs_errors(caplog): + client = CallbackClient() + future = RemoteFuture("request:late", client) + future._set_result("ok") + observed = [] + + future.add_done_callback(lambda done: observed.append(done.result())) + + assert observed == ["ok"] + assert client.single_callbacks == [] + + def broken_callback(done): + raise RuntimeError(f"late callback failed for {done.request_id}") + + future.add_done_callback(broken_callback) + assert "late callback failed for request:late" in caplog.text + + +def test_future_local_cancellation_never_contacts_server(): + client = CallbackClient() + future = RemoteFuture("request:2", client) + callbacks = [] + future.add_done_callback(lambda done: callbacks.append(done.cancelled())) + + assert future.cancel() + assert future.cancelled() + assert future.submission_state is SubmissionState.CANCELLED + assert client.cancel_calls == 0 + assert callbacks == [True] + done, not_done = wait([future], timeout=0.1) + assert done == {future} + assert not not_done + assert list(as_completed([future], timeout=0.1)) == [future] diff --git a/tests/remote/test_unix_lifecycle.py b/tests/remote/test_unix_lifecycle.py new file mode 100644 index 0000000..681c597 --- /dev/null +++ b/tests/remote/test_unix_lifecycle.py @@ -0,0 +1,465 @@ +import errno +import socket +import stat +import threading +import time +from pathlib import Path + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteLimits, + ServerState, + UnixAddress, + UnixListener, +) +from deadpool.remote import _transport +from tests.remote.tasks import multiply + + +def pool_factory() -> deadpool.Deadpool: + return deadpool.Deadpool(max_workers=1, mp_context="forkserver") + + +def test_unix_context_sets_mode_and_removes_owned_socket(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + with DeadpoolServer( + pool_factory, + listeners=[UnixListener(path, mode=0o620)], + ) as server: + assert stat.S_IMODE(path.stat().st_mode) == 0o620 + with DeadpoolClient(UnixAddress(path)) as client: + assert client.submit(multiply, 6, 7).result(5) == 42 + client.shutdown() + assert server.bound_addresses == (path,) + + assert server.state is ServerState.STOPPED + assert not path.exists() + + +@pytest.mark.parametrize("path_kind", ["file", "symlink"]) +def test_failed_startup_is_stopped_repeatable_and_non_destructive( + tmp_path: Path, path_kind: str +) -> None: + path = tmp_path / "pool.sock" + if path_kind == "file": + path.write_text("keep") + else: + path.symlink_to(tmp_path / "missing-target") + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + + with pytest.raises(ValueError, match="unexpected Unix socket path"): + server.start() + assert server.state is ServerState.STOPPED + with pytest.raises(ValueError, match="unexpected Unix socket path"): + server.start() + assert path.is_symlink() if path_kind == "symlink" else path.read_text() == "keep" + + +def test_thread_start_failure_is_terminal_and_repeatable( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + + def fail_thread_start(thread: threading.Thread) -> None: + raise RuntimeError("thread unavailable") + + monkeypatch.setattr(threading.Thread, "start", fail_thread_start) + with pytest.raises(RuntimeError, match="thread unavailable"): + server.start() + + assert server.state is ServerState.STOPPED + assert server.ready.is_set() + assert server._stopped.is_set() + with pytest.raises(RuntimeError, match="thread unavailable"): + server.start() + server.shutdown() + + +def test_direct_serve_forever_propagates_startup_failure(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + path.write_text("keep") + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + + with pytest.raises(ValueError, match="unexpected Unix socket path"): + server.serve_forever() + + assert server.state is ServerState.STOPPED + assert path.read_text() == "keep" + + +def test_concurrent_start_waits_for_shared_readiness(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer( + pool_factory, + listeners=[UnixListener(path)], + limits=RemoteLimits(handshake_timeout=0.02), + ) + initialize = server._initialize + initialization_entered = threading.Event() + allow_initialization = threading.Event() + second_returned = threading.Event() + errors = [] + + def delayed_initialize() -> None: + initialization_entered.set() + assert allow_initialization.wait(5) + initialize() + + def start_server(returned: threading.Event | None = None) -> None: + try: + server.start() + except BaseException as error: + errors.append(error) + finally: + if returned is not None: + returned.set() + + server._initialize = delayed_initialize + first = threading.Thread(target=start_server) + second = threading.Thread(target=start_server, args=(second_returned,)) + try: + first.start() + assert initialization_entered.wait(5) + second.start() + assert not second_returned.wait(0.1) + assert first.is_alive() + assert errors == [] + allow_initialization.set() + first.join(5) + second.join(5) + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert server.state is ServerState.RUNNING + finally: + allow_initialization.set() + if first.ident is not None: + first.join(5) + if second.ident is not None: + second.join(5) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_start_keeps_a_published_readiness_outcome_during_shutdown( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + readiness_observed = threading.Event() + release_waiter = threading.Event() + start_errors = [] + wait_for_ready = server.ready.wait + + def delayed_wait(timeout=None) -> bool: + result = wait_for_ready(timeout) + readiness_observed.set() + assert release_waiter.wait(5) + return result + + monkeypatch.setattr(server.ready, "wait", delayed_wait) + + def start_server() -> None: + try: + server.start() + except BaseException as error: + start_errors.append(error) + + startup_thread = threading.Thread(target=start_server) + try: + startup_thread.start() + assert readiness_observed.wait(5) + assert server.state is ServerState.RUNNING + server.shutdown(cancel_futures=True, deadline=5) + assert server.state is ServerState.STOPPED + finally: + release_waiter.set() + startup_thread.join(5) + server.shutdown(cancel_futures=True, deadline=5) + + assert not startup_thread.is_alive() + assert start_errors == [] + + +def test_shutdown_during_startup_cleans_unpublished_resources( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "pool.sock" + pool_created = threading.Event() + allow_factory_return = threading.Event() + shutdown_complete = threading.Event() + startup_errors = [] + + class RecordingPool: + def __init__(self) -> None: + self.shutdown_calls = [] + + def shutdown(self, wait=True, *, cancel_futures=False) -> None: + self.shutdown_calls.append((wait, cancel_futures)) + + pool = RecordingPool() + close_listener = _transport.BoundListener.close + + def close_then_fail(listener: _transport.BoundListener) -> None: + close_listener(listener) + raise OSError("simulated listener cleanup failure") + + monkeypatch.setattr(_transport.BoundListener, "close", close_then_fail) + + def slow_pool_factory() -> RecordingPool: + pool_created.set() + assert allow_factory_return.wait(5) + return pool + + server = DeadpoolServer(slow_pool_factory, listeners=[UnixListener(path)]) + + def start_server() -> None: + try: + server.start() + except BaseException as error: + startup_errors.append(error) + + def stop_server() -> None: + server.shutdown() + shutdown_complete.set() + + startup_thread = threading.Thread(target=start_server) + shutdown_thread = threading.Thread(target=stop_server) + try: + startup_thread.start() + assert pool_created.wait(5) + shutdown_thread.start() + deadline = time.monotonic() + 2 + while server.state is not ServerState.STOPPING and time.monotonic() < deadline: + time.sleep(0.005) + assert server.state is ServerState.STOPPING + assert not shutdown_complete.wait(0.05) + finally: + allow_factory_return.set() + startup_thread.join(5) + shutdown_thread.join(5) + + assert not startup_thread.is_alive() + assert not shutdown_thread.is_alive() + assert len(startup_errors) == 1 + assert isinstance(startup_errors[0], RuntimeError) + assert str(startup_errors[0]) == "remote server stopped before becoming ready" + assert pool.shutdown_calls == [(True, True)] + assert server.state is ServerState.STOPPED + assert server.bound_addresses == () + assert not path.exists() + + +def test_optional_listener_failure_does_not_prevent_valid_listener( + tmp_path: Path, +) -> None: + blocked = tmp_path / "blocked" + blocked.write_text("keep") + valid = tmp_path / "valid.sock" + server = DeadpoolServer( + pool_factory, + listeners=[UnixListener(blocked, optional=True), UnixListener(valid)], + ).start() + client = DeadpoolClient(UnixAddress(valid)) + try: + assert server.bound_addresses == (valid,) + assert client.submit(multiply, 2, 4).result(5) == 8 + assert blocked.read_text() == "keep" + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_stale_listener_requires_force_unlink_and_then_has_normal_lifecycle( + tmp_path: Path, +) -> None: + path = tmp_path / "stale.sock" + stale = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + stale.bind(str(path)) + stale.close() + original = path.lstat() + + refusing = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + with pytest.raises(FileExistsError): + refusing.start() + assert (path.lstat().st_dev, path.lstat().st_ino) == ( + original.st_dev, + original.st_ino, + ) + + server = DeadpoolServer( + pool_factory, + listeners=[UnixListener(path, stale_policy="force_unlink")], + ).start() + client = DeadpoolClient(UnixAddress(path)) + try: + assert client.submit(multiply, 6, 7).result(5) == 42 + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + assert not path.exists() + + with DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) as restarted: + with DeadpoolClient(UnixAddress(path)) as restarted_client: + assert restarted_client.submit(multiply, 2, 4).result(5) == 8 + assert restarted.state is ServerState.RUNNING + assert not path.exists() + + +def test_stale_probe_does_not_unlink_a_replacement( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "stale.sock" + stale = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + stale.bind(str(path)) + stale.close() + real_socket = socket.socket + + class RacingProbe: + winner: socket.socket | None = None + + def settimeout(self, timeout: float) -> None: + return None + + def connect(self, address: str) -> None: + path.unlink() + self.winner = real_socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.winner.bind(address) + self.winner.listen() + raise ConnectionRefusedError + + def close(self) -> None: + return None + + probe = RacingProbe() + replacement_was_verified = False + + class ReplacementVerifier: + def settimeout(self, timeout: float) -> None: + return None + + def connect(self, address: str) -> None: + nonlocal replacement_was_verified + replacement_was_verified = True + + def close(self) -> None: + return None + + verifier = ReplacementVerifier() + calls = 0 + + def socket_factory(*args: object, **kwargs: object) -> object: + nonlocal calls + calls += 1 + if calls == 1: + return probe + if calls == 2: + return verifier + return real_socket(*args, **kwargs) + + monkeypatch.setattr(_transport.socket, "socket", socket_factory) + try: + with pytest.raises(OSError, match="Address already in use"): + _transport.bind_listener(UnixListener(path, stale_policy="force_unlink")) + assert replacement_was_verified + assert probe.winner is not None + verifier = real_socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + verifier.connect(str(path)) + finally: + verifier.close() + finally: + if probe.winner is not None: + probe.winner.close() + path.unlink(missing_ok=True) + + +def test_live_force_unlink_listener_is_preserved(tmp_path: Path) -> None: + path = tmp_path / "live.sock" + winner = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + winner.bind(str(path)) + winner.listen() + server = DeadpoolServer( + pool_factory, + listeners=[UnixListener(path, stale_policy="force_unlink")], + ) + try: + with pytest.raises(OSError, match="already live"): + server.start() + assert server.state is ServerState.STOPPED + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + probe.connect(str(path)) + finally: + probe.close() + finally: + winner.close() + path.unlink(missing_ok=True) + + +def test_insecure_socket_directory_is_rejected_without_creating_path( + tmp_path: Path, +) -> None: + directory = tmp_path / "insecure" + directory.mkdir(mode=0o777) + directory.chmod(0o777) + path = directory / "pool.sock" + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + try: + with pytest.raises(PermissionError, match="group/world writable"): + server.start() + assert server.state is ServerState.STOPPED + assert not path.exists() + finally: + directory.chmod(0o700) + + +def test_shutdown_preserves_replacement_at_listener_path(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]).start() + path.unlink() + path.write_text("replacement") + + server.shutdown(cancel_futures=True, deadline=5) + + assert path.read_text() == "replacement" + + +def test_bind_failure_never_unlinks_a_racing_winner( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A loser must not clean a pathname it never successfully bound.""" + path = tmp_path / "pool.sock" + real_socket = socket.socket + + class RacingSocket: + winner: socket.socket | None = None + + def bind(self, address: str) -> None: + self.winner = real_socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.winner.bind(address) + self.winner.listen() + raise OSError(errno.EADDRINUSE, "address in use") + + def close(self) -> None: + return None + + racing = RacingSocket() + monkeypatch.setattr(_transport.socket, "socket", lambda *args, **kwargs: racing) + + with pytest.raises(OSError, match="address in use"): + _transport.bind_listener(UnixListener(path)) + + assert path.exists() + probe = real_socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + probe.connect(str(path)) + finally: + probe.close() + assert racing.winner is not None + racing.winner.close() + path.unlink() diff --git a/tests/test_package_api.py b/tests/test_package_api.py new file mode 100644 index 0000000..fa3ad78 --- /dev/null +++ b/tests/test_package_api.py @@ -0,0 +1,16 @@ +import deadpool + + +def test_package_preserves_local_api(): + assert deadpool.__version__ == "2026.6.1" + assert deadpool.Deadpool.__module__ == "deadpool._pool" + assert deadpool.PrioritizedItem is deadpool._pool.PrioritizedItem + assert callable(deadpool.trim_memory) + + +def test_remote_api_is_importable(): + from deadpool.remote import DeadpoolClient, DeadpoolServer, RemoteFuture + + assert DeadpoolClient + assert DeadpoolServer + assert RemoteFuture