From 8578b8b0db81f2c8acc53b29c33675ab302dab27 Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:55:59 -0400 Subject: [PATCH 1/7] feat: Add the idle RPC and IdleTime message to the wire protocol Add an idle unary RPC returning a new IdleTime message (a double seconds field) to the Worker gRPC service, and re-export IdleTime from wool.protocol alongside the other wire messages. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/proto/wire.proto | 8 ++++++++ wool/src/wool/protocol/__init__.py | 2 ++ wool/src/wool/protocol/_wire.py | 2 ++ 3 files changed, 12 insertions(+) diff --git a/wool/proto/wire.proto b/wool/proto/wire.proto index de36f498..b9dcda09 100644 --- a/wool/proto/wire.proto +++ b/wool/proto/wire.proto @@ -86,6 +86,7 @@ message Nack { service Worker { rpc dispatch (stream Request) returns (stream Response); rpc stop (StopRequest) returns (Void); + rpc idle (Void) returns (IdleTime); } message Request { @@ -117,6 +118,13 @@ message StopRequest { float timeout = 1; } +message IdleTime { + // Seconds the worker has been continuously idle: time since the + // in-flight task set last emptied (startup counts as empty). + // 0 while any task is in flight; measured on a monotonic clock. + double seconds = 1; +} + message WorkerMetadata { // Unique identifier for the worker instance string uid = 1; diff --git a/wool/src/wool/protocol/__init__.py b/wool/src/wool/protocol/__init__.py index 4029d5a9..81ec5b00 100644 --- a/wool/src/wool/protocol/__init__.py +++ b/wool/src/wool/protocol/__init__.py @@ -13,6 +13,7 @@ from wool.protocol._wire import ChainManifest as ChainManifest from wool.protocol._wire import ChannelOptions as ChannelOptions from wool.protocol._wire import ContextVar as ContextVar +from wool.protocol._wire import IdleTime as IdleTime from wool.protocol._wire import Message as Message from wool.protocol._wire import Nack as Nack from wool.protocol._wire import Request as Request @@ -35,6 +36,7 @@ "ChainManifest", "ChannelOptions", "ContextVar", + "IdleTime", "Message", "Nack", "Request", diff --git a/wool/src/wool/protocol/_wire.py b/wool/src/wool/protocol/_wire.py index 1c551e32..d9874b42 100644 --- a/wool/src/wool/protocol/_wire.py +++ b/wool/src/wool/protocol/_wire.py @@ -14,6 +14,7 @@ from wool.protocol.wire_pb2 import ChainManifest from wool.protocol.wire_pb2 import ChannelOptions from wool.protocol.wire_pb2 import ContextVar + from wool.protocol.wire_pb2 import IdleTime from wool.protocol.wire_pb2 import Message from wool.protocol.wire_pb2 import Nack from wool.protocol.wire_pb2 import Request @@ -57,6 +58,7 @@ def __call__(servicer, server) -> None: ... "ChainManifest", "ChannelOptions", "ContextVar", + "IdleTime", "Message", "Nack", "Request", From bf722d75703781bef00b9ce7c3c554e271755b9f Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:55:59 -0400 Subject: [PATCH 2/7] feat: Track and report continuous worker idle time WorkerService records a monotonic timestamp of when its in-flight task set last became empty, with worker startup counting as the initial empty state, and answers the idle RPC with the seconds elapsed since, or zero while any task is in flight. The timestamp flips on the docket's empty and non-empty edges, so polling never enters the docket and cannot disturb the measurement it reads. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/src/wool/runtime/worker/service.py | 48 +++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/wool/src/wool/runtime/worker/service.py b/wool/src/wool/runtime/worker/service.py index bf233714..4088dddc 100644 --- a/wool/src/wool/runtime/worker/service.py +++ b/wool/src/wool/runtime/worker/service.py @@ -4,6 +4,7 @@ import logging import pickle import threading +import time from contextlib import asynccontextmanager from dataclasses import dataclass from inspect import isawaitable @@ -140,11 +141,17 @@ class WorkerService(protocol.WorkerServicer): _stopped: asyncio.Event _stopping: asyncio.Event _loop_pool: ResourcePool[tuple[asyncio.AbstractEventLoop, threading.Thread]] + # Monotonic timestamp at which the docket last became empty, or + # ``None`` while any task is in flight. + _idle_since: float | None def __init__(self, *, backpressure: BackpressureLike | None = None): self._stopped = asyncio.Event() self._stopping = asyncio.Event() self._docket = set() + # Worker startup counts as the initial empty state, so idle is + # measured from construction. + self._idle_since = time.monotonic() self._backpressure = backpressure # Strong refs to live ``session.cancel()`` tasks scheduled # from ``_propagate_cancel_on_done`` (a gRPC-internal-thread @@ -557,6 +564,36 @@ async def stop( await self._stop(timeout=request.timeout) return protocol.Void() + async def idle( + self, + request: protocol.Void, + context: ServicerContext | None, + ) -> protocol.IdleTime: + """Report how long the worker has been continuously idle. + + Idle is the number of seconds since the in-flight task set + (`_docket`) last became empty, with worker startup counting as + the initial empty state. While any task is in flight the + reported idle time is zero, and the count resets whenever work + resumes. Measured against a monotonic clock so a wall-clock + adjustment cannot distort it. + + Polling this RPC creates no `DispatchSession` and never touches + the docket, so a caller cannot disturb the measurement it reads. + + :param request: + The empty protobuf request. + :param context: + The `grpc.aio.ServicerContext` for this request. + :returns: + An `IdleTime` carrying the continuous idle duration in + seconds. + """ + seconds = ( + 0.0 if self._idle_since is None else time.monotonic() - self._idle_since + ) + return protocol.IdleTime(seconds=seconds) + @staticmethod def _create_worker_loop( key, @@ -703,10 +740,21 @@ async def _tracked( StatusCode.UNAVAILABLE, "Worker service is shutting down" ) self._docket.add(session) + # Empty -> non-empty edge: work has resumed, so the worker is + # no longer idle. The add and this check share a single + # main-loop turn (no ``await`` between), so the transition is + # race-free against concurrent dispatches. + if len(self._docket) == 1: + self._idle_since = None try: yield finally: self._docket.discard(session) + # Non-empty -> empty edge: the in-flight set just drained, + # so start counting idle from now. Same single-turn + # atomicity as the add-side edge above. + if not self._docket: + self._idle_since = time.monotonic() async def _stop(self, *, timeout: float | None = 0) -> None: if timeout is not None and timeout < 0: From 58d6da9062dc1619334242e6adb002fae9874993 Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:55:59 -0400 Subject: [PATCH 3/7] feat: Add idle_time and stop to WorkerConnection Give the direct single-worker connection an idle_time poll that returns the worker's continuous idle seconds (validating a positive deadline) and a stop control that sends the stop RPC with a shutdown grace period. A worker predating the idle RPC answers UNIMPLEMENTED, which surfaces as the new IdleUnavailable; it descends from WoolError rather than RpcError so an absent capability is not mistaken for an RPC-health fault and does not evict a healthy worker. Export IdleUnavailable from the wool package. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/src/wool/__init__.py | 2 + wool/src/wool/runtime/worker/connection.py | 93 +++++++++++++++++++++- 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/wool/src/wool/__init__.py b/wool/src/wool/__init__.py index 98a1e740..017edd7f 100644 --- a/wool/src/wool/__init__.py +++ b/wool/src/wool/__init__.py @@ -48,6 +48,7 @@ from wool.runtime.worker.base import Worker from wool.runtime.worker.base import WorkerFactory from wool.runtime.worker.base import WorkerLike +from wool.runtime.worker.connection import IdleUnavailable from wool.runtime.worker.connection import RpcError from wool.runtime.worker.connection import TransientRpcError from wool.runtime.worker.connection import UnexpectedResponse @@ -110,6 +111,7 @@ "DiscoveryPublisherLike", "DiscoverySubscriberLike", "Factory", + "IdleUnavailable", "IneffectiveLeaseWarning", "IneffectiveQuorumTimeoutWarning", "LanDiscovery", diff --git a/wool/src/wool/runtime/worker/connection.py b/wool/src/wool/runtime/worker/connection.py index c48831b4..f1120dc8 100644 --- a/wool/src/wool/runtime/worker/connection.py +++ b/wool/src/wool/runtime/worker/connection.py @@ -17,6 +17,7 @@ import wool from wool import protocol +from wool.exceptions import WoolError from wool.runtime.resourcepool import ResourcePool from wool.runtime.routine.task import Task from wool.runtime.serializer import Serializer @@ -121,9 +122,29 @@ class TransientRpcError(RpcError): """ +# public +class IdleUnavailable(WoolError): + """Raised when a worker does not implement the idle RPC. + + A worker that predates the idle capability answers the idle RPC + with gRPC ``UNIMPLEMENTED``; `WorkerConnection.idle_time` + translates that into this typed signal so a polling client can + detect an old worker and treat idle reporting as unavailable, + distinct from a transient hiccup or an unhealthy peer. It descends + from `wool.WoolError` — not `RpcError` — because an absent + capability is not an RPC-health fault, so ``except RpcError`` does + not catch it. + """ + + # public class WorkerConnection: - """gRPC connection to a worker for task dispatch. + """Direct single-worker control surface over a pooled gRPC channel. + + Exposes `dispatch` (task execution), `idle_time` (poll the worker's + continuous idle duration), and `stop` (shut the remote worker + down). ``close`` is distinct: it releases this connection's local + pooled channel, whereas `stop` terminates the remote worker. Acquires pooled gRPC channels keyed by ``(target, credentials, options)``. Each `dispatch` call obtains a reference-counted @@ -329,6 +350,76 @@ async def close(self): except KeyError: pass + async def idle_time(self, *, timeout: float | None = None) -> float: + """Query how long the remote worker has been continuously idle. + + Returns the worker's reported continuous idle duration; see + `WorkerService.idle` for how idle is defined and measured. The + call draws a channel from this connection's pool, inheriting + its credential and secure-channel handling. + + :param timeout: + gRPC call deadline in seconds for the poll. ``None`` + applies no deadline. + :returns: + The worker's continuous idle duration in seconds. + :raises ValueError: + If ``timeout`` is not positive. + :raises IdleUnavailable: + If the worker predates the idle RPC (gRPC ``UNIMPLEMENTED``). + :raises TransientRpcError: + If the worker returns a transient RPC error + (``UNAVAILABLE``, ``DEADLINE_EXCEEDED``, or + ``RESOURCE_EXHAUSTED``). + :raises RpcError: + If the worker returns any other non-transient RPC error. + """ + if timeout is not None and timeout <= 0: + raise ValueError("Idle timeout must be positive") + async with _channel_pool.get(self._key) as channel: + try: + response = await channel.stub.idle(protocol.Void(), timeout=timeout) + except grpc.RpcError as error: + code = error.code() + details = error.details() or str(error) + if code == grpc.StatusCode.UNIMPLEMENTED: + raise IdleUnavailable( + "Worker does not implement idle reporting" + ) from error + if code in self.TRANSIENT_ERRORS: + raise TransientRpcError(code, details) from error + raise RpcError(code, details) from error + return response.seconds + + async def stop(self, *, grace: float | None = None) -> None: + """Stop the remote worker. + + Sends the stop RPC over a channel drawn from this connection's + pool, inheriting its credential and secure-channel handling. + + :param grace: + The worker's shutdown grace period in seconds, carried in + the stop request. ``None`` (the wire default of ``0``) + cancels in-flight tasks immediately; a positive value + waits that long for a graceful drain; a negative value + waits indefinitely. + :raises TransientRpcError: + If the worker returns a transient RPC error + (``UNAVAILABLE``, ``DEADLINE_EXCEEDED``, or + ``RESOURCE_EXHAUSTED``). + :raises RpcError: + If the worker returns any other non-transient RPC error. + """ + async with _channel_pool.get(self._key) as channel: + try: + await channel.stub.stop(protocol.StopRequest(timeout=grace)) + except grpc.RpcError as error: + code = error.code() + details = error.details() or str(error) + if code in self.TRANSIENT_ERRORS: + raise TransientRpcError(code, details) from error + raise RpcError(code, details) from error + async def _handshake( self, call: _DispatchCall, From 9f6a2f3625855861255a04e2cc42c9e9408b8c6c Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:55:59 -0400 Subject: [PATCH 4/7] refactor: Rename the worker shutdown grace parameter to grace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rename the stop() shutdown grace period from timeout to grace on Worker and WorkerLike so it no longer reads as a client-side deadline and no longer collides with the dispatch and idle timeout on the same surface. The old timeout keyword is kept as a deprecated alias — passing it still works and emits a DeprecationWarning — so existing callers are not broken. The wire field StopRequest.timeout and the startup deadline _start keep the name timeout. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/src/wool/runtime/worker/README.md | 4 ++-- wool/src/wool/runtime/worker/base.py | 33 +++++++++++++++++++++----- wool/src/wool/runtime/worker/pool.py | 2 +- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/wool/src/wool/runtime/worker/README.md b/wool/src/wool/runtime/worker/README.md index 54876744..e012d2db 100644 --- a/wool/src/wool/runtime/worker/README.md +++ b/wool/src/wool/runtime/worker/README.md @@ -139,14 +139,14 @@ class ResilientWorker(LocalWorker): await super()._start(timeout=timeout) self._monitor_task = asyncio.create_task(self._monitor()) - async def _stop(self, timeout=None): + async def _stop(self, grace=None): if self._monitor_task: self._monitor_task.cancel() try: await self._monitor_task except asyncio.CancelledError: pass - await super()._stop(timeout=timeout) + await super()._stop(grace=grace) async def _restart(self): """Replace the dead process, reusing the original port.""" diff --git a/wool/src/wool/runtime/worker/base.py b/wool/src/wool/runtime/worker/base.py index 5bd9a4ba..20060a44 100644 --- a/wool/src/wool/runtime/worker/base.py +++ b/wool/src/wool/runtime/worker/base.py @@ -3,6 +3,7 @@ import functools import inspect import uuid +import warnings from abc import ABC from abc import abstractmethod from dataclasses import dataclass @@ -295,11 +296,15 @@ async def start(self, *, timeout: float | None = None): """ ... - async def stop(self, *, timeout: float | None = None): + async def stop(self, *, grace: float | None = None, timeout: float | None = None): """Stop the worker and unregister it from the pool. + :param grace: + The worker's shutdown grace period in seconds — how long to + wait for in-flight tasks to drain before cancelling them. :param timeout: - Maximum time in seconds to wait for worker shutdown. + Deprecated alias for ``grace``, retained for backwards + compatibility; passing it emits a ``DeprecationWarning``. :raises RuntimeError: If the worker has not been started. """ @@ -326,7 +331,7 @@ async def _start(self, timeout): # Start your worker process self._info = WorkerMetadata(...) - async def _stop(self, timeout): + async def _stop(self, grace): # Clean shutdown ... @@ -406,17 +411,33 @@ async def start(self, *, timeout: float | None = None): assert self._info @final - async def stop(self, *, timeout: float | None = None): + async def stop(self, *, grace: float | None = None, timeout: float | None = None): """Stop the worker and unregister it from the pool. This method is a final implementation that calls the abstract `_stop` method to gracefully shut down the worker process and unregister it from the registrar service. + + :param grace: + The worker's shutdown grace period in seconds — how long to + wait for in-flight tasks to drain before cancelling them. + :param timeout: + Deprecated alias for ``grace``, retained for backwards + compatibility; passing it emits a ``DeprecationWarning``. """ + if timeout is not None: + warnings.warn( + "The 'timeout' parameter of Worker.stop is deprecated; " + "use 'grace' instead.", + DeprecationWarning, + stacklevel=2, + ) + if grace is None: + grace = timeout if not self._started: raise RuntimeError("Worker has not been started") try: - await self._stop(timeout) + await self._stop(grace) finally: self._started = False @@ -433,7 +454,7 @@ async def _start(self, timeout: float | None): ... @abstractmethod - async def _stop(self, timeout: float | None): + async def _stop(self, grace: float | None): """Implementation-specific worker shutdown logic. Subclasses must implement this method to handle the graceful diff --git a/wool/src/wool/runtime/worker/pool.py b/wool/src/wool/runtime/worker/pool.py index 7663dea5..d171f5fa 100644 --- a/wool/src/wool/runtime/worker/pool.py +++ b/wool/src/wool/runtime/worker/pool.py @@ -648,7 +648,7 @@ async def start(worker): async def stop(worker): await publisher_svc.publish("worker-dropped", worker.metadata) - await worker.stop(timeout=self._shutdown_timeout) + await worker.stop(grace=self._shutdown_timeout) task = asyncio.create_task(start(worker)) tasks.append(task) From 83e0f867cf96f8ae542f704cbe5bfbf0452e91ff Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:56:00 -0400 Subject: [PATCH 5/7] refactor: Route LocalWorker stop through WorkerConnection LocalWorker._stop now issues the stop RPC through a WorkerConnection, reusing its credential handling, secure-channel construction, and channel pooling instead of building a one-off stub, and releases the pooled channel afterward. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/src/wool/runtime/worker/local.py | 31 ++++++++++++--------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/wool/src/wool/runtime/worker/local.py b/wool/src/wool/runtime/worker/local.py index e84e7038..3dbf4fd6 100644 --- a/wool/src/wool/runtime/worker/local.py +++ b/wool/src/wool/runtime/worker/local.py @@ -4,12 +4,10 @@ from typing import TYPE_CHECKING from typing import Any -import grpc.aio - -from wool import protocol from wool.runtime.worker.auth import WorkerCredentials from wool.runtime.worker.base import Worker from wool.runtime.worker.base import WorkerOptions +from wool.runtime.worker.connection import WorkerConnection from wool.runtime.worker.process import WorkerProcess if TYPE_CHECKING: @@ -134,7 +132,7 @@ async def _start(self, timeout: float | None): if self._info is None: raise RuntimeError("Worker process failed to start - no metadata") - async def _stop(self, timeout: float | None): + async def _stop(self, grace: float | None): """Stop the worker process and unregister it from the pool. Unregisters the worker from the registrar service, gracefully @@ -142,22 +140,21 @@ async def _stop(self, timeout: float | None): up the registrar service. If graceful shutdown fails, the process is forcefully terminated. - For workers configured with :class:`WorkerCredentials`, uses - the client credentials to establish a secure connection for the - stop operation. For insecure workers, uses an insecure channel. + Credential and secure-channel handling for the stop RPC lives + on `WorkerConnection.stop`. """ if self._worker_process.is_alive(): assert self.address - # Create appropriate channel based on available credentials - if self._credentials is not None: - credentials = self._credentials.client_credentials() - channel = grpc.aio.secure_channel(self.address, credentials) - else: - channel = grpc.aio.insecure_channel(self.address) - + # Route the stop RPC through WorkerConnection; ``close`` + # releases the pooled channel this stop acquired. + credentials = ( + self._credentials.client_credentials() + if self._credentials is not None + else None + ) + connection = WorkerConnection(self.address, credentials=credentials) try: - stub = protocol.WorkerStub(channel) - await stub.stop(protocol.StopRequest(timeout=timeout)) + await connection.stop(grace=grace) finally: - await channel.close() + await connection.close() From 7b7f626985d3b875222728ef70567e535ba7a1d5 Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:56:00 -0400 Subject: [PATCH 6/7] test: Cover worker idle time and the connection control surface Add unit and property-based coverage for the IdleTime wire message, WorkerService idle tracking across concurrent drains and pre-tracking rejections, WorkerConnection.idle_time and stop (status-code mapping, timeout validation, and value fidelity), IdleUnavailable's WoolError lineage and public export, and the LocalWorker stop routing. Track the grace rename across the worker suite and cover the deprecated timeout alias on Worker.stop. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/tests/protocol/test_wire.py | 41 ++ wool/tests/runtime/worker/test_base.py | 33 +- wool/tests/runtime/worker/test_connection.py | 406 +++++++++++++++++++ wool/tests/runtime/worker/test_local.py | 50 ++- wool/tests/runtime/worker/test_pool.py | 8 +- wool/tests/runtime/worker/test_service.py | 265 +++++++++++- wool/tests/test_exceptions.py | 12 +- wool/tests/test_public.py | 1 + 8 files changed, 802 insertions(+), 14 deletions(-) diff --git a/wool/tests/protocol/test_wire.py b/wool/tests/protocol/test_wire.py index 0784119e..5939c8eb 100644 --- a/wool/tests/protocol/test_wire.py +++ b/wool/tests/protocol/test_wire.py @@ -1,10 +1,14 @@ import pytest +from hypothesis import given +from hypothesis import settings +from hypothesis import strategies as st from wool import protocol EXPECTED_MESSAGE_EXPORTS = [ "Ack", "ChainManifest", + "IdleTime", "Message", "Nack", "Request", @@ -251,6 +255,43 @@ def test_nack_without_exception(self): parsed.ParseFromString(nack.SerializeToString()) assert parsed.HasField("exception") is False + def test_idle_time_fields(self): + """Test IdleTime carries the idle duration in seconds. + + Given: + A seconds value, and the default construction. + When: + An IdleTime message is constructed with and without a value. + Then: + The seconds field should hold the value and default to 0.0. + """ + # Arrange, act, & assert + assert protocol.IdleTime(seconds=42.5).seconds == 42.5 + assert protocol.IdleTime().seconds == 0.0 + + @settings(max_examples=100) + @given(seconds=st.floats(width=64, allow_nan=False, allow_infinity=False)) + def test_idle_time_roundtrip(self, seconds): + """Test IdleTime.seconds survives the wire-format round-trip. + + Given: + Any finite double value for the idle duration. + When: + An IdleTime message is serialized and re-parsed. + Then: + The seconds field should equal the original exactly — a + proto double round-trips losslessly. + """ + # Arrange + message = protocol.IdleTime(seconds=seconds) + + # Act + parsed = protocol.IdleTime() + parsed.ParseFromString(message.SerializeToString()) + + # Assert + assert parsed.seconds == seconds + class TestOneofBehavior: """Tests for Request/Response oneof semantics.""" diff --git a/wool/tests/runtime/worker/test_base.py b/wool/tests/runtime/worker/test_base.py index e18daefa..5e7cacf1 100644 --- a/wool/tests/runtime/worker/test_base.py +++ b/wool/tests/runtime/worker/test_base.py @@ -366,7 +366,7 @@ async def _start(self, timeout: float | None): extra=MappingProxyType(self._extra), ) - async def _stop(self, timeout: float | None): + async def _stop(self, grace: float | None): """Mock stop implementation.""" pass @@ -638,11 +638,40 @@ async def test_stop_calls_implementation_stop(self, mocker): ) # Act - await worker.stop(timeout=60.0) + await worker.stop(grace=60.0) # Assert mock_stop.assert_called_once_with(60.0) + @pytest.mark.asyncio + async def test_stop_forwards_deprecated_timeout_as_grace(self, mocker): + """Test stop accepts the deprecated timeout alias for grace. + + Given: + A started Worker with a mocked _stop + When: + stop is called with the deprecated timeout keyword + Then: + It should emit a DeprecationWarning and forward the value to + _stop as the grace period, preserving backwards + compatibility. + """ + # Arrange + worker = ConcreteWorker() + await worker.start() + mock_stop = mocker.patch.object( + worker, + "_stop", + new_callable=mocker.AsyncMock, + ) + + # Act + with pytest.warns(DeprecationWarning): + await worker.stop(timeout=7.5) + + # Assert + mock_stop.assert_called_once_with(7.5) + @pytest.mark.asyncio async def test_stop_allows_restart(self): """Test Worker stop method allows worker to be restarted. diff --git a/wool/tests/runtime/worker/test_connection.py b/wool/tests/runtime/worker/test_connection.py index 1934eec5..57565317 100644 --- a/wool/tests/runtime/worker/test_connection.py +++ b/wool/tests/runtime/worker/test_connection.py @@ -21,6 +21,7 @@ from wool.runtime.routine.task import Task from wool.runtime.routine.task import WorkerProxyLike from wool.runtime.worker.base import ChannelOptions +from wool.runtime.worker.connection import IdleUnavailable from wool.runtime.worker.connection import RpcError from wool.runtime.worker.connection import TransientRpcError from wool.runtime.worker.connection import UnexpectedResponse @@ -30,6 +31,28 @@ from .conftest import PicklableMock +class _MockRpcError(grpc.RpcError): + """A ``grpc.RpcError`` carrying a controllable ``code``/``details``. + + Defined at module scope so property tests can raise a realistic + gRPC error for any status code without redefining the class per + example. ``str(self)`` is non-empty so the ``details() or + str(error)`` fallback in :class:`WorkerConnection` is observable + when ``details`` is empty. + """ + + def __init__(self, code: grpc.StatusCode, details: str | None): + super().__init__(f"") + self._code = code + self._details = details + + def code(self) -> grpc.StatusCode: + return self._code + + def details(self) -> str | None: + return self._details + + class MyAppError(Exception): """Custom user exception subclass defined at module scope. @@ -394,6 +417,389 @@ def details(self): async for _ in await connection.dispatch(sample_task): pass + @pytest.mark.asyncio + async def test_idle_time_should_return_seconds_when_worker_responds( + self, mocker: MockerFixture + ): + """Test idle returns the worker's reported idle seconds. + + Given: + A connection whose worker answers the idle RPC with an + IdleTime + When: + idle is awaited + Then: + It should return the reported seconds as a float. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock(return_value=protocol.IdleTime(seconds=42.5)) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + + # Act + result = await connection.idle_time() + + # Assert + assert result == 42.5 + mock_stub.idle.assert_awaited_once() + + @pytest.mark.asyncio + async def test_idle_time_should_raise_idle_unavailable_when_unimplemented( + self, mocker: MockerFixture + ): + """Test idle surfaces IdleUnavailable for a worker without the RPC. + + Given: + A connection whose worker answers the idle RPC with gRPC + UNIMPLEMENTED (a worker predating idle reporting) + When: + idle is awaited + Then: + It should raise IdleUnavailable so the caller can treat idle + reporting as unavailable for this worker. + """ + + # Arrange + class MockRpcError(grpc.RpcError): + def code(self): + return grpc.StatusCode.UNIMPLEMENTED + + def details(self): + return "method idle not implemented" + + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock(side_effect=MockRpcError()) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + + # Act & assert + with pytest.raises(IdleUnavailable): + await connection.idle_time() + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given( + code=st.sampled_from( + [c for c in grpc.StatusCode if c is not grpc.StatusCode.OK] + ), + details=st.one_of(st.none(), st.just(""), st.text()), + ) + async def test_idle_time_should_map_status_code_to_typed_exception( + self, mocker: MockerFixture, code: grpc.StatusCode, details: str | None + ): + """Test idle maps every gRPC status code to the right exception. + + Given: + A connection whose worker answers the idle RPC with any gRPC + error status code and any details + When: + idle is awaited + Then: + It should raise IdleUnavailable for UNIMPLEMENTED (chained + from the gRPC error and not an RpcError), TransientRpcError + for transient codes, and RpcError otherwise — carrying the + code and details, falling back to str(error) when details + are empty. + """ + # Arrange + error = _MockRpcError(code, details) + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock(side_effect=error) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example — the module pool persists across + # Hypothesis examples within one test function, so evict any + # channel a prior example cached under this key. + await clear_channel_pool() + + # Act + with pytest.raises(Exception) as exc_info: + await connection.idle_time() + + # Assert + raised = exc_info.value + if code is grpc.StatusCode.UNIMPLEMENTED: + assert isinstance(raised, IdleUnavailable) + assert not isinstance(raised, RpcError) + assert raised.__cause__ is error + else: + if code in WorkerConnection.TRANSIENT_ERRORS: + assert isinstance(raised, TransientRpcError) + else: + assert isinstance(raised, RpcError) + assert not isinstance(raised, TransientRpcError) + assert raised.code is code + assert raised.details == (details or str(error)) + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given(timeout=st.floats(max_value=0.0, allow_nan=False, allow_infinity=False)) + async def test_idle_time_should_reject_non_positive_timeout( + self, mocker: MockerFixture, timeout: float + ): + """Test idle_time rejects a non-positive timeout with ValueError. + + Given: + A connection whose worker answers the idle RPC, and any + non-positive timeout + When: + idle_time is awaited with that timeout + Then: + It should raise ValueError without calling the stub, matching + dispatch's timeout validation. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock(return_value=protocol.IdleTime(seconds=1.0)) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example (see idle-map note above). + await clear_channel_pool() + + # Act & assert + with pytest.raises(ValueError): + await connection.idle_time(timeout=timeout) + mock_stub.idle.assert_not_awaited() + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given( + timeout=st.one_of( + st.none(), + st.floats( + min_value=0.0, + exclude_min=True, + allow_nan=False, + allow_infinity=False, + ), + ) + ) + async def test_idle_time_should_forward_positive_timeout( + self, mocker: MockerFixture, timeout: float | None + ): + """Test idle_time forwards None or a positive timeout to the stub. + + Given: + A connection whose worker answers the idle RPC, and None or + any positive timeout + When: + idle_time is awaited with that timeout + Then: + It should call the stub with a Void request and that timeout + as the gRPC deadline. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock(return_value=protocol.IdleTime(seconds=1.0)) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example (see idle-map note above). + await clear_channel_pool() + + # Act + result = await connection.idle_time(timeout=timeout) + + # Assert + assert result == 1.0 + mock_stub.idle.assert_awaited_once() + call = mock_stub.idle.await_args + assert isinstance(call.args[0], protocol.Void) + assert call.kwargs["timeout"] == timeout + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given(seconds=st.floats(width=64, allow_nan=False, allow_infinity=False)) + async def test_idle_time_should_return_reported_seconds_verbatim( + self, mocker: MockerFixture, seconds: float + ): + """Test idle returns the worker's reported seconds losslessly. + + Given: + A worker answering the idle RPC with any double value + When: + idle is awaited + Then: + It should return exactly that value (including 0.0) — the + double survives the round-trip through the connection. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.idle = mocker.AsyncMock( + return_value=protocol.IdleTime(seconds=seconds) + ) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example (see idle-map note above). + await clear_channel_pool() + + # Act + result = await connection.idle_time() + + # Assert + assert result == seconds + + @pytest.mark.asyncio + async def test_stop_should_send_stop_request_with_timeout( + self, mocker: MockerFixture + ): + """Test stop sends a StopRequest carrying the grace timeout. + + Given: + A connection to a worker that acknowledges the stop RPC + When: + stop is awaited with a timeout + Then: + It should send a StopRequest with that timeout and return + None. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.stop = mocker.AsyncMock(return_value=protocol.Void()) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + + # Act + result = await connection.stop(grace=5.0) + + # Assert + assert result is None + mock_stub.stop.assert_awaited_once() + sent = mock_stub.stop.await_args.args[0] + assert isinstance(sent, protocol.StopRequest) + assert sent.timeout == 5.0 + + @pytest.mark.asyncio + async def test_stop_should_send_zero_timeout_when_none(self, mocker: MockerFixture): + """Test stop sends the wire-default timeout when none is given. + + Given: + A connection to a worker that acknowledges the stop RPC + When: + stop is awaited with no timeout argument + Then: + It should send a StopRequest whose timeout is 0.0 (the proto + default), which the worker reads as cancel-immediately. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.stop = mocker.AsyncMock(return_value=protocol.Void()) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + + # Act + await connection.stop() + + # Assert + sent = mock_stub.stop.await_args.args[0] + assert isinstance(sent, protocol.StopRequest) + assert sent.timeout == 0.0 + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given( + code=st.sampled_from( + [c for c in grpc.StatusCode if c is not grpc.StatusCode.OK] + ), + details=st.one_of(st.none(), st.just(""), st.text()), + ) + async def test_stop_should_map_status_code_to_typed_exception( + self, mocker: MockerFixture, code: grpc.StatusCode, details: str | None + ): + """Test stop maps every gRPC status code to the right exception. + + Given: + A connection whose worker answers the stop RPC with any gRPC + error status code and any details + When: + stop is awaited + Then: + It should raise TransientRpcError for transient codes and + RpcError otherwise (including UNIMPLEMENTED, which stop does + not special-case), carrying the code and details, falling + back to str(error) when details are empty. + """ + # Arrange + error = _MockRpcError(code, details) + mock_stub = mocker.MagicMock() + mock_stub.stop = mocker.AsyncMock(side_effect=error) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example (see idle-map note above). + await clear_channel_pool() + + # Act + with pytest.raises(RpcError) as exc_info: + await connection.stop() + + # Assert + raised = exc_info.value + if code in WorkerConnection.TRANSIENT_ERRORS: + assert isinstance(raised, TransientRpcError) + else: + assert not isinstance(raised, TransientRpcError) + assert raised.code is code + assert raised.details == (details or str(error)) + + @pytest.mark.asyncio + @settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given(timeout=st.floats(width=32, allow_nan=False, allow_infinity=False)) + async def test_stop_should_forward_timeout_without_validation( + self, mocker: MockerFixture, timeout: float + ): + """Test stop forwards any timeout onto the StopRequest verbatim. + + Given: + A connection whose worker acknowledges the stop RPC, and any + float32 timeout including non-positive values + When: + stop is awaited with that timeout + Then: + It should send a StopRequest carrying that timeout (a + negative value rides the wire unchanged) without raising + ValueError — stop does not validate its timeout. + """ + # Arrange + mock_stub = mocker.MagicMock() + mock_stub.stop = mocker.AsyncMock(return_value=protocol.Void()) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + connection = WorkerConnection("localhost:50051") + # Fresh channel per example (see idle-map note above). + await clear_channel_pool() + + # Act + result = await connection.stop(grace=timeout) + + # Assert + assert result is None + sent = mock_stub.stop.await_args.args[0] + assert isinstance(sent, protocol.StopRequest) + assert sent.timeout == pytest.approx(timeout) + @pytest.mark.asyncio @settings(suppress_health_check=[HealthCheck.function_scoped_fixture]) @given(timeout=st.floats(max_value=0.0, allow_nan=False, allow_infinity=False)) diff --git a/wool/tests/runtime/worker/test_local.py b/wool/tests/runtime/worker/test_local.py index 8eec1630..e4ba9f0b 100644 --- a/wool/tests/runtime/worker/test_local.py +++ b/wool/tests/runtime/worker/test_local.py @@ -377,6 +377,46 @@ async def test_stop_sends_grpc_stop_request(self, mocker): # Assert mock_stub.stop.assert_called_once() + @pytest.mark.asyncio + async def test_stop_forwards_grace_timeout(self, mocker): + """Test _stop forwards its grace timeout onto the StopRequest. + + Given: + A started LocalWorker with an alive process + When: + stop() is called with a grace timeout + Then: + It should send a StopRequest carrying that timeout through + the pooled WorkerConnection. + """ + # Arrange + mock_process = mocker.MagicMock(spec=WorkerProcess) + mock_process.address = "127.0.0.1:50051" + mock_process.metadata = _make_metadata() + mock_process.start.return_value = None + mock_process.is_alive.return_value = True + + mocker.patch.object(local_module, "WorkerProcess", return_value=mock_process) + + worker = LocalWorker() + await worker.start() + + mock_channel = mocker.MagicMock() + mock_channel.close = mocker.AsyncMock() + mock_stub = mocker.MagicMock() + mock_stub.stop = mocker.AsyncMock() + + mocker.patch.object(grpc.aio, "insecure_channel", return_value=mock_channel) + mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) + + # Act + await worker.stop(grace=7.5) + + # Assert + sent = mock_stub.stop.await_args.args[0] + assert isinstance(sent, protocol.StopRequest) + assert sent.timeout == 7.5 + @pytest.mark.asyncio async def test_stop_does_nothing_if_process_not_alive(self, mocker): """Test _stop does nothing if worker process is not alive. @@ -409,7 +449,7 @@ async def test_stop_does_nothing_if_process_not_alive(self, mocker): mock_channel_fn.assert_not_called() @pytest.mark.asyncio - async def test_stop_handles_grpc_errors_gracefully(self, mocker): + async def test_stop_releases_channel_when_grpc_call_raises(self, mocker): """Test _stop handles gRPC errors without crashing. Given: @@ -439,10 +479,13 @@ async def test_stop_handles_grpc_errors_gracefully(self, mocker): mocker.patch.object(grpc.aio, "insecure_channel", return_value=mock_channel) mocker.patch.object(protocol, "WorkerStub", return_value=mock_stub) - # Act & assert + # Act with pytest.raises(Exception, match="gRPC error"): await worker.stop() + # Assert — the pooled channel is released even when stop raises + mock_channel.close.assert_awaited() + # === NEW CREDENTIAL TESTS === def test___init___with_worker_credentials(self, worker_credentials): @@ -613,7 +656,8 @@ async def test_stop_without_credentials_insecure_connection(self, mocker): await worker.stop() # Assert - mock_insecure_channel.assert_called_once_with("127.0.0.1:50051") + mock_insecure_channel.assert_called_once() + assert mock_insecure_channel.call_args[0][0] == "127.0.0.1:50051" @pytest.mark.asyncio async def test_stop_with_one_way_tls(self, mocker, worker_credentials_one_way): diff --git a/wool/tests/runtime/worker/test_pool.py b/wool/tests/runtime/worker/test_pool.py index 9609a9e5..fe7b7f33 100644 --- a/wool/tests/runtime/worker/test_pool.py +++ b/wool/tests/runtime/worker/test_pool.py @@ -1415,7 +1415,7 @@ def hanging_factory(*tags, credentials=None): worker = mocker.MagicMock(spec=LocalWorker) worker.start = mocker.AsyncMock() - async def hang(*, timeout=None): + async def hang(*, grace=None): await asyncio.sleep(60) worker.stop = hang @@ -1456,7 +1456,7 @@ def hanging_factory(*tags, credentials=None): worker = mocker.MagicMock(spec=LocalWorker) worker.start = mocker.AsyncMock() - async def hang(*, timeout=None): + async def hang(*, grace=None): await asyncio.sleep(60) worker.stop = hang @@ -1499,7 +1499,7 @@ def hanging_factory(*tags, credentials=None): worker = mocker.MagicMock(spec=LocalWorker) worker.start = mocker.AsyncMock() - async def hang(*, timeout=None): + async def hang(*, grace=None): try: await asyncio.sleep(60) except asyncio.CancelledError: @@ -1583,7 +1583,7 @@ def mixed_factory(*tags, credentials=None): worker.metadata = _make_worker_metadata() if not workers_built: - async def hang(*, timeout=None): + async def hang(*, grace=None): await asyncio.sleep(60) worker.stop = hang diff --git a/wool/tests/runtime/worker/test_service.py b/wool/tests/runtime/worker/test_service.py index f9c56295..77c3a686 100644 --- a/wool/tests/runtime/worker/test_service.py +++ b/wool/tests/runtime/worker/test_service.py @@ -27,17 +27,49 @@ from .conftest import PicklableMock -def make_task(callable, *, proxy_id="test-proxy-id"): +def make_task(callable, *, args=(), kwargs=None, proxy_id="test-proxy-id"): """Build a `wool.Task` wrapping *callable* with a throwaway worker proxy.""" return Task( id=uuid4(), callable=callable, - args=(), - kwargs={}, + args=args, + kwargs=kwargs or {}, proxy=PicklableMock(spec=WorkerProxyLike, id=proxy_id), ) +# Registry of per-task control events for tests that need several +# independently-releasable in-flight tasks (the single global +# `_control_event` below only drives one). Keyed by an opaque token +# carried in the task's args so the worker-side routine can look up its +# own event. Cross-loop/cross-thread safe (threading.Event); tests +# populate and clear it around each use. +_keyed_events: dict[str, threading.Event] = {} + + +class _KeyedTaskError(Exception): + """Failure raised by a keyed controllable task drawn to raise. + + Module-level so cloudpickle can resolve it when the routine's + exception ships back across the dispatch stream. + """ + + +async def _keyed_task(key: str, should_raise: bool): + """Controllable task keyed by *key*. + + Waits on ``_keyed_events[key]`` (set by the test to release it), + then returns *key* or raises `_KeyedTaskError`, per *should_raise*. + Module-level so cloudpickle can serialize it for dispatch; runs on + the worker loop in the same process, so it shares ``_keyed_events``. + """ + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, _keyed_events[key].wait) + if should_raise: + raise _KeyedTaskError(key) + return key + + @pytest.fixture(scope="function") def grpc_interceptors(): return [VersionInterceptor()] @@ -5085,6 +5117,233 @@ async def failing_aenter(self): assert "_UnpicklableRejected" in str(shipped) assert "unpicklable parse failure" in str(shipped) + @pytest.mark.asyncio + async def test_idle_should_report_elapsed_since_startup_when_no_task_dispatched( + self, grpc_aio_stub + ): + """Test the idle RPC reports time accrued since worker startup. + + Given: + A `WorkerService` that has never dispatched a task, so its + in-flight set has been empty since construction + When: + The idle RPC is polled twice with a short wait between + Then: + It should report a positive, non-decreasing idle duration + measured from startup. + """ + # Arrange + service = WorkerService() + + # Act + async with grpc_aio_stub(servicer=service) as stub: + await asyncio.sleep(0.05) + first = await stub.idle(protocol.Void()) + await asyncio.sleep(0.05) + second = await stub.idle(protocol.Void()) + + # Assert + assert first.seconds >= 0.04 + assert second.seconds >= first.seconds + + @pytest.mark.asyncio + async def test_idle_should_report_zero_when_task_in_flight( + self, service_fixture, mock_worker_proxy_cache + ): + """Test the idle RPC reports zero while a task is executing. + + Given: + A `WorkerService` with one task dispatched and still in + flight + When: + The idle RPC is polled + Then: + It should report zero seconds, since the in-flight set is + non-empty. + """ + # Arrange, act, & assert + async with service_fixture as (service, event, stub): + response = await stub.idle(protocol.Void()) + assert response.seconds == 0.0 + + @pytest.mark.asyncio + async def test_idle_should_reset_when_task_completes( + self, grpc_aio_stub, mock_worker_proxy_cache + ): + """Test the idle count resets once the in-flight set drains. + + Given: + A `WorkerService` that has accrued idle time since startup, + then dispatches and completes a task + When: + The idle RPC is polled before dispatch, during execution, + and after the task completes + Then: + It should report accrued idle before dispatch, zero while + in flight, and a smaller value than before after the set + drains — proving the count reset when work drained. + """ + # Arrange + global _control_event + _control_event = threading.Event() + wool_task = make_task(_controllable_task) + request = protocol.Request(task=wool_task.to_protobuf()) + service = WorkerService() + + # Act + try: + async with grpc_aio_stub(servicer=service) as stub: + await asyncio.sleep(0.1) + before = await stub.idle(protocol.Void()) + + stream = stub.dispatch() + await stream.write(request) + await stream.done_writing() + async for response in stream: + assert response.HasField("ack") + break + during = await stub.idle(protocol.Void()) + + _control_event.set() + [r async for r in stream] + after = await stub.idle(protocol.Void()) + finally: + if _control_event and not _control_event.is_set(): + _control_event.set() + _control_event = None + + # Assert + assert before.seconds >= 0.08 + assert during.seconds == 0.0 + assert after.seconds < before.seconds + + @pytest.mark.asyncio + @settings( + max_examples=20, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given(modes=st.lists(st.booleans(), min_size=1, max_size=4)) + async def test_idle_should_stay_zero_until_the_last_task_drains( + self, grpc_aio_stub, mock_worker_proxy_cache, modes + ): + """Test idle is zero until the last of N in-flight tasks drains. + + Given: + A `WorkerService` with N tasks (1-4) dispatched and in + flight, each drawn to return or raise + When: + The tasks are completed one at a time and idle is polled + after each completion + Then: + It should report zero while any task remains in flight + (including after each non-final completion) and reset to a + positive value only after the last task drains — regardless + of how many tasks there are or how each one finishes. + """ + # Arrange + keys = [uuid4().hex for _ in modes] + for key in keys: + _keyed_events[key] = threading.Event() + service = WorkerService() + + # Act & assert + try: + async with grpc_aio_stub(servicer=service) as stub: + streams = [] + for key, should_raise in zip(keys, modes): + wool_task = make_task(_keyed_task, args=(key, should_raise)) + stream = stub.dispatch() + await stream.write(protocol.Request(task=wool_task.to_protobuf())) + await stream.done_writing() + async for response in stream: + assert response.HasField("ack") + break + streams.append(stream) + + # All N in flight -> idle is zero. + assert (await stub.idle(protocol.Void())).seconds == 0.0 + + # Complete one at a time; idle stays zero until the last. + for index, (key, stream) in enumerate(zip(keys, streams)): + _keyed_events[key].set() + [_ async for _ in stream] # drain this task fully + if index < len(keys) - 1: + assert (await stub.idle(protocol.Void())).seconds == 0.0 + + # Last task drained -> idle reset and now accruing. + await asyncio.sleep(0.02) + assert (await stub.idle(protocol.Void())).seconds > 0 + finally: + for key in keys: + _keyed_events[key].set() + _keyed_events.pop(key, None) + if not service.stopped.is_set(): + await service.stop(protocol.StopRequest(), None) + + @pytest.mark.asyncio + @pytest.mark.parametrize("kind", ["backpressure", "malformed_id", "sync_callable"]) + async def test_idle_should_ignore_dispatches_rejected_before_tracking( + self, grpc_aio_stub, mock_worker_proxy_cache, kind + ): + """Test a dispatch rejected before tracking never resets idle. + + Given: + A `WorkerService` that has accrued idle, and a dispatch that + is rejected before it enters the in-flight set — backpressure + (RESOURCE_EXHAUSTED), a malformed task id, or a sync callable + When: + The rejected dispatch is attempted and idle is polled before + and after + Then: + It should keep counting from startup — the rejected task + never enters the docket, so idle is never reset to zero. + """ + # Arrange + if kind == "backpressure": + service = WorkerService(backpressure=lambda ctx: True) + else: + service = WorkerService() + + async def returns(): + return "unreached" + + def sync_callable(): + return "unreached" + + if kind == "sync_callable": + wool_task = make_task(sync_callable) + else: + wool_task = make_task(returns) + request = protocol.Request(task=wool_task.to_protobuf()) + if kind == "malformed_id": + request.task.id = "not-a-uuid" + + # Act & assert + try: + async with grpc_aio_stub(servicer=service) as stub: + await asyncio.sleep(0.05) + before = (await stub.idle(protocol.Void())).seconds + + stream = stub.dispatch() + await stream.write(request) + await stream.done_writing() + if kind == "backpressure": + with pytest.raises(grpc.RpcError) as exc_info: + [_ async for _ in stream] + assert exc_info.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + else: + responses = [r async for r in stream] + assert responses and responses[0].HasField("nack") + + after = (await stub.idle(protocol.Void())).seconds + + assert before > 0 + assert after >= before + finally: + if not service.stopped.is_set(): + await service.stop(protocol.StopRequest(), None) + class TestBackpressureContext: """Tests for `wool.runtime.worker.service.BackpressureContext`.""" diff --git a/wool/tests/test_exceptions.py b/wool/tests/test_exceptions.py index 5a66ae6e..bf6e73ad 100644 --- a/wool/tests/test_exceptions.py +++ b/wool/tests/test_exceptions.py @@ -1,6 +1,7 @@ import pytest from wool import AdvertiseHostError +from wool import IdleUnavailable from wool import IneffectiveLeaseWarning from wool import IneffectiveQuorumTimeoutWarning from wool import LoopbackAdvertisementWarning @@ -15,7 +16,14 @@ class TestWoolError: Fully qualified name: wool.exceptions.WoolError """ - def test___init___should_subclass_wool_error(self): + @pytest.mark.parametrize( + "exception", + [ + AdvertiseHostError, + IdleUnavailable, + ], + ) + def test___init___should_subclass_wool_error(self, exception): """Test wool's typed exceptions descend from WoolError. Given: @@ -27,7 +35,7 @@ def test___init___should_subclass_wool_error(self): catches all wool-raised errors. """ # Act & assert - assert issubclass(AdvertiseHostError, WoolError) + assert issubclass(exception, WoolError) class TestWoolWarning: diff --git a/wool/tests/test_public.py b/wool/tests/test_public.py index d572aaba..67398036 100644 --- a/wool/tests/test_public.py +++ b/wool/tests/test_public.py @@ -31,6 +31,7 @@ def test_public_api_completeness_should_match_expected_surface(): """ # Arrange expected_public_api = { + "IdleUnavailable", "RpcError", "TransientRpcError", "UnexpectedResponse", From 47d2083b370eb7f3f7e34c33f7394608dfebf14b Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 1 Jul 2026 21:56:00 -0400 Subject: [PATCH 7/7] test: Add integration coverage for worker idle time over the wire Exercise idle accrual from startup, zero-while-busy, reset-on-drain, IdleUnavailable against a servicer without the RPC, and the poll-then-retire and LocalWorker.stop flows against real subprocess workers across the insecure, mTLS, and one-way-TLS transports. Claude-Session: https://claude.ai/code/session_01LEWxpdabdvMq4evG6Zfq8x --- wool/tests/integration/test_worker_idle.py | 346 +++++++++++++++++++++ 1 file changed, 346 insertions(+) create mode 100644 wool/tests/integration/test_worker_idle.py diff --git a/wool/tests/integration/test_worker_idle.py b/wool/tests/integration/test_worker_idle.py new file mode 100644 index 00000000..a9815341 --- /dev/null +++ b/wool/tests/integration/test_worker_idle.py @@ -0,0 +1,346 @@ +import asyncio +import contextlib +import uuid +from contextlib import asynccontextmanager + +import grpc +import grpc.aio +import pytest + +from wool import protocol +from wool.runtime.discovery.local import LocalDiscovery +from wool.runtime.worker.connection import IdleUnavailable +from wool.runtime.worker.connection import RpcError +from wool.runtime.worker.connection import TransientRpcError +from wool.runtime.worker.connection import WorkerConnection +from wool.runtime.worker.local import LocalWorker +from wool.runtime.worker.pool import WorkerPool + +from . import routines +from .conftest import CredentialType +from .conftest import _DirectDiscovery + +pytestmark = pytest.mark.integration + + +def _client_credentials(creds): + """Resolve WorkerCredentials to the client-side channel credentials. + + Mirrors ``LocalWorker._stop``: secure workers use their client + credentials, insecure workers use ``None``. + """ + return creds.client_credentials() if creds is not None else None + + +async def _poll(predicate, *, tries=600, interval=0.05): + """Await until the sync *predicate* returns truthy, or fail.""" + for _ in range(tries): + if predicate(): + return + await asyncio.sleep(interval) + raise AssertionError("condition not met within timeout") + + +async def _poll_coro(predicate, *, tries=200, interval=0.05): + """Await until the async *predicate* returns truthy, or fail.""" + for _ in range(tries): + if await predicate(): + return + await asyncio.sleep(interval) + raise AssertionError("condition not met within timeout") + + +@asynccontextmanager +async def _bare_worker(creds): + """Start a real LocalWorker and yield it plus a direct-connection + factory. No pool — for tests that only poll idle/stop by address. + """ + conns = [] + worker = LocalWorker(credentials=creds) + await worker.start() + + def connect(): + conn = WorkerConnection(worker.address, credentials=_client_credentials(creds)) + conns.append(conn) + return conn + + try: + yield worker, connect + finally: + for conn in conns: + with contextlib.suppress(Exception): + await conn.close() + with contextlib.suppress(Exception): + await worker.stop() + + +@asynccontextmanager +async def _worker_with_pool(creds): + """Start a real LocalWorker behind a DURABLE pool (so routines can + be dispatched to it) and yield the worker, the pool, and a + direct-connection factory. Teardown is defensive. + """ + conns = [] + namespace = f"idle-{uuid.uuid4().hex[:12]}" + with LocalDiscovery(namespace) as discovery: + worker = LocalWorker(credentials=creds) + await worker.start() + + def connect(): + conn = WorkerConnection( + worker.address, credentials=_client_credentials(creds) + ) + conns.append(conn) + return conn + + try: + publisher = discovery.publisher + async with publisher: + await publisher.publish("worker-added", worker.metadata) + try: + pool = WorkerPool( + discovery=_DirectDiscovery(discovery), credentials=creds + ) + async with pool: + yield worker, pool, connect + finally: + with contextlib.suppress(Exception): + await publisher.publish("worker-dropped", worker.metadata) + finally: + for conn in conns: + with contextlib.suppress(Exception): + await conn.close() + with contextlib.suppress(Exception): + await worker.stop() + + +class TestWorkerIdleReporting: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "cred", + [CredentialType.INSECURE, CredentialType.MTLS, CredentialType.ONE_WAY], + ) + async def test_idle_time_should_accrue_from_startup_over_the_real_wire( + self, credentials_map, retry_grpc_internal, cred + ): + """Test idle accrues from startup, over each transport. + + Given: + A freshly started real worker that has never been dispatched + a task, reached over an insecure, mTLS, or one-way-TLS + channel + When: + A direct WorkerConnection polls idle twice with a wait + between + Then: + Both readings should be positive and non-decreasing — idle + counts from startup and inherits the connection's + credential handling. + """ + + async def body(): + # Arrange + async with _bare_worker(credentials_map[cred]) as (worker, connect): + conn = connect() + + # Act + first = await conn.idle_time() + await asyncio.sleep(0.1) + second = await conn.idle_time() + + # Assert + assert first >= 0.0 + assert second >= first + assert second > 0.0 + + await retry_grpc_internal(body) + + @pytest.mark.asyncio + async def test_idle_time_should_report_zero_while_a_task_is_in_flight( + self, credentials_map, retry_grpc_internal, tmp_path + ): + """Test idle is zero while a dispatched routine runs on the worker. + + Given: + A real worker with one in-flight routine (confirmed via its + "started" sentinel) + When: + idle is polled through a direct WorkerConnection + Then: + It should report exactly zero — the docket is non-empty. + """ + + async def body(): + async with _worker_with_pool(credentials_map[CredentialType.INSECURE]) as ( + worker, + pool, + connect, + ): + conn = connect() + sentinel = tmp_path / "in-flight.txt" + dispatch = asyncio.create_task( + routines.cancellable_sleep(str(sentinel), 30.0) + ) + try: + # Arrange + await _poll( + lambda: sentinel.exists() and sentinel.read_text() == "started" + ) + + # Act & assert + assert await conn.idle_time() == 0.0 + finally: + dispatch.cancel() + with contextlib.suppress(BaseException): + await dispatch + + await retry_grpc_internal(body) + + @pytest.mark.asyncio + async def test_idle_time_should_reset_after_the_in_flight_set_drains( + self, credentials_map, retry_grpc_internal, tmp_path + ): + """Test idle resets once the worker's in-flight set drains. + + Given: + A real worker that has accrued idle, then runs one routine + to completion + When: + idle is polled before dispatch and after the routine drains + Then: + It should report accrued idle before, and a smaller value + after — proving the count reset at the drain rather than + continuing to accrue from startup. + """ + + async def body(): + async with _worker_with_pool(credentials_map[CredentialType.INSECURE]) as ( + worker, + pool, + connect, + ): + conn = connect() + await asyncio.sleep(0.2) + before = await conn.idle_time() + + # Act + sentinel = tmp_path / "drain.txt" + await routines.cancellable_sleep(str(sentinel), 0.3) + + # Wait for the docket-drain to be reflected (idle > 0), + # then confirm it counts from the drain, not from startup. + async def _reset(): + return await conn.idle_time() > 0.0 + + await _poll_coro(_reset) + after = await conn.idle_time() + + # Assert + assert before > 0.0 + assert after < before + + await retry_grpc_internal(body) + + @pytest.mark.asyncio + async def test_idle_time_should_raise_idle_unavailable_for_a_legacy_worker( + self, retry_grpc_internal + ): + """Test idle surfaces IdleUnavailable against a legacy worker. + + Given: + A real gRPC server whose Worker servicer does not implement + idle (a stand-in for a worker predating the RPC), so the + framework answers UNIMPLEMENTED + When: + A real WorkerConnection polls idle against it + Then: + It should raise IdleUnavailable — and not an RpcError — over + the real channel. + """ + + async def body(): + # Arrange + server = grpc.aio.server() + protocol.add_WorkerServicer_to_server(protocol.WorkerServicer(), server) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + conn = WorkerConnection(f"127.0.0.1:{port}") + try: + # Act & assert + with pytest.raises(IdleUnavailable) as exc_info: + await conn.idle_time() + assert not isinstance(exc_info.value, RpcError) + finally: + await conn.close() + await server.stop(None) + + await retry_grpc_internal(body) + + +class TestWorkerControlSurface: + @pytest.mark.asyncio + @pytest.mark.parametrize("cred", [CredentialType.INSECURE, CredentialType.MTLS]) + async def test_poll_then_retire_should_terminate_the_worker( + self, credentials_map, retry_grpc_internal, cred + ): + """Test the poll-then-retire flow stops the worker. + + Given: + An idle real worker and a direct WorkerConnection to it + When: + The caller polls idle, then calls stop, then polls idle again + Then: + stop returns and the worker terminates — a subsequent idle + eventually raises TransientRpcError (the server is gone). + """ + + async def body(): + # Arrange + async with _bare_worker(credentials_map[cred]) as (worker, connect): + conn = connect() + assert await conn.idle_time() >= 0.0 + + # Act + await conn.stop() + + # Assert + async def _unreachable(): + try: + await conn.idle_time() + return False + except TransientRpcError: + return True + + await _poll_coro(_unreachable, tries=200, interval=0.1) + + await retry_grpc_internal(body) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "cred", + [CredentialType.INSECURE, CredentialType.MTLS, CredentialType.ONE_WAY], + ) + async def test_local_worker_stop_should_succeed_over_each_transport( + self, credentials_map, retry_grpc_internal, cred + ): + """Test LocalWorker.stop succeeds through the refactored path. + + Given: + A started real LocalWorker under each credential type + When: + worker.stop is awaited (now routed through + WorkerConnection.stop and close) + Then: + It should complete without error, exercising the refactored + _stop's credential and secure-channel handling end-to-end. + """ + + async def body(): + # Arrange + worker = LocalWorker(credentials=credentials_map[cred]) + await worker.start() + + # Act & assert (completes without raising) + await worker.stop() + + await retry_grpc_internal(body)