ringo-task-queue 0.1.0.dev7__py3-none-win_amd64.whl → 0.1.0.dev9__py3-none-win_amd64.whl

ringo_task_queue/client.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import os
 import secrets
+import sys
 from datetime import timedelta
 from itertools import islice
 from typing import TYPE_CHECKING, Any, Awaitable, Callable, Coroutine, Iterable, Literal
@@ -14,6 +15,7 @@ from grpc_health.v1 import health_pb2, health_pb2_grpc
 
 from . import _pb
 from ._proto.ringo.v1 import queue_pb2, queue_pb2_grpc
+from ._supervisor_policy import TaskTimeouts
 from ._version import PROTOCOL_MAJOR, PROTOCOL_MINOR, SDK_VERSION
 from .concurrency import AdaptiveConcurrency, ResourceConcurrency
 from .errors import (
@@ -172,9 +174,7 @@ class _Connection:
             ):
                 if key not in configured:
                     self._channel_options.append((key, value))
-        _check_compression_channel_option_conflict(
-            self._channel_options, compression
-        )
+        _check_compression_channel_option_conflict(self._channel_options, compression)
         # Channel-level default compression; survives every reconnect because
         # _open_channel() re-applies it to each new aio channel. ``None`` means
         # unspecified, so legacy raw channel options decide the encoding.
@@ -192,6 +192,8 @@ class _Connection:
     # -- lifecycle ---------------------------------------------------------
 
     def _open_channel(self) -> None:
+        if self._closed:
+            raise RingoError("connection is closed")
         if self._tls:
             credentials = grpc.ssl_channel_credentials()
             self._channel = grpc.aio.secure_channel(
@@ -226,10 +228,17 @@ class _Connection:
             call = stub.AttachLocal(request, metadata=self._metadata())
             response = await call.read()
             if response is grpc.aio.EOF or response is None:
-                raise UnavailableError("shared daemon closed AttachLocal during handshake")
+                raise UnavailableError(
+                    "shared daemon closed AttachLocal during handshake"
+                )
+            if self._closed:
+                call.cancel()
+                raise RingoError("connection is closed")
             if response.generation != self._local_generation:
                 call.cancel()
-                raise UnavailableError("shared daemon generation changed during AttachLocal")
+                raise UnavailableError(
+                    "shared daemon generation changed during AttachLocal"
+                )
         except grpc.aio.AioRpcError as exc:
             raise map_rpc_error(exc) from exc
         self._attach_call = call
@@ -285,6 +294,8 @@ class _Connection:
         deadline = asyncio.get_running_loop().time() + timeout
         last: Exception | None = None
         while True:
+            if self._closed:
+                raise RingoError("connection is closed")
             try:
                 response = await health_stub.Check(
                     health_pb2.HealthCheckRequest(),
@@ -302,6 +313,8 @@ class _Connection:
 
     async def reset(self, endpoint: str) -> None:
         """Point the connection at a new endpoint (local daemon restart)."""
+        if self._closed:
+            raise RingoError("connection is closed")
         await self._close_attach()
         self._endpoint, self._tls = _parse_endpoint(endpoint)
         channel, self._channel = self._channel, None
@@ -325,6 +338,8 @@ class _Connection:
             )
         except grpc.aio.AioRpcError as exc:
             raise map_rpc_error(exc) from exc
+        if self._closed:
+            raise RingoError("connection is closed")
         if response.protocol_major != PROTOCOL_MAJOR:
             raise IncompatibleVersionError(
                 f"server speaks protocol major {response.protocol_major},"
@@ -344,7 +359,7 @@ class _Connection:
     # -- helpers -----------------------------------------------------------
 
     def _require_stub(self) -> queue_pb2_grpc.QueueServiceStub:
-        if self._stub is None:
+        if self._closed or self._stub is None:
             raise RingoError("connection is closed")
         return self._stub
 
@@ -379,7 +394,6 @@ class _Connection:
         *,
         timeout: float = DEFAULT_RPC_TIMEOUT,
         retry_safe: bool = False,
-        context_attr: str | None = "context",
     ) -> Any:
         """Issue a single unary RPC.
 
@@ -391,9 +405,7 @@ class _Connection:
         pass ``retry_safe=True`` and may be re-sent, each time with a freshly
         negotiated ``RequestContext``.
         """
-        rpc, sent_token = await self._begin_unary(
-            name, request, timeout=timeout, context_attr=context_attr
-        )
+        rpc, sent_token = await self._begin_unary(name, request, timeout=timeout)
         cause: BaseException | None = None
         err: RingoError | None = None
         try:
@@ -423,9 +435,7 @@ class _Connection:
             async with self._lock:
                 rotated = self._session_token != sent_token
             if rotated:
-                rpc, _ = await self._begin_unary(
-                    name, request, timeout=timeout, context_attr=context_attr
-                )
+                rpc, _ = await self._begin_unary(name, request, timeout=timeout)
                 try:
                     return await rpc
                 except grpc.aio.AioRpcError as retry_exc:
@@ -439,9 +449,7 @@ class _Connection:
         await self._recover_best_effort()
         if not retry_safe:
             raise err from cause
-        rpc, _ = await self._begin_unary(
-            name, request, timeout=timeout, context_attr=context_attr
-        )
+        rpc, _ = await self._begin_unary(name, request, timeout=timeout)
         try:
             return await rpc
         except grpc.aio.AioRpcError as retry_exc:
@@ -453,7 +461,6 @@ class _Connection:
         request: Any,
         *,
         timeout: float,
-        context_attr: str | None,
     ) -> tuple[Any, str]:
         """Snapshot a stub and fresh session context outside reconnect races."""
         async with self._lock:
@@ -466,8 +473,12 @@ class _Connection:
                     "attempt history requires protocol minor >= 5 (GetTaskAttempts);"
                     f" negotiated minor is {self._session_minor}"
                 )
-            if context_attr is not None:
-                getattr(request, context_attr).CopyFrom(self.request_context())
+            if name == "ProbeQueue" and self._session_minor < 7:
+                raise IncompatibleVersionError(
+                    "queue eligibility probe requires protocol minor >= 7 (ProbeQueue);"
+                    f" negotiated minor is {self._session_minor}"
+                )
+            request.context.CopyFrom(self.request_context())
             stub = self._require_stub()
             token = self._session_token
             rpc = getattr(stub, name)(
@@ -545,7 +556,9 @@ class RingoClient:
         self._daemon = daemon
         self._shared_daemon = shared_daemon
         if daemon is not None and shared_daemon is not None:
-            raise InvalidArgumentError("daemon and shared_daemon are mutually exclusive")
+            raise InvalidArgumentError(
+                "daemon and shared_daemon are mutually exclusive"
+            )
         # ``_conn`` is built in start(): a local client learns the daemon's
         # random loopback port only after the daemon announces readiness.
         self._conn: _Connection | None = None
@@ -553,6 +566,9 @@ class RingoClient:
         self._workers: set[Worker] = set()
         self._started = False
         self._closed = False
+        # Start, reconnect and final resource cleanup share one lifecycle lane.
+        # Reconnect enters from the connection's session lock; cleanup must not
+        # acquire that lock while holding this one.
         self._start_lock = asyncio.Lock()
         self._close_lock = asyncio.Lock()
         self._close_task: asyncio.Task[None] | None = None
@@ -560,10 +576,19 @@ class RingoClient:
 
     # -- lifecycle ---------------------------------------------------------
 
+    def _check_open(self) -> None:
+        if self._closed:
+            raise RingoError("client is closed")
+
     async def start(self) -> "RingoClient":
         async with self._start_lock:
-            if self._closed:
-                raise RingoError("client is closed")
+            self._check_open()
+            # Register before allocating resources: a runtime STOP fence rejects
+            # late starts, and an in-flight start is visible to unified drain.
+            module = sys.modules.get("ringo_task_queue._supervisor_runtime")
+            runtime = None if module is None else module.current_runtime()
+            if runtime is not None:
+                runtime.register_client(self)
             if not self._started:
                 try:
                     for attempt in range(2 if self._shared_daemon is not None else 1):
@@ -571,11 +596,20 @@ class RingoClient:
                             self._endpoint = await self._daemon.start()
                         elif self._shared_daemon is not None:
                             self._endpoint = await self._shared_daemon.start()
+                        self._check_open()
                         conn = _Connection(
                             self._endpoint or "",
                             self._token,
-                            local_token=(self._shared_daemon.token if self._shared_daemon else None),
-                            local_generation=(self._shared_daemon.generation if self._shared_daemon else None),
+                            local_token=(
+                                self._shared_daemon.token
+                                if self._shared_daemon
+                                else None
+                            ),
+                            local_generation=(
+                                self._shared_daemon.generation
+                                if self._shared_daemon
+                                else None
+                            ),
                             channel_options=self._channel_options,
                             compression=self._compression,
                         )
@@ -589,13 +623,19 @@ class RingoClient:
                             )
                         try:
                             await conn.connect(
-                                check_health=self._daemon is not None or self._shared_daemon is not None
+                                check_health=self._daemon is not None
+                                or self._shared_daemon is not None
                             )
+                            self._check_open()
                             break
                         except Exception:
                             await conn.close()
                             self._conn = None
-                            if self._shared_daemon is None or attempt > 0:
+                            if (
+                                self._closed
+                                or self._shared_daemon is None
+                                or attempt > 0
+                            ):
                                 raise
                             # Stale PID/port, wrong token or generation is never
                             # accepted as identity. Re-elect once, boundedly.
@@ -620,34 +660,47 @@ class RingoClient:
         return conn
 
     async def _reconnect(self) -> None:
-        """Reconnect after a failed transport and negotiate a fresh session."""
-        conn = self._require_conn()
-        daemon = self._daemon
-        shared = self._shared_daemon
-        if daemon is not None and daemon.is_dead:
-            if self._daemon_restart_count >= _MAX_DAEMON_RESTARTS:
-                raise UnavailableError(
-                    f"local daemon restart budget ({_MAX_DAEMON_RESTARTS}) exhausted"
-                )
-            self._daemon_restart_count += 1
-            delay = 0.3 * (2 ** (self._daemon_restart_count - 1))
-            await asyncio.sleep(delay)
-            endpoint = await daemon.start()
-            await conn.reset(endpoint)
-        elif shared is not None:
-            try:
-                endpoint = await shared.refresh()
+        """Rotate resources in the lifecycle lane; close wins at every boundary."""
+        async with self._start_lock:
+            self._check_open()
+            conn = self._require_conn()
+            daemon = self._daemon
+            shared = self._shared_daemon
+            if daemon is not None and daemon.is_dead:
+                if self._daemon_restart_count >= _MAX_DAEMON_RESTARTS:
+                    raise UnavailableError(
+                        f"local daemon restart budget ({_MAX_DAEMON_RESTARTS}) exhausted"
+                    )
+                self._daemon_restart_count += 1
+                delay = 0.3 * (2 ** (self._daemon_restart_count - 1))
+                await asyncio.sleep(delay)
+                self._check_open()
+                endpoint = await daemon.start()
+                self._check_open()
                 await conn.reset(endpoint)
-                conn.set_shared_session(
-                    shared.token or "", shared.generation or "", self._shared_stream_lost
-                )
-                await conn.attach_local()
-                await conn.negotiate()
-            except Exception:
-                shared.invalidate_current()
-                raise
-            return
-        await conn.negotiate()
+                self._check_open()
+            elif shared is not None:
+                try:
+                    endpoint = await shared.refresh()
+                    self._check_open()
+                    await conn.reset(endpoint)
+                    self._check_open()
+                    conn.set_shared_session(
+                        shared.token or "",
+                        shared.generation or "",
+                        self._shared_stream_lost,
+                    )
+                    await conn.attach_local()
+                    self._check_open()
+                    await conn.negotiate()
+                    self._check_open()
+                except Exception:
+                    if not self._closed:
+                        shared.invalidate_current()
+                    raise
+                return
+            await conn.negotiate()
+            self._check_open()
 
     async def _shared_stream_lost(self) -> None:
         """Proactively reconnect an idle shared client after AttachLocal EOF."""
@@ -662,7 +715,9 @@ class RingoClient:
             if self._closed:
                 return
             try:
-                await asyncio.sleep(_BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)])
+                await asyncio.sleep(
+                    _BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)]
+                )
                 async with conn._lock:
                     await self._reconnect()
                 return
@@ -694,18 +749,30 @@ class RingoClient:
             except Exception as exc:  # keep cleaning up the rest
                 errors.append(exc)
         self._workers.clear()
-        # 2. Close the gRPC channel.
-        conn, self._conn = self._conn, None
+        # 2. Fence the current channel first: this interrupts a pending handshake
+        # without acquiring the session lock (reconnect holds it). Work remains
+        # usable through the worker grace period above.
+        conn = self._conn
         if conn is not None:
             try:
                 await conn.close()
             except Exception as exc:
                 errors.append(exc)
-        # 3. Shut the local daemon down via the lifecycle pipe, then wait.
-        try:
-            await self._stop_daemon()
-        except Exception as exc:
-            errors.append(exc)
+        # 3. An in-progress daemon.start cannot be abandoned. Wait for its
+        # bounded completion, then reconcile resources under the lifecycle lane.
+        # Closed checks prevent it from publishing a new channel in the meantime.
+        async with self._start_lock:
+            conn, self._conn = self._conn, None
+            if conn is not None:
+                try:
+                    await conn.close()
+                except Exception as exc:
+                    errors.append(exc)
+            try:
+                await self._stop_daemon()
+            except Exception as exc:
+                errors.append(exc)
+            self._started = False
         if errors:
             raise ExceptionGroup("errors while closing Ringo client", errors)
 
@@ -857,7 +924,6 @@ class TaskQueue:
         if spec is None:
             spec = TaskSpec(**kwargs)
         request = queue_pb2.EnqueueRequest(queue=self.name, task=spec.to_proto())
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call("Enqueue", request)
         return EnqueueResult.from_proto(response.result)
 
@@ -882,7 +948,6 @@ class TaskQueue:
             request = queue_pb2.EnqueueManyRequest(
                 queue=self.name, tasks=[s.to_proto() for s in chunk]
             )
-            request.context.CopyFrom(self._client._require_conn().request_context())
             response = await self._client._require_conn().call("EnqueueMany", request)
             all_results.extend(EnqueueResult.from_proto(r) for r in response.results)
             total_created += response.created_count
@@ -900,7 +965,6 @@ class TaskQueue:
         if not task_id:
             raise InvalidArgumentError("task_id is required")
         request = queue_pb2.GetTaskRequest(queue=self.name, task_id=task_id)
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call(
             "GetTask", request, retry_safe=True
         )
@@ -910,7 +974,6 @@ class TaskQueue:
         if not unique_key:
             raise InvalidArgumentError("unique_key is required")
         request = queue_pb2.GetTaskRequest(queue=self.name, unique_key=unique_key)
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call(
             "GetTask", request, retry_safe=True
         )
@@ -1009,7 +1072,6 @@ class TaskQueue:
         request = queue_pb2.GetQueueStatsRequest(
             queue=self.name, task_types=selected or []
         )
-        request.context.CopyFrom(conn.request_context())
         response = await conn.call("GetQueueStats", request, retry_safe=True)
         return QueueStats.from_proto(response)
 
@@ -1023,7 +1085,6 @@ class TaskQueue:
             raise InvalidArgumentError("unique_key is required")
         request = queue_pb2.RetryTaskRequest(queue=self.name)
         request.unique_key = unique_key
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call("RetryTask", request)
         return Task.from_proto(response.task)
 
@@ -1033,7 +1094,6 @@ class TaskQueue:
             raise InvalidArgumentError("task_id is required")
         request = queue_pb2.RetryTaskRequest(queue=self.name)
         request.task_id = task_id
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call("RetryTask", request)
         return Task.from_proto(response.task)
 
@@ -1067,7 +1127,6 @@ class TaskQueue:
             request.wait_timeout.CopyFrom(_pb.td_to_dur(wait_timeout))
         if lease_duration is not None:
             request.lease_duration.CopyFrom(_pb.td_to_dur(lease_duration))
-        request.context.CopyFrom(self._client._require_conn().request_context())
         response = await self._client._require_conn().call("Claim", request)
         return [LeasedTask(self, t) for t in response.tasks]
 
@@ -1083,6 +1142,7 @@ class TaskQueue:
         lease_duration: timedelta = timedelta(seconds=60),
         grace_period: timedelta = timedelta(seconds=30),
         on_idle: Callable[[], Any] | None = None,
+        task_timeouts: TaskTimeouts | None = None,
     ) -> "Worker":
         """Create a credit-controlled worker over the frozen Work stream.
 
@@ -1114,7 +1174,9 @@ class TaskQueue:
             raise InvalidArgumentError(
                 "concurrency must be an integer, ResourceConcurrency, or AdaptiveConcurrency"
             )
-        if isinstance(concurrency, int) and (concurrency < 1 or concurrency > 2**31 - 1):
+        if isinstance(concurrency, int) and (
+            concurrency < 1 or concurrency > 2**31 - 1
+        ):
             raise InvalidArgumentError(
                 "initial fixed concurrency must be between 1 and 2**31-1"
             )
@@ -1124,6 +1186,8 @@ class TaskQueue:
             raise InvalidArgumentError("lease_duration must not be negative")
         if grace_period < timedelta(0):
             raise InvalidArgumentError("grace_period must not be negative")
+        if task_timeouts is not None and not isinstance(task_timeouts, TaskTimeouts):
+            raise InvalidArgumentError("task_timeouts must be TaskTimeouts or None")
         worker = Worker(
             self,
             handler,
@@ -1133,6 +1197,7 @@ class TaskQueue:
             grace_period=grace_period,
             on_idle=on_idle,
             weight=weight,
+            task_timeouts=task_timeouts,
         )
         self._client._register_worker(worker)
         return worker
@@ -1167,7 +1232,6 @@ class LeasedTask:
     async def _terminal(self, name: str, request: Any) -> None:
         request.lease.CopyFrom(self._ref())
         conn = self._queue._client._require_conn()
-        request.context.CopyFrom(conn.request_context())
         await conn.call(name, request)
 
     async def ack(self) -> None:
@@ -1193,6 +1257,5 @@ class LeasedTask:
         request.extension.CopyFrom(_pb.td_to_dur(extension))
         request.lease.CopyFrom(self._ref())
         conn = self._queue._client._require_conn()
-        request.context.CopyFrom(conn.request_context())
         response = await conn.call("ExtendLease", request)
         self._lease_until = _pb.ts_to_dt(response.lease.until)

ringo_task_queue/shared.py

@@ -8,6 +8,7 @@ import os
 import re
 import stat
 import subprocess
+import threading
 import time
 from pathlib import Path
 from typing import Any, cast
@@ -131,6 +132,46 @@ def validate_shared_state(value: object) -> dict[str, object] | None:
     return value
 
 
+class _DetachedProcess:
+    """Observe a shared child without giving an event loop kill ownership.
+
+    asyncio subprocess transports kill running children when finalized. A
+    shared daemon must instead outlive its launcher's loop, object and process.
+    The daemon reaper retains Popen until natural exit, reaping POSIX zombies
+    without joining the interpreter or its default asyncio executor.
+    """
+
+    def __init__(self, process: subprocess.Popen[bytes]) -> None:
+        self._process = process
+        threading.Thread(
+            target=process.wait,
+            name=f"ringo-shared-reaper-{process.pid}",
+            daemon=True,
+        ).start()
+
+    @classmethod
+    def spawn(cls, command: list[str], **options: Any) -> _DetachedProcess:
+        return cls(subprocess.Popen(command, **options))
+
+    @property
+    def pid(self) -> int:
+        return self._process.pid
+
+    @property
+    def returncode(self) -> int | None:
+        return self._process.poll()
+
+    def kill(self) -> None:
+        # Explicit test/recovery callers still target this exact Popen child.
+        # Neither object finalization nor SharedLocalDaemon.stop calls this.
+        self._process.kill()
+
+    async def wait(self) -> int:
+        while (code := self._process.poll()) is None:
+            await asyncio.sleep(0.05)
+        return code
+
+
 class SharedLocalDaemon:
     """Discover or start a shared-local daemon without owning its process.
 
@@ -149,7 +190,7 @@ class SharedLocalDaemon:
         self.data_dir = Path(data_dir).absolute()
         self._daemon_path = daemon_path
         self._startup_timeout = startup_timeout
-        self._launcher: asyncio.subprocess.Process | None = None
+        self._launcher: _DetachedProcess | None = None
         self.endpoint: str | None = None
         self.token: str | None = None
         self.generation: str | None = None
@@ -240,18 +281,29 @@ class SharedLocalDaemon:
             # DEVNULL is intentional: no lifecycle stdin and no inherited pipe
             # can keep either process alive or block after this SDK exits.
             detach: dict[str, Any] = (
-                {"creationflags": subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP}  # type: ignore[attr-defined]
+                {"creationflags": (
+                    subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore[attr-defined]
+                    | (subprocess.CREATE_BREAKAWAY_FROM_JOB  # type: ignore[attr-defined]
+                       if os.environ.get("RINGO_SUPERVISOR_JOB") == "1" else 0)
+                )}
                 if os.name == "nt" else {"start_new_session": True}
             )
-            self._launcher = await asyncio.create_subprocess_exec(
-                *command,
-                stdin=asyncio.subprocess.DEVNULL,
-                stdout=asyncio.subprocess.DEVNULL,
-                stderr=asyncio.subprocess.DEVNULL,
+            # No asyncio transport: its finalizer would kill a still-running
+            # shared daemon when this client/loop is collected or exits.
+            self._launcher = _DetachedProcess.spawn(
+                command,
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                close_fds=True,
                 **detach,
             )
         except OSError as exc:
-            raise RingoError(f"failed to spawn shared daemon {binary}: {exc}") from exc
+            required = (
+                " (supervisor Job breakaway required; no inherited-Job fallback)"
+                if os.name == "nt" and os.environ.get("RINGO_SUPERVISOR_JOB") == "1" else ""
+            )
+            raise RingoError(f"failed to spawn shared daemon {binary}{required}: {exc}") from exc
 
     async def _endpoint_reachable(self, endpoint: str) -> bool:
         try: