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/__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/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", 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/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, 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() 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) 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: 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) 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",