From c8eba56a85454dece0e16495905dce634c4fb3c1 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Wed, 29 Jul 2026 16:34:52 +0200 Subject: [PATCH 01/24] feat: add remote Deadpool executor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Big picture Add an experimental remote executor so independent processes can share one server-owned Deadpool over Unix-domain or TCP sockets while retaining the standard `concurrent.futures.Executor` programming model. The implementation keeps application payloads opaque in the broker, makes distributed ownership and cancellation outcomes explicit, and bounds network-facing queues, frames, connections, and retained results. The existing single-file module is migrated into a package so the remote client/server implementation can be separated behind focused internal modules without breaking current `import deadpool` users. ## Changes - Preserve the existing local executor API through package-level re-exports and compatibility forwarding while moving the pool implementation to `deadpool/_pool.py`. - Add `DeadpoolClient`, `DeadpoolServer`, and `RemoteFuture`, including callable and registered-task submission, ordered mapping, health/status/statistics, request identity, fork protection, and explicit client/server lifecycle handling. - Add a private experimental wire protocol with safe JSON control headers, opaque chunked payloads, digest validation, bounded framing, typed failures, and ordered full-duplex transport handling. - Support protected Unix-domain sockets, explicit loopback TCP, and caller-configured mutual TLS with bounded authentication and connection admission. - Add broker-owned priority and principal-fair scheduling, queue and execution timeouts, ordinary/group/hard cancellation, terminal arbitration, and bounded result retention with acknowledgement cleanup. - Add serializer limits, worker-side opaque invocation decoding, remote traceback fallbacks, and distinct transport, compatibility, queue, process, cancellation, and result exceptions. - Document the remote API, pickle trust boundary, experimental wire status, and intentionally deferred protocol extensions. - Add package compatibility, protocol, transport, cancellation, timeout, security, retention, callback, fork, and end-to-end integration coverage. ## Validation - `nox -s test-3.14` — 103 passed. - Ruff checks for the new package and tests passed. - Wheel construction succeeded and includes all remote executor modules. --- README.rst | 66 +- deadpool/__init__.py | 36 + deadpool.py => deadpool/_pool.py | 64 +- deadpool/remote/__init__.py | 93 ++ deadpool/remote/_future.py | 181 ++++ deadpool/remote/_protocol.py | 392 ++++++++ deadpool/remote/_scheduler.py | 81 ++ deadpool/remote/_transport.py | 186 ++++ deadpool/remote/_worker.py | 80 ++ deadpool/remote/client.py | 889 ++++++++++++++++++ deadpool/remote/config.py | 222 +++++ deadpool/remote/errors.py | 121 +++ deadpool/remote/serializer.py | 59 ++ deadpool/remote/server.py | 1384 +++++++++++++++++++++++++++++ tests/__init__.py | 1 + tests/remote/__init__.py | 1 + tests/remote/conftest.py | 25 + tests/remote/tasks.py | 5 + tests/remote/test_cancellation.py | 186 ++++ tests/remote/test_fork_timeout.py | 45 + tests/remote/test_integration.py | 337 +++++++ tests/remote/test_protocol.py | 80 ++ tests/remote/test_tcp.py | 129 +++ tests/remote/test_units.py | 35 + tests/test_package_api.py | 16 + 25 files changed, 4707 insertions(+), 7 deletions(-) create mode 100644 deadpool/__init__.py rename deadpool.py => deadpool/_pool.py (94%) create mode 100644 deadpool/remote/__init__.py create mode 100644 deadpool/remote/_future.py create mode 100644 deadpool/remote/_protocol.py create mode 100644 deadpool/remote/_scheduler.py create mode 100644 deadpool/remote/_transport.py create mode 100644 deadpool/remote/_worker.py create mode 100644 deadpool/remote/client.py create mode 100644 deadpool/remote/config.py create mode 100644 deadpool/remote/errors.py create mode 100644 deadpool/remote/serializer.py create mode 100644 deadpool/remote/server.py create mode 100644 tests/__init__.py create mode 100644 tests/remote/__init__.py create mode 100644 tests/remote/conftest.py create mode 100644 tests/remote/tasks.py create mode 100644 tests/remote/test_cancellation.py create mode 100644 tests/remote/test_fork_timeout.py create mode 100644 tests/remote/test_integration.py create mode 100644 tests/remote/test_protocol.py create mode 100644 tests/remote/test_tcp.py create mode 100644 tests/remote/test_units.py create mode 100644 tests/test_package_api.py diff --git a/README.rst b/README.rst index e1a808b..a84870f 100644 --- a/README.rst +++ b/README.rst @@ -763,6 +763,70 @@ 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. + +.. warning:: + + The default pickle callable mode is remote code execution by design. TLS + authenticates and encrypts peers; it does not sandbox Python deserialization + or worker execution. Only authorize mutually trusted clients. + +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 +901,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..28f0635 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,10 +711,37 @@ 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, + ) -> Future: + """Submit using a framework-owned Future. + + Remote brokerage uses this narrow private seam to arbitrate queued + cancellation against worker dispatch without changing local Future + semantics. + """ + 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), ) ) self._statistics.tasks_received.increment() @@ -720,7 +768,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..5e6d275 --- /dev/null +++ b/deadpool/remote/__init__.py @@ -0,0 +1,93 @@ +"""Experimental socket-based access to a server-owned Deadpool. + +Security warning +---------------- +Pickle callable mode intentionally executes code supplied by authenticated +clients. TLS authenticates peers; it does not sandbox Python deserialization or +execution. Use this package only across a mutually trusted boundary. + +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..76951f0 --- /dev/null +++ b/deadpool/remote/_future.py @@ -0,0 +1,181 @@ +"""Process-local Future representing one remote request.""" + +from __future__ import annotations + +import concurrent.futures +import os +import threading +import weakref +from enum import Enum +from typing import TYPE_CHECKING + +from .errors import RemoteForkedProcessError + +if TYPE_CHECKING: + from .client import DeadpoolClient + + +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[tuple[object, bool]] = [] + + @property + def pid(self) -> int | None: + self._check_process() + return self._pid + + def cancel(self) -> bool: + 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 = super().cancel() + 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) -> None: + self._check_process() + client = self._client() + reserved = client is not None + if client is not None: + client._reserve_callback() + with self._state_lock: + if not self.done(): + self._remote_callbacks.append((fn, reserved)) + return + self._schedule_callback(fn, reserved) + + 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 + super().cancel() + self._dispatch_callbacks() + + 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 not None and all(reserved for _, reserved in callbacks): + client._schedule_callbacks( + [callback for callback, _ in callbacks], + self, + ) + return + for callback, reserved in callbacks: + self._schedule_callback(callback, reserved) + + def _schedule_callback(self, callback, reserved: bool) -> None: + client = self._client() + if client is None or not reserved: + callback(self) + else: + client._schedule_callback(callback, self) + + 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..a6c0815 --- /dev/null +++ b/deadpool/remote/_protocol.py @@ -0,0 +1,392 @@ +"""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 + + +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 _json_loads(data: bytes, limits: RemoteLimits) -> dict: + try: + value = json.loads( + data.decode("utf-8"), + object_pairs_hook=_unique_object, + parse_constant=lambda value: (_ for _ in ()).throw( + ValueError(f"non-finite JSON number {value}") + ), + ) + except (UnicodeDecodeError, json.JSONDecodeError, 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 > 12: + 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 value.bit_length() > 63: + raise RemoteProtocolError("control integer exceeds 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 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) + if len(payload) > limits.max_message_bytes: + raise RemoteProtocolError("message payload exceeds configured limit") + frame_size = limits.max_frame_payload_bytes + count = max(1, (len(payload) + frame_size - 1) // frame_size) + if count > limits.max_chunks: + raise RemoteProtocolError("message requires too many chunks") + 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] + frame_control = { + "message_id": message_id, + "index": index, + "count": count, + "total": len(payload), + "digest": digest, + "control": message.control if index == 0 else None, + } + control = _json_dumps(frame_control, 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] = {} + + 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") + state.parts.append(payload) + state.received += 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] + 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 with a monotonic progress 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 + deadline = time.monotonic() + timeout + + +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..baf1478 --- /dev/null +++ b/deadpool/remote/_scheduler.py @@ -0,0 +1,81 @@ +"""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: + 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) + 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..06ce4e0 --- /dev/null +++ b/deadpool/remote/_transport.py @@ -0,0 +1,186 @@ +"""Unix and TCP socket lifecycle shared by client and server.""" + +from __future__ import annotations + +import os +import socket +import stat +import struct +from contextlib import contextmanager +from dataclasses import dataclass + +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 + and ( + current.st_dev, + current.st_ino, + ) + == self.unix_identity + ): + path.unlink() + + +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) + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + probe.settimeout(0.2) + probe.connect(str(path)) + except (ConnectionRefusedError, FileNotFoundError, socket.timeout): + path.unlink(missing_ok=True) + else: + raise OSError(f"Unix socket is already live: {path}") + finally: + probe.close() + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + with _umask(0o077): + sock.bind(str(path)) + 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) + identity = path.stat() + return BoundListener(config, sock, path, (identity.st_dev, identity.st_ino)) + except BaseException: + sock.close() + path.unlink(missing_ok=True) + raise + + +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..b692297 --- /dev/null +++ b/deadpool/remote/_worker.py @@ -0,0 +1,80 @@ +"""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={ + "module": type(error).__module__, + "type": type(error).__qualname__, + "message": _safe_repr(error), + "traceback": "".join(traceback.format_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..74ebb5f --- /dev/null +++ b/deadpool/remote/client.py @@ -0,0 +1,889 @@ +"""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, + 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 PickleSerializer, Serializer + +logger = logging.getLogger("deadpool.remote") + + +class DeadpoolClient(concurrent.futures.Executor): + """A process-local client whose submitted work runs in a server-owned pool.""" + + 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( + "Deadpool remote callable mode uses pickle and grants trusted " + "clients code execution as the worker account" + ) + 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.server_id: str | None = None + self.server_epoch: str | None = None + self.session_id: str | None = None + self.negotiated_limits: dict = {} + self._connect() + + 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, + } + 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" + ) + 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 {}) + 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: + 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.get(token) + if rpc is not None: + rpc[1].update(message.control) + 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: + acknowledge = False + 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, + ) + ) + acknowledge = True + except BaseException as error: + future._set_exception( + error + if isinstance(error, RemoteExecutorError) + else RemoteProtocolError(str(error), request_id=future.request_id) + ) + finally: + if acknowledge: + try: + self._enqueue_control( + MessageType.RESULT_ACK, {"request_id": future.request_id} + ) + 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) + self._enqueue_control( + MessageType.CANCEL_REQUEST, + {"request_id": future.request_id, "hard": hard, "token": token}, + ) + try: + 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" + self._enqueue_control(kind, {**control, key: token}) + try: + if not event.wait(timeout or self.control_timeout): + raise TimeoutError("remote control request timed out") + 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()) + if sock is not None: + try: + sock.close() + except OSError: + pass + for event, response in rpc: + response["error"] = str(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 _reserve_callback(self) -> None: + self._callback_slots.acquire() + + def _schedule_callback(self, callback, future: RemoteFuture) -> None: + try: + self._callback_executor.submit( + _run_reserved_callback, + self._callback_slots, + callback, + future, + ) + except RuntimeError: + self._callback_slots.release() + _call_callback(callback, future) + + def _schedule_callbacks(self, callbacks: list, future: RemoteFuture) -> None: + try: + self._callback_executor.submit( + _run_callbacks, + self._callback_slots, + callbacks, + future, + ) + except RuntimeError: + for callback in callbacks: + self._callback_slots.release() + _call_callback(callback, future) + + 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() + + 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, + ) + 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 and not isinstance(values["group_id"], str): + raise TypeError("deadpool_group_id must be a string or None") + 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_callbacks( + semaphore: threading.BoundedSemaphore, + callbacks: list, + future: RemoteFuture, +) -> None: + for callback in callbacks: + semaphore.release() + _call_callback(callback, future) + + +def _run_reserved_callback( + semaphore: threading.BoundedSemaphore, + callback, + future: RemoteFuture, +) -> None: + semaphore.release() + _call_callback(callback, future) + + +def _call_callback(callback, future: RemoteFuture) -> None: + try: + callback(future) + except BaseException: + import logging + + logging.getLogger("deadpool.remote").exception( + "exception calling RemoteFuture 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..424653b --- /dev/null +++ b/deadpool/remote/config.py @@ -0,0 +1,222 @@ +"""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 + + +@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) + + +@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 {"127.0.0.1", "localhost"}: + raise ValueError("insecure TCP listeners are restricted to loopback") + + +@dataclass(frozen=True, slots=True) +class RemoteLimits: + """Finite defaults bound every user-controlled queue and byte buffer.""" + + 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 + 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..70dca83 --- /dev/null +++ b/deadpool/remote/serializer.py @@ -0,0 +1,59 @@ +"""Payload serializers. + +Pickle callable mode is remote code execution by design and is only appropriate +between mutually trusted, compatible Python processes. +""" + +from __future__ import annotations + +import pickle +from typing import Protocol, runtime_checkable + + +class SerializationLimitError(ValueError): ... + + +@runtime_checkable +class Serializer(Protocol): + name: str + protocol_version: str + + def dumps(self, value: object, *, limit: int) -> bytes: ... + + def loads(self, payload: bytes) -> object: ... + + +class PickleSerializer: + 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..cf00a6f --- /dev/null +++ b/deadpool/remote/server.py @@ -0,0 +1,1384 @@ +"""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, send_message +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 PickleSerializer, Serializer + +logger = logging.getLogger("deadpool.remote") + + +class ServerState(str, Enum): + CREATED = "CREATED" + 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._begin_dispatch(record, pid) + self._after_submit = lambda pid: server._commit_running(record, pid) + self._submit_failed = lambda pid: server._abort_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. Pickle callable mode grants trusted clients code + execution as the worker account. + """ + + 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( + "Deadpool remote callable mode uses pickle and grants trusted " + "clients code execution as the worker account" + ) + 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 not None: + return self + self._serve_thread = threading.Thread( + target=self.serve_forever, name="deadpool.remote.server", daemon=False + ) + self._serve_thread.start() + self.wait_ready(self.limits.handshake_timeout) + 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.state == ServerState.RUNNING + + def serve_forever(self) -> None: + try: + self._initialize() + except BaseException as error: + self._startup_error = error + self._close_bound() + self.ready.set() + self._stopped.set() + if threading.current_thread() is not self._serve_thread: + raise + return + self._stopped.wait() + + def _initialize(self) -> None: + with self._lock: + if self.state != ServerState.CREATED: + 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: + for item in bound: + item.close() + try: + pool.shutdown(wait=True, cancel_futures=True) + except BaseException: + pass + raise + with self._lock: + 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() + + 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 [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", + "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, + ) + 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 _begin_dispatch(self, record: _Request, pid: int) -> bool: + # Hold the arbitration lock across the small pipe send. Cancellation, + # queue timeout, and dispatch can therefore have exactly one winner. + self._condition.acquire() + if record.state != RequestState.ACCEPTED_QUEUED: + self._condition.release() + return False + return True + + def _commit_running(self, record: _Request, pid: int) -> None: + try: + record.state = RequestState.RUNNING + self._statistics["remote_tasks_running"] += 1 + record.connection.send( + MessageType.RUNNING, + { + "request_id": record.request_id, + "pid": pid, + "worker_id": f"worker:{pid}", + }, + ) + finally: + self._condition.release() + + def _abort_dispatch(self, record: _Request, pid: int) -> None: + # The worker pipe did not accept bytes, so the request remains queued + # and Deadpool may retry it with another worker under the same ID. + self._condition.release() + + 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] + self._terminal_locked( + record, state, kind, outcome.descriptor or {}, 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"} + ): + for record in 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) + 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 _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 + ) + with self._condition: + if self.state == ServerState.STOPPED: + return + if self.state == ServerState.CREATED: + self.state = ServerState.STOPPED + self._stopped.set() + return + 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 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 _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", + } + 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") + 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/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/conftest.py b/tests/remote/conftest.py new file mode 100644 index 0000000..b4a92f9 --- /dev/null +++ b/tests/remote/conftest.py @@ -0,0 +1,25 @@ +from functools import partial + +import pytest + +import deadpool +from deadpool.remote import DeadpoolClient, DeadpoolServer, UnixAddress, UnixListener +from tests.remote.tasks import multiply + + +@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..c6c7f69 --- /dev/null +++ b/tests/remote/tasks.py @@ -0,0 +1,5 @@ +"""Importable registered tasks used by forkserver worker tests.""" + + +def multiply(left, right): + return left * right diff --git a/tests/remote/test_cancellation.py b/tests/remote/test_cancellation.py new file mode 100644 index 0000000..502872c --- /dev/null +++ b/tests/remote/test_cancellation.py @@ -0,0 +1,186 @@ +import time +from concurrent.futures import CancelledError +from functools import partial + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteQueueTimeout, + ServerState, + UnixAddress, + UnixListener, +) + + +def delayed(value, delay): + time.sleep(delay) + return value + + +def marker(path): + path.write_text("ran") + + +def test_queued_cancellation_prevents_execution(tmp_path): + socket_path = tmp_path / "pool.sock" + 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: + running = client.submit(delayed, "done", 0.3) + queued = client.submit(marker, marker_path) + deadline = time.monotonic() + 2 + while client.get_status(queued) == "UNKNOWN" and time.monotonic() < deadline: + time.sleep(0.01) + assert queued.cancel() + assert queued.cancelled() + assert running.result(timeout=5) == "done" + assert not marker_path.exists() + finally: + 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" + 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(delayed, "done", 0.2, deadpool_group_id="batch") + queued = client.submit(delayed, "late", 0.01, deadpool_group_id="batch") + deadline = time.monotonic() + 3 + while not running.running() and time.monotonic() < deadline: + time.sleep(0.01) + counts = client.cancel_group("batch") + assert counts == { + "cancelled": 1, + "running": 1, + "terminal": 0, + "unknown": 0, + } + deadline = time.monotonic() + 2 + while not queued.done() and time.monotonic() < deadline: + time.sleep(0.01) + assert queued.cancelled() + assert running.result(timeout=5) == "done" + finally: + 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" + 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(delayed, "done", 0.25) + queued = client.submit(delayed, "late", 0.01, deadpool_queue_timeout=0.02) + with pytest.raises(RemoteQueueTimeout): + queued.result(timeout=5) + assert running.result(timeout=5) == "done" + finally: + 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" + 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(delayed, "done", 0.2) + deadline = time.monotonic() + 3 + while not future.running() and time.monotonic() < deadline: + time.sleep(0.01) + assert future.running() + assert not future.cancel() + assert future.result(timeout=5) == "done" + finally: + 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" + 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(delayed, "never", 5) + deadline = time.monotonic() + 3 + while not future.running() and time.monotonic() < deadline: + time.sleep(0.01) + assert future.running() + assert future.cancel_and_kill_if_running() + assert future.done() + assert future.cancelled() + with pytest.raises(CancelledError): + future.result() + assert ( + client.submit(delayed, "recovered", 0.01).result(timeout=5) == "recovered" + ) + finally: + client.shutdown(cancel_futures=True) + server.shutdown(cancel_futures=True, deadline=5) + + +def test_server_shutdown_policy_can_be_strengthened(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)) + future = client.submit(delayed, "never", 5) + deadline = time.monotonic() + 3 + while not future.running() and time.monotonic() < deadline: + time.sleep(0.01) + server.shutdown(wait=False, cancel_futures=False, deadline=5) + server.shutdown(wait=False, cancel_futures=True, deadline=0.05) + deadline = time.monotonic() + 3 + while server.state != ServerState.STOPPED and time.monotonic() < deadline: + time.sleep(0.01) + assert server.state == ServerState.STOPPED + assert future.done() + client.shutdown() + + +def test_nonblocking_client_and_server_shutdown(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)) + future = client.submit(delayed, "done", 0.2) + + started = time.monotonic() + client.shutdown(wait=False) + assert time.monotonic() - started < 0.15 + assert future.result(timeout=5) == "done" + + started = time.monotonic() + server.shutdown(wait=False, deadline=5) + assert time.monotonic() - started < 0.15 + deadline = time.monotonic() + 5 + while server.state != ServerState.STOPPED and time.monotonic() < deadline: + time.sleep(0.01) + assert server.state == ServerState.STOPPED + assert not socket_path.exists() 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..6ecaf2c --- /dev/null +++ b/tests/remote/test_integration.py @@ -0,0 +1,337 @@ +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 MessageType +from deadpool.remote.serializer import PickleSerializer + + +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 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])) + + +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) + + +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_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 + client._socket.shutdown(socket.SHUT_RDWR) + client._socket.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_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..be30bf2 --- /dev/null +++ b/tests/remote/test_protocol.py @@ -0,0 +1,80 @@ +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, + send_message, +) + + +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_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) diff --git a/tests/remote/test_tcp.py b/tests/remote/test_tcp.py new file mode 100644 index 0000000..cdb38f7 --- /dev/null +++ b/tests/remote/test_tcp.py @@ -0,0 +1,129 @@ +from functools import partial +import socket +import ssl +import time + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + RemoteCompatibilityError, + RemoteExecutorError, + RemoteExecutorUnavailable, + RemoteLimits, + TcpAddress, + TcpListener, +) + + +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_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_plaintext_tcp_requires_explicit_loopback_opt_in(): + with pytest.raises(ValueError): + TcpAddress("127.0.0.1", 1) + with pytest.raises(ValueError): + TcpListener("0.0.0.0", 0, insecure=True) diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py new file mode 100644 index 0000000..a6a6138 --- /dev/null +++ b/tests/remote/test_units.py @@ -0,0 +1,35 @@ +import pytest + +from deadpool.remote import ( + AcceptanceCertainty, + ExecutionCertainty, + PickleSerializer, + RemoteQueueFull, + SerializationLimitError, +) +from deadpool.remote._scheduler import FairScheduler + + +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"] 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 From 5df5e614efa9031d8cd1607b7ceec5516409e880 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Wed, 29 Jul 2026 16:39:27 +0200 Subject: [PATCH 02/24] Fix lint --- tests/remote/conftest.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/remote/conftest.py b/tests/remote/conftest.py index b4a92f9..a9cc296 100644 --- a/tests/remote/conftest.py +++ b/tests/remote/conftest.py @@ -22,4 +22,3 @@ def remote_pair(tmp_path): finally: client.shutdown(wait=True, cancel_futures=True) server.shutdown(wait=True, cancel_futures=True, deadline=5) - From f3a03c480dbdb7a9448f04c7c1b5bb25ae311b46 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Wed, 29 Jul 2026 22:57:15 +0200 Subject: [PATCH 03/24] fix: harden remote executor concurrency and coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Big picture Strengthen the experimental remote executor with realistic behavioral coverage across its public client, server, and Future interfaces. The expanded suite exercises process, network, TLS, scheduling, cancellation, shutdown, and Unix socket lifecycles rather than targeting individual implementation lines merely to improve the percentage. The new scenarios exposed several real races and lifecycle defects. Fix those issues alongside the tests that reproduce them, and configure Coverage.py to attribute work performed by multiprocessing, forked, and spawned children correctly. Branch-aware coverage rises from roughly 79% to 88%, while statement coverage rises above 90%. ## Changes - Make submission registration atomic with transport failure and shutdown so a Future cannot be stranded after sender exit or concurrent client closure. - Resolve control RPCs exactly once so a valid response cannot be overwritten by a subsequent disconnect, and ensure failed queueing releases RPC bookkeeping. - Make all concurrent server startup callers observe the same readiness or startup failure, preserve failed startup state, and avoid request-table mutation while applying disconnect cancellation. - Correct signed 64-bit protocol validation and reject inconsistent empty group identifiers before submission. - Harden Unix listener cleanup with inode/device identity checks so failed bind and shutdown paths preserve visible replacement sockets and unrelated files. - Configure coverage collection through `.coveragerc` for threads, multiprocessing, raw forks, subprocesses, and `_exit`, erase stale data before collection, and require a Coverage.py version that supports process patches. - Add public end-to-end scenarios for spawned clients sharing one server pool, mutual TLS authentication and certificate rejection, chunked payloads, graceful draining, worker death and recovery, admission limits, scheduling fairness, disconnect policies, callbacks, and control/status operations. - Add focused protocol, serializer, scheduler, configuration, and Future tests where low-level resource and state-machine invariants cannot be exercised reliably through the public API. - Add test-only TLS credentials and document that they are fixtures with no production trust. - Replace timing-sensitive cancellation assertions with explicit barriers, bounded polling, and robust child-process cleanup. ## Validation - `nox -s test-3.14` — 172 passed. - `nox -s testcov-3.14` — 172 passed; 88% branch-aware coverage. - `nox -s test-3.10 -- tests/remote -q` — 105 passed. - Ruff, Black, and `git diff --check` passed. - Independent concurrency and test-quality review approved the final changes. --- .coveragerc | 10 +- deadpool/remote/_protocol.py | 4 +- deadpool/remote/_transport.py | 45 ++- deadpool/remote/client.py | 35 +- deadpool/remote/server.py | 25 +- noxfile.py | 6 +- pyproject.toml | 2 +- tests/remote/certs/README.txt | 1 + tests/remote/certs/ca.pem | 19 + tests/remote/certs/client-key.pem | 28 ++ tests/remote/certs/client.pem | 20 + tests/remote/certs/rogue-client-key.pem | 28 ++ tests/remote/certs/rogue-client.pem | 20 + tests/remote/certs/server-key.pem | 28 ++ tests/remote/certs/server.pem | 20 + tests/remote/conftest.py | 44 +++ tests/remote/tasks.py | 45 ++- tests/remote/test_cancellation.py | 253 ++++++++---- tests/remote/test_client_api.py | 428 +++++++++++++++++++++ tests/remote/test_integration.py | 6 +- tests/remote/test_protocol.py | 120 ++++++ tests/remote/test_public_end_to_end.py | 122 ++++++ tests/remote/test_scheduling_resilience.py | 251 ++++++++++++ tests/remote/test_tcp.py | 111 +++++- tests/remote/test_units.py | 121 ++++++ tests/remote/test_unix_lifecycle.py | 296 ++++++++++++++ 26 files changed, 1978 insertions(+), 110 deletions(-) create mode 100644 tests/remote/certs/README.txt create mode 100644 tests/remote/certs/ca.pem create mode 100644 tests/remote/certs/client-key.pem create mode 100644 tests/remote/certs/client.pem create mode 100644 tests/remote/certs/rogue-client-key.pem create mode 100644 tests/remote/certs/rogue-client.pem create mode 100644 tests/remote/certs/server-key.pem create mode 100644 tests/remote/certs/server.pem create mode 100644 tests/remote/test_client_api.py create mode 100644 tests/remote/test_public_end_to_end.py create mode 100644 tests/remote/test_scheduling_resilience.py create mode 100644 tests/remote/test_unix_lifecycle.py 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/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index a6c0815..0dc05fd 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -116,8 +116,8 @@ def _validate_json(value: object, limits: RemoteLimits, depth: int = 0) -> None: elif isinstance(value, bool) or value is None: return elif isinstance(value, int): - if value.bit_length() > 63: - raise RemoteProtocolError("control integer exceeds 64-bit range") + 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") diff --git a/deadpool/remote/_transport.py b/deadpool/remote/_transport.py index 06ce4e0..a7e2fb5 100644 --- a/deadpool/remote/_transport.py +++ b/deadpool/remote/_transport.py @@ -8,6 +8,7 @@ import struct from contextlib import contextmanager from dataclasses import dataclass +from pathlib import Path from .config import PeerInfo, TcpAddress, TcpListener, UnixAddress, UnixListener @@ -29,15 +30,8 @@ def close(self) -> None: current = path.lstat() except FileNotFoundError: current = None - if ( - current is not None - and ( - current.st_dev, - current.st_ino, - ) - == self.unix_identity - ): - path.unlink() + if current is not None: + _unlink_if_identity(path, current, self.unix_identity) def bind_listener(config: UnixListener | TcpListener) -> BoundListener: @@ -102,15 +96,23 @@ def _bind_unix(config: UnixListener) -> BoundListener: probe.settimeout(0.2) probe.connect(str(path)) except (ConnectionRefusedError, FileNotFoundError, socket.timeout): - path.unlink(missing_ok=True) + try: + current = path.lstat() + except FileNotFoundError: + pass + else: + _unlink_if_identity(path, current, (info.st_dev, info.st_ino)) else: raise OSError(f"Unix socket is already live: {path}") finally: probe.close() 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, @@ -119,14 +121,31 @@ def _bind_unix(config: UnixListener) -> BoundListener: ) os.chmod(path, config.mode) sock.listen(config.backlog) - identity = path.stat() - return BoundListener(config, sock, path, (identity.st_dev, identity.st_ino)) + return BoundListener(config, sock, path, identity) except BaseException: sock.close() - path.unlink(missing_ok=True) + if identity is not None: + try: + current = path.lstat() + except FileNotFoundError: + pass + else: + _unlink_if_identity(path, current, identity) raise +def _unlink_if_identity( + path: Path, current: os.stat_result, identity: tuple[int, int] +) -> None: + """Compare identity immediately before unlinking a Unix listener path. + + This protects replacements visible to the final ``lstat``. A portable + ``lstat``/``unlink`` sequence still has an unavoidable TOCTOU window. + """ + if (current.st_dev, current.st_ino) == identity: + path.unlink() + + def _validate_unix_directory(path) -> None: info = path.lstat() if not stat.S_ISDIR(info.st_mode) or stat.S_ISLNK(info.st_mode): diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py index 74ebb5f..9b953ac 100644 --- a/deadpool/remote/client.py +++ b/deadpool/remote/client.py @@ -228,6 +228,13 @@ def _submit( ) 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) @@ -378,9 +385,10 @@ def _receive(self, message: Message) -> None: token = message.control.get("token") or message.control.get("nonce") if isinstance(token, str): with self._lock: - rpc = self._rpc.get(token) + rpc = self._rpc.pop(token, None) + if rpc is not None: + rpc[1].update(message.control) if rpc is not None: - rpc[1].update(message.control) rpc[0].set() elif message.kind == MessageType.GOAWAY: self._connection_lost(RemoteConnectionLost("server closed the session")) @@ -493,11 +501,11 @@ def _cancel(self, future: RemoteFuture, *, hard: bool) -> bool: "connection is unavailable", request_id=future.request_id ) self._rpc[token] = (event, response) - self._enqueue_control( - MessageType.CANCEL_REQUEST, - {"request_id": future.request_id, "hard": hard, "token": token}, - ) 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", @@ -572,10 +580,13 @@ def _rpc_request( raise RemoteConnectionLost("connection is unavailable") self._rpc[token] = (event, response) key = "nonce" if kind == MessageType.PING else "token" - self._enqueue_control(kind, {**control, key: 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: @@ -599,13 +610,14 @@ def _connection_lost(self, error: BaseException) -> None: 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.close() except OSError: pass for event, response in rpc: - response["error"] = str(error) + response["_transport_error"] = error event.set() for future in futures: with self._lock: @@ -821,8 +833,11 @@ def _submission_options(kwargs: dict, default_submission_timeout: float) -> dict values["submission_timeout"] = _positive_timeout( values["submission_timeout"], "submission_timeout" ) - if values["group_id"] is not None and not isinstance(values["group_id"], str): - raise TypeError("deadpool_group_id must be a string or None") + 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 diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index cf00a6f..89bb8de 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -287,12 +287,16 @@ def bound_addresses(self) -> tuple[object, ...]: def start(self) -> "DeadpoolServer": with self._lock: - if self._serve_thread is not None: - return self - self._serve_thread = threading.Thread( - target=self.serve_forever, name="deadpool.remote.server", daemon=False - ) - self._serve_thread.start() + if self._serve_thread is None: + self._serve_thread = threading.Thread( + target=self.serve_forever, + name="deadpool.remote.server", + daemon=False, + ) + self._serve_thread.start() + # Every caller observes the same readiness or startup failure. In + # particular, a concurrent caller must not return merely because the + # serve thread has been created. self.wait_ready(self.limits.handshake_timeout) return self @@ -307,7 +311,9 @@ def serve_forever(self) -> None: try: self._initialize() except BaseException as error: - self._startup_error = error + with self._lock: + self._startup_error = error + self.state = ServerState.STOPPED self._close_bound() self.ready.set() self._stopped.set() @@ -1096,7 +1102,10 @@ def _disconnect(self, connection: _ServerConnection) -> None: and not connection.clean_close and self.disconnect_policy in {"cancel_queued", "terminate"} ): - for record in session.requests.values(): + # 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( 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/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 index a9cc296..99e61f1 100644 --- a/tests/remote/conftest.py +++ b/tests/remote/conftest.py @@ -1,4 +1,7 @@ +from dataclasses import dataclass from functools import partial +from pathlib import Path +import ssl import pytest @@ -6,6 +9,47 @@ 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): diff --git a/tests/remote/tasks.py b/tests/remote/tasks.py index c6c7f69..d55056d 100644 --- a/tests/remote/tasks.py +++ b/tests/remote/tasks.py @@ -1,5 +1,48 @@ """Importable registered tasks used by forkserver worker tests.""" +import os +import time +from pathlib import Path -def multiply(left, right): + +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 index 502872c..ff13a11 100644 --- a/tests/remote/test_cancellation.py +++ b/tests/remote/test_cancellation.py @@ -1,4 +1,6 @@ +import threading import time +from collections.abc import Callable from concurrent.futures import CancelledError from functools import partial @@ -8,179 +10,294 @@ from deadpool.remote import ( DeadpoolClient, DeadpoolServer, + RemoteExecutorError, RemoteQueueTimeout, ServerState, + SubmissionState, UnixAddress, UnixListener, ) - - -def delayed(value, delay): - time.sleep(delay) - return value +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" - marker_path = tmp_path / "marker" + 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(delayed, "done", 0.3) - queued = client.submit(marker, marker_path) - deadline = time.monotonic() + 2 - while client.get_status(queued) == "UNKNOWN" and time.monotonic() < deadline: - time.sleep(0.01) + 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() + release.touch() assert running.result(timeout=5) == "done" - assert not marker_path.exists() + 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(delayed, "done", 0.2, deadpool_group_id="batch") - queued = client.submit(delayed, "late", 0.01, deadpool_group_id="batch") - deadline = time.monotonic() + 3 - while not running.running() and time.monotonic() < deadline: - time.sleep(0.01) - counts = client.cancel_group("batch") - assert counts == { + 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, } - deadline = time.monotonic() + 2 - while not queued.done() and time.monotonic() < deadline: - time.sleep(0.01) + 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(delayed, "done", 0.25) - queued = client.submit(delayed, "late", 0.01, deadpool_queue_timeout=0.02) + 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_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(delayed, "done", 0.2) - deadline = time.monotonic() + 3 - while not future.running() and time.monotonic() < deadline: - time.sleep(0.01) - assert future.running() + 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(delayed, "never", 5) - deadline = time.monotonic() + 3 - while not future.running() and time.monotonic() < deadline: - time.sleep(0.01) - assert future.running() + 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 ( - client.submit(delayed, "recovered", 0.01).result(timeout=5) == "recovered" - ) + 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_policy_can_be_strengthened(tmp_path): +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)) - future = client.submit(delayed, "never", 5) - deadline = time.monotonic() + 3 - while not future.running() and time.monotonic() < deadline: - time.sleep(0.01) - server.shutdown(wait=False, cancel_futures=False, deadline=5) - server.shutdown(wait=False, cancel_futures=True, deadline=0.05) - deadline = time.monotonic() + 3 - while server.state != ServerState.STOPPED and time.monotonic() < deadline: - time.sleep(0.01) - assert server.state == ServerState.STOPPED - assert future.done() - client.shutdown() - - -def test_nonblocking_client_and_server_shutdown(tmp_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)) - future = client.submit(delayed, "done", 0.2) - - started = time.monotonic() - client.shutdown(wait=False) - assert time.monotonic() - started < 0.15 - assert future.result(timeout=5) == "done" - - started = time.monotonic() - server.shutdown(wait=False, deadline=5) - assert time.monotonic() - started < 0.15 - deadline = time.monotonic() + 5 - while server.state != ServerState.STOPPED and time.monotonic() < deadline: - time.sleep(0.01) - assert server.state == ServerState.STOPPED - assert not socket_path.exists() + 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..6628c88 --- /dev/null +++ b/tests/remote/test_client_api.py @@ -0,0 +1,428 @@ +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, + 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) + + +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 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_integration.py b/tests/remote/test_integration.py index 6ecaf2c..a2098ae 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -268,8 +268,10 @@ def omit_result_ack(kind, control): try: assert client.submit(add, 1, 2).result(timeout=5) == 3 assert server.get_statistics()["remote_retained_outcomes"] == 1 - client._socket.shutdown(socket.SHUT_RDWR) - client._socket.close() + 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"] diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index be30bf2..12c67e1 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -1,3 +1,4 @@ +import hashlib import socket import struct import threading @@ -12,7 +13,9 @@ Message, MessageReader, MessageType, + _json_loads, send_message, + validate_control, ) @@ -78,3 +81,120 @@ def test_remote_limits_reject_unbounded_values(): 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_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_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..937a1a5 --- /dev/null +++ b/tests/remote/test_scheduling_resilience.py @@ -0,0 +1,251 @@ +import socket +import time +from functools import partial +from pathlib import Path +from typing import Callable + +import pytest + +import deadpool +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_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 index cdb38f7..8adba88 100644 --- a/tests/remote/test_tcp.py +++ b/tests/remote/test_tcp.py @@ -2,7 +2,6 @@ import socket import ssl import time - import pytest import deadpool @@ -12,10 +11,12 @@ RemoteCompatibilityError, RemoteExecutorError, RemoteExecutorUnavailable, + Principal, RemoteLimits, TcpAddress, TcpListener, ) +from tests.remote.tasks import multiply def identity(value): @@ -37,6 +38,99 @@ def test_insecure_loopback_tcp_uses_same_protocol(): 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"), @@ -122,6 +216,21 @@ def test_tls_configuration_requires_mutual_verification(): 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) diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py index a6a6138..48f1903 100644 --- a/tests/remote/test_units.py +++ b/tests/remote/test_units.py @@ -4,10 +4,14 @@ 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(): @@ -33,3 +37,120 @@ def test_scheduler_is_strict_priority_and_principal_fair(): scheduler.put("urgent", priority=1, principal="a") assert [scheduler.pop() for _ in range(4)] == ["urgent", "a1", "b1", "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 + + def _reserve_callback(self): + return None + + def _schedule_callbacks(self, callbacks, future): + for callback in callbacks: + callback(future) + + def _schedule_callback(self, callback, future): + 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_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] diff --git a/tests/remote/test_unix_lifecycle.py b/tests/remote/test_unix_lifecycle.py new file mode 100644 index 0000000..7f40f62 --- /dev/null +++ b/tests/remote/test_unix_lifecycle.py @@ -0,0 +1,296 @@ +import errno +import socket +import stat +import threading +from pathlib import Path + +import pytest + +import deadpool +from deadpool.remote import ( + DeadpoolClient, + DeadpoolServer, + 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_concurrent_start_waits_for_shared_readiness(tmp_path: Path) -> None: + path = tmp_path / "pool.sock" + server = DeadpoolServer(pool_factory, listeners=[UnixListener(path)]) + 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) + 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_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() + calls = 0 + + def socket_factory(*args: object, **kwargs: object) -> object: + nonlocal calls + calls += 1 + return probe if calls == 1 else 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 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() From 18696be7103e19f714ac5dd6e64c8ce491b32580 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Wed, 29 Jul 2026 23:22:54 +0200 Subject: [PATCH 04/24] Fix CI tests --- deadpool/remote/_transport.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/deadpool/remote/_transport.py b/deadpool/remote/_transport.py index a7e2fb5..2cc134f 100644 --- a/deadpool/remote/_transport.py +++ b/deadpool/remote/_transport.py @@ -2,6 +2,7 @@ from __future__ import annotations +import errno import os import socket import stat @@ -101,7 +102,12 @@ def _bind_unix(config: UnixListener) -> BoundListener: except FileNotFoundError: pass else: - _unlink_if_identity(path, current, (info.st_dev, info.st_ino)) + if not _unlink_if_identity(path, current, (info.st_dev, info.st_ino)): + raise OSError( + errno.EADDRINUSE, + os.strerror(errno.EADDRINUSE), + str(path), + ) else: raise OSError(f"Unix socket is already live: {path}") finally: @@ -136,14 +142,17 @@ def _bind_unix(config: UnixListener) -> BoundListener: def _unlink_if_identity( path: Path, current: os.stat_result, identity: tuple[int, int] -) -> None: - """Compare identity immediately before unlinking a Unix listener path. +) -> bool: + """Remove a Unix listener path only when it retains ``identity``. - This protects replacements visible to the final ``lstat``. A portable - ``lstat``/``unlink`` sequence still has an unavoidable TOCTOU window. + 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: - path.unlink() + if (current.st_dev, current.st_ino) != identity: + return False + path.unlink() + return True def _validate_unix_directory(path) -> None: From 9c62ae5d194069144809617570bfbbb926d60abd Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Wed, 29 Jul 2026 23:35:27 +0200 Subject: [PATCH 05/24] Fix CI tests --- deadpool/remote/_transport.py | 48 ++++++++++++++++++----------- tests/remote/test_unix_lifecycle.py | 21 ++++++++++++- 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/deadpool/remote/_transport.py b/deadpool/remote/_transport.py index 2cc134f..f80140f 100644 --- a/deadpool/remote/_transport.py +++ b/deadpool/remote/_transport.py @@ -92,26 +92,21 @@ def _bind_unix(config: UnixListener) -> BoundListener: raise ValueError(f"refusing unexpected Unix socket path type: {path}") if config.stale_policy != "force_unlink": raise FileExistsError(path) - probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + 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: - probe.settimeout(0.2) - probe.connect(str(path)) - except (ConnectionRefusedError, FileNotFoundError, socket.timeout): - try: - current = path.lstat() - except FileNotFoundError: - pass - else: - if not _unlink_if_identity(path, current, (info.st_dev, info.st_ino)): - raise OSError( - errno.EADDRINUSE, - os.strerror(errno.EADDRINUSE), - str(path), - ) + current = path.lstat() + except FileNotFoundError: + pass else: - raise OSError(f"Unix socket is already live: {path}") - finally: - probe.close() + 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: @@ -140,6 +135,23 @@ def _bind_unix(config: UnixListener) -> BoundListener: 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: diff --git a/tests/remote/test_unix_lifecycle.py b/tests/remote/test_unix_lifecycle.py index 7f40f62..35c633f 100644 --- a/tests/remote/test_unix_lifecycle.py +++ b/tests/remote/test_unix_lifecycle.py @@ -186,17 +186,36 @@ 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 - return probe if calls == 1 else real_socket(*args, **kwargs) + 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: From 73892e2860aa3b4ecba9519a0bbc0a9967731515 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 00:38:47 +0200 Subject: [PATCH 06/24] fix: retain live remote sessions until disconnect --- deadpool/remote/server.py | 4 ---- tests/remote/test_integration.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index 89bb8de..d294a00 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -1156,10 +1156,6 @@ def _remove_request_locked(self, record: _Request) -> None: record.terminal_payload = b"" session.requests.pop(record.request_id, None) self._requests.pop(record.request_id, None) - 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 _purge_session_locked(self, session: _Session) -> None: for record in list(session.requests.values()): diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index a2098ae..833bd34 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -245,6 +245,35 @@ def omit_result_ack(kind, control): 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( From fd4b4e0f33891ba42e94f62c0e55edafb7827d6f Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 00:52:14 +0200 Subject: [PATCH 07/24] fix: arbitrate remote server startup shutdown --- deadpool/remote/server.py | 148 ++++++++++++++++++-------- tests/remote/test_unix_lifecycle.py | 155 +++++++++++++++++++++++++++- 2 files changed, 257 insertions(+), 46 deletions(-) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index d294a00..4fdce2c 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -47,6 +47,7 @@ class ServerState(str, Enum): CREATED = "CREATED" + STARTING = "STARTING" RUNNING = "RUNNING" DRAINING = "DRAINING" STOPPING = "STOPPING" @@ -288,16 +289,26 @@ def bound_addresses(self) -> tuple[object, ...]: 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, ) - self._serve_thread.start() - # Every caller observes the same readiness or startup failure. In - # particular, a concurrent caller must not return merely because the - # serve thread has been created. - self.wait_ready(self.limits.handshake_timeout) + 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: @@ -305,9 +316,19 @@ def wait_ready(self, timeout: float | None = None) -> bool: raise TimeoutError("remote server did not become ready") if self._startup_error is not None: raise self._startup_error - return self.state == ServerState.RUNNING + 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: @@ -317,14 +338,19 @@ def serve_forever(self) -> None: self._close_bound() self.ready.set() self._stopped.set() - if threading.current_thread() is not self._serve_thread: + if not background_start: raise return self._stopped.wait() def _initialize(self) -> None: - with self._lock: - if self.state != ServerState.CREATED: + 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] = [] @@ -338,28 +364,49 @@ def _initialize(self) -> None: if not bound: raise OSError("no listeners could be bound") except BaseException: - for item in bound: - item.close() - try: - pool.shutdown(wait=True, cancel_futures=True) - except BaseException: - pass + self._discard_startup_resources(pool, bound) raise - with self._lock: - self._pool = pool - self._bound = bound - self.state = ServerState.RUNNING - threading.Thread( - target=self._dispatch_loop, name="deadpool.remote.broker", daemon=True - ).start() + 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: - threading.Thread( - target=self._accept_loop, - args=(item,), - name="deadpool.remote.accept", - daemon=True, - ).start() - self.ready.set() + 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}: @@ -1191,30 +1238,41 @@ def shutdown( 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.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() + 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 wait and self._shutdown_thread is not threading.current_thread(): + 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: diff --git a/tests/remote/test_unix_lifecycle.py b/tests/remote/test_unix_lifecycle.py index 35c633f..48b41cf 100644 --- a/tests/remote/test_unix_lifecycle.py +++ b/tests/remote/test_unix_lifecycle.py @@ -2,6 +2,7 @@ import socket import stat import threading +import time from pathlib import Path import pytest @@ -10,6 +11,7 @@ from deadpool.remote import ( DeadpoolClient, DeadpoolServer, + RemoteLimits, ServerState, UnixAddress, UnixListener, @@ -57,9 +59,46 @@ def test_failed_startup_is_stopped_repeatable_and_non_destructive( assert path.is_symlink() if path_kind == "symlink" else path.read_text() == "keep" -def test_concurrent_start_waits_for_shared_readiness(tmp_path: Path) -> None: +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() @@ -88,6 +127,8 @@ def start_server(returned: threading.Event | None = None) -> None: 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) @@ -104,6 +145,118 @@ def start_server(returned: threading.Event | None = None) -> None: 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: From 17ec6218bb078a37a529dfe2975de732fdb2969a Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 00:53:21 +0200 Subject: [PATCH 08/24] fix: restrict insecure TCP clients to loopback --- deadpool/remote/config.py | 7 ++++++- tests/remote/test_tcp.py | 4 +++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/deadpool/remote/config.py b/deadpool/remote/config.py index 424653b..ada0d43 100644 --- a/deadpool/remote/config.py +++ b/deadpool/remote/config.py @@ -9,6 +9,9 @@ from typing import Protocol, runtime_checkable +_INSECURE_LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost"}) + + @dataclass(frozen=True, slots=True) class UnixAddress: path: str | Path @@ -38,6 +41,8 @@ def __post_init__(self) -> None: 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) @@ -91,7 +96,7 @@ def __post_init__(self) -> None: 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 {"127.0.0.1", "localhost"}: + if self.insecure and self.host not in _INSECURE_LOOPBACK_HOSTS: raise ValueError("insecure TCP listeners are restricted to loopback") diff --git a/tests/remote/test_tcp.py b/tests/remote/test_tcp.py index 8adba88..df47010 100644 --- a/tests/remote/test_tcp.py +++ b/tests/remote/test_tcp.py @@ -234,5 +234,7 @@ def test_tls_configuration_rejects_versions_older_than_tls_1_2(): def test_plaintext_tcp_requires_explicit_loopback_opt_in(): with pytest.raises(ValueError): TcpAddress("127.0.0.1", 1) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="addresses are restricted to loopback"): + TcpAddress("example.com", 443, insecure=True) + with pytest.raises(ValueError, match="listeners are restricted to loopback"): TcpListener("0.0.0.0", 0, insecure=True) From 77429192a47b0d90ecaa7e54e894cd060cbcdf34 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 00:54:54 +0200 Subject: [PATCH 09/24] fix: notify stdlib waiters of remote cancellation --- deadpool/remote/_future.py | 11 +++++++++-- tests/remote/test_cancellation.py | 6 +++++- tests/remote/test_units.py | 6 ++++++ 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/deadpool/remote/_future.py b/deadpool/remote/_future.py index 76951f0..eec6d61 100644 --- a/deadpool/remote/_future.py +++ b/deadpool/remote/_future.py @@ -52,7 +52,7 @@ def cancel(self) -> bool: return False if self.submission_state == SubmissionState.LOCAL_PENDING: self.submission_state = SubmissionState.CANCELLED - cancelled = super().cancel() + cancelled = self._cancel_and_notify_waiters() self._dispatch_callbacks() return cancelled client = self._client() @@ -148,9 +148,16 @@ def _set_cancelled(self) -> None: if self.done(): return self.submission_state = SubmissionState.CANCELLED - super().cancel() + 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, [] diff --git a/tests/remote/test_cancellation.py b/tests/remote/test_cancellation.py index ff13a11..ee84452 100644 --- a/tests/remote/test_cancellation.py +++ b/tests/remote/test_cancellation.py @@ -1,7 +1,7 @@ import threading import time from collections.abc import Callable -from concurrent.futures import CancelledError +from concurrent.futures import CancelledError, as_completed, wait from functools import partial import pytest @@ -69,6 +69,10 @@ def test_queued_cancellation_prevents_execution(tmp_path): 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" diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py index 48f1903..53bb959 100644 --- a/tests/remote/test_units.py +++ b/tests/remote/test_units.py @@ -1,3 +1,5 @@ +from concurrent.futures import as_completed, wait + import pytest from deadpool.remote import ( @@ -154,3 +156,7 @@ def test_future_local_cancellation_never_contacts_server(): 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] From 3689c9cc6ad88b804cb8f0866279a4f6238fc9d9 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 01:04:34 +0200 Subject: [PATCH 10/24] fix: reject incompatible remote wire limits --- README.rst | 4 ++ deadpool/remote/_protocol.py | 20 ++++++++ deadpool/remote/client.py | 15 +++++- deadpool/remote/server.py | 20 +++++++- tests/remote/test_integration.py | 78 +++++++++++++++++++++++++++++++- 5 files changed, 134 insertions(+), 3 deletions(-) diff --git a/README.rst b/README.rst index a84870f..53af1c3 100644 --- a/README.rst +++ b/README.rst @@ -809,6 +809,10 @@ 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. + .. warning:: The default pickle callable mode is remote code execution by design. TLS diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index 0dc05fd..a171e7f 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -27,6 +27,26 @@ MINOR = 0 _PREFIX = struct.Struct("!4sBBBBII") _FLAG_CHUNKED = 1 +_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): diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py index 9b953ac..0d7f22e 100644 --- a/deadpool/remote/client.py +++ b/deadpool/remote/client.py @@ -25,6 +25,8 @@ Message, MessageReader, MessageType, + _validate_wire_limits, + _wire_limits, send_message, validate_control, ) @@ -146,6 +148,7 @@ def _connect(self) -> None: "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) @@ -167,6 +170,13 @@ def _connect(self) -> None: 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: @@ -186,7 +196,10 @@ def _connect(self) -> None: 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 {}) + 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 diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index 4fdce2c..02f42fe 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -21,7 +21,16 @@ from deadpool import Future as LocalFuture from deadpool import ProcessError, TimeoutError -from ._protocol import MAJOR, MINOR, Message, MessageReader, MessageType, send_message +from ._protocol import ( + MAJOR, + MINOR, + Message, + MessageReader, + MessageType, + _validate_wire_limits, + _wire_limits, + send_message, +) from ._scheduler import FairScheduler from ._transport import ( BoundListener, @@ -496,6 +505,12 @@ def _connection_loop(self, connection: _ServerConnection) -> None: 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") @@ -603,6 +618,7 @@ def _handshake(self, connection: _ServerConnection, hello: dict) -> None: "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, @@ -1379,6 +1395,7 @@ def _validate_hello(hello: dict) -> None: "capabilities", "max_result_bytes", "wire", + "wire_limits", } if not required <= set(hello): raise RemoteProtocolError("HELLO is missing required fields") @@ -1421,6 +1438,7 @@ def _validate_hello(hello: dict) -> None: 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) diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index 833bd34..4a5c5a2 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -22,7 +22,13 @@ UnixAddress, UnixListener, ) -from deadpool.remote._protocol import MessageType +from deadpool.remote._protocol import ( + Message, + MessageReader, + MessageType, + _wire_limits, + send_message, +) from deadpool.remote.serializer import PickleSerializer @@ -108,6 +114,76 @@ def test_duplicate_live_client_instance_is_rejected(remote_pair, monkeypatch): 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( From f6063f7be98790f6ecf54286918daedee3c8248c Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 01:22:01 +0200 Subject: [PATCH 11/24] fix: bound remote error control descriptors --- deadpool/remote/_protocol.py | 76 +++++++++++++++++++++++------ deadpool/remote/_worker.py | 7 +-- deadpool/remote/server.py | 82 +++++++++++++++++++++++++++++++- tests/remote/test_integration.py | 44 +++++++++++++++++ 4 files changed, 188 insertions(+), 21 deletions(-) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index a171e7f..07af9d9 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -173,6 +173,58 @@ def validate_control(control: dict, limits: RemoteLimits) -> None: _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, @@ -180,26 +232,22 @@ def send_message( ) -> None: """Send one logical message as one or more independently bounded frames.""" payload = memoryview(message.payload) - if len(payload) > limits.max_message_bytes: - raise RemoteProtocolError("message payload exceeds configured limit") frame_size = limits.max_frame_payload_bytes - count = max(1, (len(payload) + frame_size - 1) // frame_size) - if count > limits.max_chunks: - raise RemoteProtocolError("message requires too many chunks") + 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] - frame_control = { - "message_id": message_id, - "index": index, - "count": count, - "total": len(payload), - "digest": digest, - "control": message.control if index == 0 else None, - } - control = _json_dumps(frame_control, limits) + 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, diff --git a/deadpool/remote/_worker.py b/deadpool/remote/_worker.py index b692297..34838ac 100644 --- a/deadpool/remote/_worker.py +++ b/deadpool/remote/_worker.py @@ -54,12 +54,7 @@ def execute_opaque( except BaseException as error: return WorkerOutcome( "result_encoding_failed", - descriptor={ - "module": type(error).__module__, - "type": type(error).__qualname__, - "message": _safe_repr(error), - "traceback": "".join(traceback.format_exception(error)), - }, + descriptor=_describe_exception(error), ) return WorkerOutcome("result", payload) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index 02f42fe..c3a447c 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -30,6 +30,7 @@ _validate_wire_limits, _wire_limits, send_message, + validate_message_control, ) from ._scheduler import FairScheduler from ._transport import ( @@ -940,8 +941,14 @@ def _complete_local(self, record: _Request, outcome: object) -> None: ), } 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, outcome.descriptor or {}, outcome.payload + record, state, kind, descriptor, outcome.payload ) elif isinstance(outcome, TimeoutError): self._terminal_locked( @@ -1372,6 +1379,79 @@ def __exit__(self, exc_type, exc, traceback) -> bool: 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 diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index 4a5c5a2..6857815 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -45,6 +45,14 @@ 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 @@ -239,6 +247,42 @@ def dumps(self, value, *, limit): 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( From b92e6543428a5c1a0965c4a6a3158cb5af4c6cd3 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 01:40:19 +0200 Subject: [PATCH 12/24] fix: make remote callback registration nonblocking --- README.rst | 5 ++ deadpool/remote/_future.py | 25 ++++----- deadpool/remote/client.py | 79 ++++++++++++++++------------ deadpool/remote/config.py | 3 +- tests/remote/test_integration.py | 88 ++++++++++++++++++++++++++++++++ tests/remote/test_units.py | 3 -- 6 files changed, 149 insertions(+), 54 deletions(-) diff --git a/README.rst b/README.rst index 53af1c3..4c72913 100644 --- a/README.rst +++ b/README.rst @@ -813,6 +813,11 @@ 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 the standard nonblocking +registration contract. ``callback_queue_size`` bounds callback batches admitted +to the callback executor, not application-owned registrations waiting for +dispatch; callback backpressure never blocks transport or Future state updates. + .. warning:: The default pickle callable mode is remote code execution by design. TLS diff --git a/deadpool/remote/_future.py b/deadpool/remote/_future.py index eec6d61..049edf4 100644 --- a/deadpool/remote/_future.py +++ b/deadpool/remote/_future.py @@ -38,7 +38,7 @@ def __init__(self, request_id: str, client: DeadpoolClient) -> None: self._owner_pid = os.getpid() self._client = weakref.ref(client) self._state_lock = threading.RLock() - self._remote_callbacks: list[tuple[object, bool]] = [] + self._remote_callbacks: list[object] = [] @property def pid(self) -> int | None: @@ -69,15 +69,11 @@ def cancel_and_kill_if_running(self) -> bool: def add_done_callback(self, fn) -> None: self._check_process() - client = self._client() - reserved = client is not None - if client is not None: - client._reserve_callback() with self._state_lock: if not self.done(): - self._remote_callbacks.append((fn, reserved)) + self._remote_callbacks.append(fn) return - self._schedule_callback(fn, reserved) + self._schedule_callback(fn) def result(self, timeout: float | None = None): self._check_process() @@ -164,18 +160,15 @@ def _dispatch_callbacks(self) -> None: if not callbacks: return client = self._client() - if client is not None and all(reserved for _, reserved in callbacks): - client._schedule_callbacks( - [callback for callback, _ in callbacks], - self, - ) + if client is None: + for callback in callbacks: + callback(self) return - for callback, reserved in callbacks: - self._schedule_callback(callback, reserved) + client._schedule_callbacks(callbacks, self) - def _schedule_callback(self, callback, reserved: bool) -> None: + def _schedule_callback(self, callback) -> None: client = self._client() - if client is None or not reserved: + if client is None: callback(self) else: client._schedule_callback(callback, self) diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py index 0d7f22e..0acc8c9 100644 --- a/deadpool/remote/client.py +++ b/deadpool/remote/client.py @@ -123,11 +123,13 @@ def __init__( 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: @@ -681,33 +683,47 @@ def _forget(self, future: RemoteFuture) -> None: with self._lock: self._futures.pop(future.request_id, None) - def _reserve_callback(self) -> None: - self._callback_slots.acquire() + 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: + _call_callback(callback, future) def _schedule_callback(self, callback, future: RemoteFuture) -> None: - try: - self._callback_executor.submit( - _run_reserved_callback, - self._callback_slots, - callback, - future, - ) - except RuntimeError: - self._callback_slots.release() - _call_callback(callback, future) + self._schedule_callbacks([callback], future) def _schedule_callbacks(self, callbacks: list, future: RemoteFuture) -> None: - try: - self._callback_executor.submit( - _run_callbacks, - self._callback_slots, - callbacks, - future, - ) - except RuntimeError: - for callback in callbacks: - self._callback_slots.release() - _call_callback(callback, future) + with self._lock: + if not self._callback_dispatch_closed: + self._callback_dispatch_queue.put((callbacks, future)) + return + for callback in callbacks: + _call_callback(callback, future) def _ensure_process(self) -> None: if os.getpid() == self._owner_pid: @@ -753,6 +769,7 @@ def _ensure_process(self) -> None: 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() @@ -811,6 +828,9 @@ def _finish_shutdown(self, futures: list[RemoteFuture]) -> None: 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) @@ -881,25 +901,16 @@ def _run_bounded(semaphore, function, args: tuple, kwargs: dict) -> None: semaphore.release() -def _run_callbacks( +def _run_callback_batch( semaphore: threading.BoundedSemaphore, callbacks: list, future: RemoteFuture, ) -> None: + semaphore.release() for callback in callbacks: - semaphore.release() _call_callback(callback, future) -def _run_reserved_callback( - semaphore: threading.BoundedSemaphore, - callback, - future: RemoteFuture, -) -> None: - semaphore.release() - _call_callback(callback, future) - - def _call_callback(callback, future: RemoteFuture) -> None: try: callback(future) diff --git a/deadpool/remote/config.py b/deadpool/remote/config.py index ada0d43..4610298 100644 --- a/deadpool/remote/config.py +++ b/deadpool/remote/config.py @@ -102,7 +102,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True) class RemoteLimits: - """Finite defaults bound every user-controlled queue and byte buffer.""" + """Finite defaults bound protocol/resource queues and byte buffers.""" max_control_bytes: int = 64 * 1024 max_frame_payload_bytes: int = 1024 * 1024 @@ -125,6 +125,7 @@ class RemoteLimits: 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 diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index 6857815..56e602a 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -30,6 +30,7 @@ send_message, ) from deadpool.remote.serializer import PickleSerializer +from tests.remote.tasks import wait_then_mark def add(left, right): @@ -458,6 +459,93 @@ def first_callback(done): 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 + 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() diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py index 53bb959..3866f24 100644 --- a/tests/remote/test_units.py +++ b/tests/remote/test_units.py @@ -109,9 +109,6 @@ class CallbackClient: def __init__(self): self.cancel_calls = 0 - def _reserve_callback(self): - return None - def _schedule_callbacks(self, callbacks, future): for callback in callbacks: callback(future) From 0e9c7381d5f9660801ff5b8a88900da08cb3c520 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 01:44:22 +0200 Subject: [PATCH 13/24] style: format remote feedback fixes --- deadpool/remote/config.py | 1 - deadpool/remote/server.py | 8 ++------ tests/remote/test_integration.py | 5 ++--- tests/remote/test_unix_lifecycle.py | 5 +---- 4 files changed, 5 insertions(+), 14 deletions(-) diff --git a/deadpool/remote/config.py b/deadpool/remote/config.py index 4610298..86a806e 100644 --- a/deadpool/remote/config.py +++ b/deadpool/remote/config.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import Protocol, runtime_checkable - _INSECURE_LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost"}) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index c3a447c..8dc3965 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -947,9 +947,7 @@ def _complete_local(self, record: _Request, outcome: object) -> None: len(outcome.payload), self.limits, ) - self._terminal_locked( - record, state, kind, descriptor, outcome.payload - ) + self._terminal_locked(record, state, kind, descriptor, outcome.payload) elif isinstance(outcome, TimeoutError): self._terminal_locked( record, @@ -1411,9 +1409,7 @@ def _fit_terminal_descriptor( middle = (low + high) // 2 candidate = _truncate_utf8(original, middle) bounded[name] = candidate - if _terminal_descriptor_fits( - request_id, bounded, payload_size, limits - ): + if _terminal_descriptor_fits(request_id, bounded, payload_size, limits): best = candidate low = middle + 1 else: diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index 56e602a..03c6e40 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -133,9 +133,7 @@ def test_duplicate_live_client_instance_is_rejected(remote_pair, monkeypatch): {"max_chunks": 128}, ], ) -def test_handshake_rejects_asymmetric_wire_limits( - tmp_path, server_limit_changes -): +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"), @@ -486,6 +484,7 @@ def test_callback_registration_does_not_wait_for_dispatch_capacity(tmp_path): def register_callbacks() -> None: try: + def blocking_callback(_) -> None: callback_order.append(0) callback_started.set() diff --git a/tests/remote/test_unix_lifecycle.py b/tests/remote/test_unix_lifecycle.py index 48b41cf..681c597 100644 --- a/tests/remote/test_unix_lifecycle.py +++ b/tests/remote/test_unix_lifecycle.py @@ -234,10 +234,7 @@ def stop_server() -> None: assert pool_created.wait(5) shutdown_thread.start() deadline = time.monotonic() + 2 - while ( - server.state is not ServerState.STOPPING - and time.monotonic() < deadline - ): + 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) From 475696a573905e9fd6db7dce89f8d129fe78a439 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 02:38:12 +0200 Subject: [PATCH 14/24] fix: preserve remote executor contracts ## Big picture Keep the remote executor's public timing, scheduling, and trust guarantees intact under real backpressure and completion races. Remote queue deadlines must remain enforceable when the local pool is saturated, `RemoteFuture` must retain the standard `Future` callback contract, and the default pickle transport must clearly communicate that trust is bidirectional in every submission mode. ## Detailed changes - make broker-to-pool admission nonblocking without changing ordinary local `Deadpool.submit()` backpressure - requeue rejected admissions without losing strict priority, principal fairness, or FIFO ordering within a scheduling stream - arbitrate queue deadlines immediately before worker dispatch so expired work cannot start in an admission race - invoke callbacks registered after terminal completion synchronously while retaining isolated asynchronous dispatch for callbacks registered earlier - centralize callback exception handling and the default `PickleSerializer` trust warning - document that callable and registered-task invocations, results, and exceptions require mutually trusted clients and servers - add regressions for saturated-backlog queue timeouts, scheduler requeue ordering, and late callback timing and logging --- README.rst | 20 ++++++++----- deadpool/_pool.py | 11 +++++-- deadpool/remote/__init__.py | 8 +++-- deadpool/remote/_future.py | 30 ++++++++++++------- deadpool/remote/_scheduler.py | 19 +++++++++++- deadpool/remote/client.py | 34 +++++++-------------- deadpool/remote/serializer.py | 29 ++++++++++++++++-- deadpool/remote/server.py | 44 ++++++++++++++++++++++----- tests/remote/test_cancellation.py | 50 +++++++++++++++++++++++++++++++ tests/remote/test_integration.py | 3 ++ tests/remote/test_units.py | 36 +++++++++++++++++++++- 11 files changed, 226 insertions(+), 58 deletions(-) diff --git a/README.rst b/README.rst index 4c72913..584be4c 100644 --- a/README.rst +++ b/README.rst @@ -813,16 +813,22 @@ 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 the standard nonblocking -registration contract. ``callback_queue_size`` bounds callback batches admitted -to the callback executor, not application-owned registrations waiting for -dispatch; callback backpressure never blocks transport or Future state updates. +``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. .. warning:: - The default pickle callable mode is remote code execution by design. TLS - authenticates and encrypts peers; it does not sandbox Python deserialization - or worker execution. Only authorize mutually trusted clients. + 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 diff --git a/deadpool/_pool.py b/deadpool/_pool.py index 28f0635..b032f68 100644 --- a/deadpool/_pool.py +++ b/deadpool/_pool.py @@ -728,12 +728,16 @@ def _submit_future( 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 without changing local Future - semantics. + 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.") @@ -742,7 +746,8 @@ def _submit_future( priority=priority, item=(fn, args, kwargs, timeout, fut), sequence=next(self._submission_sequence), - ) + ), + block=block, ) self._statistics.tasks_received.increment() return fut diff --git a/deadpool/remote/__init__.py b/deadpool/remote/__init__.py index 5e6d275..dd7525a 100644 --- a/deadpool/remote/__init__.py +++ b/deadpool/remote/__init__.py @@ -2,9 +2,11 @@ Security warning ---------------- -Pickle callable mode intentionally executes code supplied by authenticated -clients. TLS authenticates peers; it does not sandbox Python deserialization or -execution. Use this package only across a mutually trusted boundary. +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, diff --git a/deadpool/remote/_future.py b/deadpool/remote/_future.py index 049edf4..1588c79 100644 --- a/deadpool/remote/_future.py +++ b/deadpool/remote/_future.py @@ -3,17 +3,20 @@ from __future__ import annotations import concurrent.futures +import logging import os import threading import weakref from enum import Enum -from typing import TYPE_CHECKING +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" @@ -38,7 +41,7 @@ def __init__(self, request_id: str, client: DeadpoolClient) -> None: self._owner_pid = os.getpid() self._client = weakref.ref(client) self._state_lock = threading.RLock() - self._remote_callbacks: list[object] = [] + self._remote_callbacks: list[Callable[[RemoteFuture], object]] = [] @property def pid(self) -> int | None: @@ -67,13 +70,20 @@ def cancel_and_kill_if_running(self) -> bool: return False return client._cancel(self, hard=True) - def add_done_callback(self, fn) -> None: + 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._schedule_callback(fn) + self._invoke_callback(fn) def result(self, timeout: float | None = None): self._check_process() @@ -162,16 +172,16 @@ def _dispatch_callbacks(self) -> None: client = self._client() if client is None: for callback in callbacks: - callback(self) + self._invoke_callback(callback) return client._schedule_callbacks(callbacks, self) - def _schedule_callback(self, callback) -> None: - client = self._client() - if client is None: + def _invoke_callback(self, callback: Callable[[RemoteFuture], object]) -> None: + """Invoke one callback through the shared exception boundary.""" + try: callback(self) - else: - client._schedule_callback(callback, self) + except BaseException: + logger.exception("exception calling RemoteFuture callback") def _check_process(self) -> None: if os.getpid() != self._owner_pid: diff --git a/deadpool/remote/_scheduler.py b/deadpool/remote/_scheduler.py index baf1478..fcc00f7 100644 --- a/deadpool/remote/_scheduler.py +++ b/deadpool/remote/_scheduler.py @@ -25,6 +25,20 @@ def __init__(self) -> None: 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() @@ -32,7 +46,10 @@ def put(self, item: T, *, priority: int, principal: str) -> None: queue = bucket.queued[principal] if not queue: bucket.principals.append(principal) - queue.append(item) + if front: + queue.appendleft(item) + else: + queue.append(item) self._size += 1 def pop(self) -> T: diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py index 0acc8c9..762b4e0 100644 --- a/deadpool/remote/client.py +++ b/deadpool/remote/client.py @@ -51,13 +51,18 @@ RemoteTaskError, SubmissionOutcomeUnknown, ) -from .serializer import PickleSerializer, Serializer +from .serializer import PICKLE_TRUST_WARNING, PickleSerializer, Serializer logger = logging.getLogger("deadpool.remote") class DeadpoolClient(concurrent.futures.Executor): - """A process-local client whose submitted work runs in a server-owned pool.""" + """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, @@ -73,10 +78,7 @@ def __init__( ) -> None: self.address = address if serializer is None: - logger.warning( - "Deadpool remote callable mode uses pickle and grants trusted " - "clients code execution as the worker account" - ) + logger.warning(PICKLE_TRUST_WARNING) self.serializer = serializer or PickleSerializer() self.authenticator = authenticator self.application_fingerprint = application_fingerprint @@ -712,10 +714,7 @@ def _callback_dispatcher_loop(self) -> None: except RuntimeError: self._callback_slots.release() for callback in callbacks: - _call_callback(callback, future) - - def _schedule_callback(self, callback, future: RemoteFuture) -> None: - self._schedule_callbacks([callback], future) + future._invoke_callback(callback) def _schedule_callbacks(self, callbacks: list, future: RemoteFuture) -> None: with self._lock: @@ -723,7 +722,7 @@ def _schedule_callbacks(self, callbacks: list, future: RemoteFuture) -> None: self._callback_dispatch_queue.put((callbacks, future)) return for callback in callbacks: - _call_callback(callback, future) + future._invoke_callback(callback) def _ensure_process(self) -> None: if os.getpid() == self._owner_pid: @@ -908,18 +907,7 @@ def _run_callback_batch( ) -> None: semaphore.release() for callback in callbacks: - _call_callback(callback, future) - - -def _call_callback(callback, future: RemoteFuture) -> None: - try: - callback(future) - except BaseException: - import logging - - logging.getLogger("deadpool.remote").exception( - "exception calling RemoteFuture callback" - ) + future._invoke_callback(callback) def _deadpool_version() -> str: diff --git a/deadpool/remote/serializer.py b/deadpool/remote/serializer.py index 70dca83..18efdd1 100644 --- a/deadpool/remote/serializer.py +++ b/deadpool/remote/serializer.py @@ -1,7 +1,10 @@ -"""Payload serializers. +"""Payload serializers for a bidirectional remote execution protocol. -Pickle callable mode is remote code execution by design and is only appropriate -between mutually trusted, compatible Python processes. +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 @@ -9,12 +12,25 @@ 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 @@ -24,6 +40,13 @@ 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) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index 8dc3965..e22f8c9 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -50,7 +50,7 @@ UnixListener, ) from .errors import RemoteProtocolError -from .serializer import PickleSerializer, Serializer +from .serializer import PICKLE_TRUST_WARNING, PickleSerializer, Serializer logger = logging.getLogger("deadpool.remote") @@ -222,8 +222,10 @@ 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. Pickle callable mode grants trusted clients code - execution as the worker account. + 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__( @@ -248,10 +250,7 @@ def __init__( self.pool_factory = pool_factory self.listeners = list(listeners) if serializer is None: - logger.warning( - "Deadpool remote callable mode uses pickle and grants trusted " - "clients code execution as the worker account" - ) + logger.warning(PICKLE_TRUST_WARNING) self.serializer = serializer or PickleSerializer() self.task_registry = dict(task_registry or {}) self.registry_fingerprint = registry_fingerprint @@ -850,7 +849,24 @@ def _dispatch_loop(self) -> None: {}, 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: @@ -886,6 +902,20 @@ def _begin_dispatch(self, record: _Request, pid: int) -> bool: if record.state != RequestState.ACCEPTED_QUEUED: self._condition.release() return False + 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() + self._condition.release() + return False return True def _commit_running(self, record: _Request, pid: int) -> None: diff --git a/tests/remote/test_cancellation.py b/tests/remote/test_cancellation.py index ee84452..a45d3bf 100644 --- a/tests/remote/test_cancellation.py +++ b/tests/remote/test_cancellation.py @@ -197,6 +197,56 @@ def test_queue_timeout_is_distinct_from_execution_timeout(tmp_path): 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" diff --git a/tests/remote/test_integration.py b/tests/remote/test_integration.py index 03c6e40..7989366 100644 --- a/tests/remote/test_integration.py +++ b/tests/remote/test_integration.py @@ -526,6 +526,9 @@ def blocking_callback(_) -> None: 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() diff --git a/tests/remote/test_units.py b/tests/remote/test_units.py index 3866f24..12a79f6 100644 --- a/tests/remote/test_units.py +++ b/tests/remote/test_units.py @@ -41,6 +41,20 @@ def test_scheduler_is_strict_priority_and_principal_fair(): 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") @@ -108,13 +122,15 @@ def test_remote_limit_cross_field_relationships(changes, message): 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): - 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 @@ -142,6 +158,24 @@ def test_future_terminal_state_and_callbacks_are_idempotent(): 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) From dcefa4fe73a797b559fbc797356e54226bca0de8 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 02:44:19 +0200 Subject: [PATCH 15/24] fix: release failed remote result completions --- deadpool/remote/client.py | 27 ++++++++------ tests/remote/test_client_api.py | 65 +++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 11 deletions(-) diff --git a/deadpool/remote/client.py b/deadpool/remote/client.py index 762b4e0..0fb31b6 100644 --- a/deadpool/remote/client.py +++ b/deadpool/remote/client.py @@ -416,7 +416,6 @@ def _receive(self, message: Message) -> None: self._connection_lost(RemoteProtocolError(str(message.control))) def _complete(self, future: RemoteFuture, message: Message) -> None: - acknowledge = False try: if message.kind == MessageType.RESULT: future._set_result(self.serializer.loads(message.payload)) @@ -471,7 +470,6 @@ def _complete(self, future: RemoteFuture, message: Message) -> None: request_id=future.request_id, ) ) - acknowledge = True except BaseException as error: future._set_exception( error @@ -479,15 +477,18 @@ def _complete(self, future: RemoteFuture, message: Message) -> None: else RemoteProtocolError(str(error), request_id=future.request_id) ) finally: - if acknowledge: - try: - self._enqueue_control( - MessageType.RESULT_ACK, {"request_id": future.request_id} - ) - finally: - self._forget(future) - with self._lock: - self._terminal_received.discard(future.request_id) + 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") @@ -629,6 +630,10 @@ def _connection_lost(self, error: BaseException) -> None: 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: diff --git a/tests/remote/test_client_api.py b/tests/remote/test_client_api.py index 6628c88..5b5b9f5 100644 --- a/tests/remote/test_client_api.py +++ b/tests/remote/test_client_api.py @@ -20,6 +20,7 @@ RemoteExecutorUnavailable, RemoteLimits, RemoteProcessError, + RemoteProtocolError, RemoteQueueFull, RemoteResultTooLarge, ServerState, @@ -93,6 +94,67 @@ def dumps(self, value: object, *, limit: int) -> bytes: 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() @@ -185,6 +247,9 @@ def wait(self, timeout: float | None = None) -> bool: return self.release_waiter.wait(timeout) class SocketStub: + def shutdown(self, how: int) -> None: + pass + def close(self) -> None: pass From 3bfa528bf3b07e5ef76af97e3aaac25e87647fba Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 02:54:44 +0200 Subject: [PATCH 16/24] fix: release broker lock before worker pipe writes --- deadpool/remote/server.py | 70 ++++++++++----------- tests/remote/test_scheduling_resilience.py | 73 ++++++++++++++++++++++ 2 files changed, 107 insertions(+), 36 deletions(-) diff --git a/deadpool/remote/server.py b/deadpool/remote/server.py index e22f8c9..5f4571c 100644 --- a/deadpool/remote/server.py +++ b/deadpool/remote/server.py @@ -129,9 +129,7 @@ class _Session: class _PoolFuture(LocalFuture): def __init__(self, server: "DeadpoolServer", record: _Request) -> None: super().__init__() - self._before_submit = lambda pid: server._begin_dispatch(record, pid) - self._after_submit = lambda pid: server._commit_running(record, pid) - self._submit_failed = lambda pid: server._abort_dispatch(record, pid) + self._before_submit = lambda pid: server._claim_dispatch(record, pid) @dataclass(slots=True) @@ -895,33 +893,39 @@ def _expire_queued_locked(self) -> None: if record.local_future is not None: record.local_future.cancel() - def _begin_dispatch(self, record: _Request, pid: int) -> bool: - # Hold the arbitration lock across the small pipe send. Cancellation, - # queue timeout, and dispatch can therefore have exactly one winner. - self._condition.acquire() - if record.state != RequestState.ACCEPTED_QUEUED: - self._condition.release() - return False - 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() - self._condition.release() - return False - return True + def _claim_dispatch(self, record: _Request, pid: int) -> bool: + """Atomically cross the queue/worker boundary before pipe I/O. - def _commit_running(self, record: _Request, pid: int) -> None: - try: - record.state = RequestState.RUNNING - self._statistics["remote_tasks_running"] += 1 + 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, { @@ -930,13 +934,7 @@ def _commit_running(self, record: _Request, pid: int) -> None: "worker_id": f"worker:{pid}", }, ) - finally: - self._condition.release() - - def _abort_dispatch(self, record: _Request, pid: int) -> None: - # The worker pipe did not accept bytes, so the request remains queued - # and Deadpool may retry it with another worker under the same ID. - self._condition.release() + return True def _pool_done(self, record: _Request, future: LocalFuture) -> None: try: diff --git a/tests/remote/test_scheduling_resilience.py b/tests/remote/test_scheduling_resilience.py index 937a1a5..d9e98d8 100644 --- a/tests/remote/test_scheduling_resilience.py +++ b/tests/remote/test_scheduling_resilience.py @@ -1,4 +1,5 @@ import socket +import threading import time from functools import partial from pathlib import Path @@ -7,6 +8,7 @@ import pytest import deadpool +from deadpool._pool import WorkerProcess from deadpool.remote import ( DeadpoolClient, DeadpoolServer, @@ -107,6 +109,77 @@ def authorize(principal: Principal, operation: str, metadata: dict) -> bool: 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( From 07033448c58044a7c8549fb6c3b2ad3a5e403003 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 02:56:08 +0200 Subject: [PATCH 17/24] fix: bound total remote frame write time --- deadpool/remote/_protocol.py | 3 +-- tests/remote/test_protocol.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index 07af9d9..f60b6cf 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -399,7 +399,7 @@ def _bounded_timeout(timeout: float, deadline: float | None) -> float: def _send_exact(sock: socket.socket, data: bytes | memoryview, timeout: float) -> None: - """Write bytes with a monotonic progress deadline instead of unbounded sendall.""" + """Write bytes within one monotonic deadline instead of unbounded sendall.""" view = memoryview(data) sent = 0 deadline = time.monotonic() + timeout @@ -422,7 +422,6 @@ def _send_exact(sock: socket.socket, data: bytes | memoryview, timeout: float) - if count == 0: raise EOFError("connection closed during frame write") sent += count - deadline = time.monotonic() + timeout def _recv_exact( diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index 12c67e1..80aca57 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -14,6 +14,7 @@ MessageReader, MessageType, _json_loads, + _send_exact, send_message, validate_control, ) @@ -64,6 +65,35 @@ def test_partial_frame_has_a_bounded_deadline(): 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) From 3792a4e5c21e93a82e0850b59b6a40353cb4cc31 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:00:10 +0200 Subject: [PATCH 18/24] fix: cap aggregate incomplete remote payloads --- deadpool/remote/_protocol.py | 5 ++++ tests/remote/test_protocol.py | 51 +++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index f60b6cf..20be3c6 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -269,6 +269,7 @@ class MessageReader: def __init__(self, limits: RemoteLimits) -> None: self.limits = limits self._messages: dict[str, _Chunks] = {} + self._incomplete_bytes = 0 def receive( self, @@ -371,14 +372,18 @@ def _accept( 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" diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index 80aca57..723cafd 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -217,6 +217,57 @@ def test_chunk_state_bounds_incomplete_messages_and_conflicting_metadata(): 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): From 3f6d8b34d6ecb67b0259e456a30a2384bf68c84a Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:01:49 +0200 Subject: [PATCH 19/24] docs: clarify remote cancellation timing --- README.rst | 6 ++++++ deadpool/remote/_future.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/README.rst b/README.rst index 584be4c..fa0ef20 100644 --- a/README.rst +++ b/README.rst @@ -820,6 +820,12 @@ 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 diff --git a/deadpool/remote/_future.py b/deadpool/remote/_future.py index 1588c79..a3dd2d5 100644 --- a/deadpool/remote/_future.py +++ b/deadpool/remote/_future.py @@ -49,6 +49,12 @@ def pid(self) -> int | None: 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(): From 7d475f0371a97862d6401d0f265e0074c4c9ac11 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:04:02 +0200 Subject: [PATCH 20/24] fix: normalize recursive control JSON errors --- deadpool/remote/_protocol.py | 2 +- tests/remote/test_protocol.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index 20be3c6..8878dd5 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -110,7 +110,7 @@ def _json_loads(data: bytes, limits: RemoteLimits) -> dict: ValueError(f"non-finite JSON number {value}") ), ) - except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: + 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): diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index 723cafd..9f8ad20 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -190,6 +190,23 @@ def test_chunk_state_machine_rejects_inconsistent_input(header, payload, 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_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() From 561809ef49cf38f7a20b6be033625ed8c0b792ae Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:06:02 +0200 Subject: [PATCH 21/24] fix: require numeric insecure loopback addresses --- deadpool/remote/config.py | 2 +- tests/remote/test_tcp.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/deadpool/remote/config.py b/deadpool/remote/config.py index 86a806e..5a72eba 100644 --- a/deadpool/remote/config.py +++ b/deadpool/remote/config.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Protocol, runtime_checkable -_INSECURE_LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost"}) +_INSECURE_LOOPBACK_HOSTS = frozenset({"127.0.0.1"}) @dataclass(frozen=True, slots=True) diff --git a/tests/remote/test_tcp.py b/tests/remote/test_tcp.py index df47010..8947aa9 100644 --- a/tests/remote/test_tcp.py +++ b/tests/remote/test_tcp.py @@ -236,5 +236,9 @@ def test_plaintext_tcp_requires_explicit_loopback_opt_in(): 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) From 5e2838f6c9f588d6f4db70973940c23b962b4e76 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:07:53 +0200 Subject: [PATCH 22/24] style: format remote feedback fixes --- deadpool/remote/_protocol.py | 7 ++++++- tests/remote/test_protocol.py | 17 +++++++++-------- tests/remote/test_scheduling_resilience.py | 4 +--- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index 8878dd5..1b3fa55 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -110,7 +110,12 @@ def _json_loads(data: bytes, limits: RemoteLimits) -> dict: ValueError(f"non-finite JSON number {value}") ), ) - except (UnicodeDecodeError, json.JSONDecodeError, RecursionError, ValueError) as error: + 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): diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index 9f8ad20..63f59c8 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -249,9 +249,7 @@ def test_chunk_state_bounds_aggregate_incomplete_payload_bytes() -> None: assert ( reader._accept( MessageType.RESULT, - chunk_header( - message_id=message_id, count=2, total=4, digest=digest - ), + chunk_header(message_id=message_id, count=2, total=4, digest=digest), b"xx", ) is None @@ -265,11 +263,14 @@ def test_chunk_state_bounds_aggregate_incomplete_payload_bytes() -> None: ) 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", count=2, total=4, digest=digest), + b"xx", + ) + is None + ) assert completing._accept( MessageType.RESULT, chunk_header( diff --git a/tests/remote/test_scheduling_resilience.py b/tests/remote/test_scheduling_resilience.py index d9e98d8..f5613d9 100644 --- a/tests/remote/test_scheduling_resilience.py +++ b/tests/remote/test_scheduling_resilience.py @@ -134,9 +134,7 @@ def blocked_submit_job(self: WorkerProcess, job: object) -> None: try: assert entered.wait(2) queued = client.submit(multiply, 6, 7) - wait_until( - lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED - ) + wait_until(lambda: queued.submission_state is SubmissionState.ACCEPTED_QUEUED) assert future.cancel_and_kill_if_running() release.set() From eb47b70858226b9fa06728174a0f9c1aecb0bff0 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:27:01 +0200 Subject: [PATCH 23/24] Fix tests --- deadpool/remote/_protocol.py | 31 +++++++++++++++++++++++++++++-- tests/remote/test_protocol.py | 22 ++++++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/deadpool/remote/_protocol.py b/deadpool/remote/_protocol.py index 1b3fa55..8e65353 100644 --- a/deadpool/remote/_protocol.py +++ b/deadpool/remote/_protocol.py @@ -27,6 +27,7 @@ MINOR = 0 _PREFIX = struct.Struct("!4sBBBBII") _FLAG_CHUNKED = 1 +_MAX_JSON_DEPTH = 12 _WIRE_LIMIT_FIELDS = ( "max_control_bytes", "max_frame_payload_bytes", @@ -101,10 +102,36 @@ class _Chunks: 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( - data.decode("utf-8"), + text, object_pairs_hook=_unique_object, parse_constant=lambda value: (_ for _ in ()).throw( ValueError(f"non-finite JSON number {value}") @@ -133,7 +160,7 @@ def _unique_object(items: list[tuple[str, object]]) -> dict: def _validate_json(value: object, limits: RemoteLimits, depth: int = 0) -> None: - if depth > 12: + 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: diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index 63f59c8..a600c81 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -196,6 +196,28 @@ def test_deep_control_json_is_a_typed_protocol_error() -> None: _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: From d6e235420d0b96e24f4cf17441bb3f7ebd33c8e0 Mon Sep 17 00:00:00 2001 From: Caleb Hattingh Date: Thu, 30 Jul 2026 03:30:43 +0200 Subject: [PATCH 24/24] Fix lint --- tests/remote/test_protocol.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/remote/test_protocol.py b/tests/remote/test_protocol.py index a600c81..466286c 100644 --- a/tests/remote/test_protocol.py +++ b/tests/remote/test_protocol.py @@ -214,7 +214,7 @@ def fail_if_called(*args: object, **kwargs: object) -> object: def test_json_nesting_check_ignores_delimiters_in_strings() -> None: - payload = b'{"value":"' + b"\\\"" + b"[" * 20 + b'"}' + payload = b'{"value":"' + b'\\"' + b"[" * 20 + b'"}' assert _json_loads(payload, small_limits()) == {"value": '"' + "[" * 20}