ringo-task-queue 0.1.0.dev2__py3-none-win_amd64.whl → 0.1.0.dev4__py3-none-win_amd64.whl

ringo_task_queue/worker.py

@@ -17,6 +17,7 @@ import grpc
 
 from . import _pb
 from ._proto.ringo.v1 import queue_pb2
+from .concurrency import AdaptiveConcurrency, ResourceConcurrency, _Member
 from .errors import (
     IncompatibleVersionError,
     LeaseLostError,
@@ -145,6 +146,10 @@ class _Assignment:
     renew_task: asyncio.Task[Any] | None = None
     alive: bool = True
     completion_started: bool = False
+    completion_success: bool = False
+    completion_congestion: bool = False
+    controller_released: bool = False
+    fixed_released: bool = False
 
     @property
     def ref(self) -> queue_pb2.LeaseRef:
@@ -170,18 +175,33 @@ class Worker:
         handler: Callable[[Task, TaskContext], Awaitable[Any]],
         *,
         task_types: Iterable[str],
-        concurrency: int,
+        concurrency: int | ResourceConcurrency | AdaptiveConcurrency,
         lease_duration: timedelta,
         grace_period: timedelta,
         on_idle: Callable[[], Any] | None = None,
+        weight: int = 1,
     ) -> None:
-        if concurrency < 1:
-            raise ValueError("concurrency must be >= 1")
+        if isinstance(concurrency, bool) or not isinstance(
+            concurrency, (int, ResourceConcurrency, AdaptiveConcurrency)
+        ):
+            raise ValueError(
+                "concurrency must be an integer, ResourceConcurrency, or AdaptiveConcurrency"
+            )
+        if isinstance(concurrency, int) and concurrency < 1:
+            raise ValueError("initial fixed concurrency must be >= 1")
+        if isinstance(weight, bool) or not isinstance(weight, int) or weight < 1:
+            raise ValueError("weight must be a positive integer")
         self._queue = queue
         self._client = queue._client
         self._handler = handler
         self._task_types = list(task_types)
-        self._concurrency = concurrency
+        self._controller: ResourceConcurrency | AdaptiveConcurrency | None = (
+            concurrency if not isinstance(concurrency, int) else None
+        )
+        self._member: _Member | None = None
+        self._concurrency = concurrency if isinstance(concurrency, int) else max(1, concurrency.capacity)
+        self._configured_concurrency = self._concurrency
+        self._weight = weight
         self._lease_duration = lease_duration
         self._grace_period = max(0.0, grace_period.total_seconds())
         self._worker_id = self._new_worker_id()
@@ -222,6 +242,24 @@ class Worker:
         self._next_request_id = 1
         self._pending: dict[int, asyncio.Future[queue_pb2.WorkResult]] = {}
         self._inflight: dict[str, _Assignment] = {}
+        # Handler lifetime, not just lease-map membership, owns local capacity.
+        # Old-stream handlers that swallow cancellation remain here until their
+        # coroutine actually settles.
+        self._handler_tasks: set[asyncio.Task[Any]] = set()
+        self._fixed_running = 0
+        # Every capacity update shares one publisher mutex and a monotonically
+        # increasing session epoch.  A late task from an old stream can never
+        # publish credit onto a replacement stream or after close.
+        self._publisher_lock = asyncio.Lock()
+        self._session_epoch = 0
+        self._publisher_closed = False
+        self._probe: tuple[_Member, int] | None = None
+        self._detached = False
+        self._controller_registered = False
+        self._fixed_dynamic_target = False
+        if self._controller is not None:
+            self._member = self._controller._attach(self, weight)
+            self._controller_registered = True
 
     @staticmethod
     def _new_worker_id() -> str:
@@ -236,8 +274,35 @@ class Worker:
 
     @property
     def concurrency(self) -> int:
+        """Current fixed target, or the controller's current capacity."""
+        if self._controller is not None:
+            return self._controller.capacity
         return self._concurrency
 
+    @property
+    def weight(self) -> int:
+        return self._weight
+
+    async def set_concurrency(self, concurrency: int) -> None:
+        """Resize a fixed worker without cancelling running handlers."""
+        if self._controller is not None:
+            raise ValueError("a Worker attached to a shared controller cannot resize itself")
+        if isinstance(concurrency, bool) or not isinstance(concurrency, int):
+            raise ValueError("concurrency must be an integer")
+        if concurrency < 0 or concurrency > 2**31 - 1:
+            raise ValueError("concurrency must be between 0 and 2**31-1")
+        if self._run_started:
+            conn = self._client._require_conn()
+            if conn._session_minor < 2:
+                raise IncompatibleVersionError(
+                    "runtime fixed concurrency requires protocol minor >= 2; "
+                    f"negotiated minor is {conn._session_minor}"
+                )
+        self._concurrency = concurrency
+        self._fixed_dynamic_target = True
+        if self._session_done is not None and not self._session_done.is_set():
+            await self._publish_fixed_capacity()
+
     async def run(self) -> None:
         """Run until :meth:`close` is called, reconnecting boundedly."""
         if self._run_started:
@@ -245,6 +310,15 @@ class Worker:
         self._run_started = True
         failures = 0
         try:
+            conn = self._client._require_conn()
+            if (
+                (self._controller is not None or self._fixed_dynamic_target)
+                and conn._session_minor < 2
+            ):
+                raise IncompatibleVersionError(
+                    "dynamic concurrency requires protocol minor >= 2; "
+                    f"negotiated minor is {conn._session_minor}"
+                )
             if self._on_idle is not None:
                 # GetQueueActivity arrived in protocol minor 1; fail fast
                 # instead of silently running without idle observation.
@@ -304,6 +378,9 @@ class Worker:
         async with self._close_lock:
             if self._close_task is None:
                 self._close_started.set()
+                # The controller must linearize detach and withdraw its probe
+                # before this gate closes publication; the close path uses the
+                # still-live stream to await that correlated WorkResult.
                 self._idle_candidate = False
                 self._idle_zero_since = None
                 self._idle_session_token = None
@@ -316,6 +393,11 @@ class Worker:
 
     async def _close_impl(self) -> None:
         try:
+            if self._controller is not None and self._member is not None:
+                await self._controller._detach(self._member)
+                self._detached = True
+            self._publisher_closed = True
+            self._session_epoch += 1
             session = self._session_done
             # Drain is best effort. A disconnected stream cannot safely report
             # any terminal result, and leases are deliberately left to expire.
@@ -357,7 +439,7 @@ class Worker:
                 self._close_finished.set()
 
     async def _wait_handlers(self) -> None:
-        tasks = [a.handler_task for a in self._inflight.values() if a.handler_task]
+        tasks = list(self._handler_tasks)
         if tasks:
             await asyncio.gather(*tasks, return_exceptions=True)
 
@@ -368,21 +450,23 @@ class Worker:
         self._idle_session_token = None
         assignments = list(self._inflight.values())
         self._inflight.clear()
+        renewals: list[asyncio.Task[Any]] = []
         for assignment in assignments:
             self._invalidate_assignment(assignment)
-        tasks: list[asyncio.Task[Any]] = []
+            handler = assignment.handler_task
+            if handler is not None and not handler.done():
+                # Cancellation is advisory.  Never let a handler that catches
+                # CancelledError block reconnect or forced close; its local
+                # capacity remains occupied until its coroutine really ends.
+                handler.cancel()
+            renewal = assignment.renew_task
+            if renewal is not None and not renewal.done():
+                renewal.cancel()
+                renewals.append(renewal)
+        if renewals:
+            await asyncio.gather(*renewals, return_exceptions=True)
         for assignment in assignments:
-            if (
-                assignment.handler_task is not None
-                and not assignment.handler_task.done()
-            ):
-                assignment.handler_task.cancel()
-                tasks.append(assignment.handler_task)
-            if assignment.renew_task is not None and not assignment.renew_task.done():
-                assignment.renew_task.cancel()
-                tasks.append(assignment.renew_task)
-        if tasks:
-            await asyncio.gather(*tasks, return_exceptions=True)
+            await self._release_capacity_if_settled(assignment)
 
     def _invalidate_assignment(self, assignment: _Assignment) -> None:
         if not assignment.alive:
@@ -391,6 +475,49 @@ class Worker:
         if assignment.context is not None:
             assignment.context._cancel()
 
+    async def _release_controller_permit(
+        self,
+        assignment: _Assignment,
+        *,
+        success: bool = False,
+        congestion: bool = False,
+    ) -> None:
+        if (
+            self._controller is None
+            or self._member is None
+            or assignment.controller_released
+        ):
+            return
+        assignment.controller_released = True
+        await self._controller._release_running(
+            self._member, success=success, congestion=congestion
+        )
+
+    async def _release_capacity_if_settled(
+        self, assignment: _Assignment, *, handler_settling: bool = False
+    ) -> None:
+        """Release capacity only as the application handler settles."""
+        if assignment.alive:
+            return
+        handler = assignment.handler_task
+        current = asyncio.current_task()
+        if handler is not None and not handler.done():
+            if handler is not current or not handler_settling:
+                return
+        if self._controller is not None:
+            await self._release_controller_permit(
+                assignment,
+                success=assignment.completion_success,
+                congestion=assignment.completion_congestion,
+            )
+            return
+        if assignment.fixed_released:
+            return
+        assignment.fixed_released = True
+        self._fixed_running = max(0, self._fixed_running - 1)
+        if not self._close_started.is_set():
+            await self._publish_fixed_capacity()
+
     async def _run_session_monitored(self) -> None:
         """Run one Work session, observing queue activity when on_idle is set.
 
@@ -437,7 +564,7 @@ class Worker:
         ):
             return False
         # Local assignments make the queue non-idle for this worker.
-        if self._inflight:
+        if self._inflight or self._handler_tasks:
             self._idle_zero_since = None
             return False
         session = self._session_done
@@ -491,6 +618,7 @@ class Worker:
         # query in-flight window cannot produce a false trigger.
         if (
             self._inflight
+            or self._handler_tasks
             or self._close_started.is_set()
             or self._force_cancel.is_set()
         ):
@@ -525,6 +653,7 @@ class Worker:
 
     async def _run_session(self) -> None:
         conn = self._client._require_conn()
+        self._session_epoch += 1
         # A new session invalidates any in-flight idle observation.
         self._idle_candidate = False
         self._idle_zero_since = None
@@ -545,7 +674,7 @@ class Worker:
             queue=self._queue.name,
             worker_id=self._worker_id,
             task_types=self._task_types,
-            max_concurrency=self._concurrency,
+            max_concurrency=max(1, self._concurrency),
         )
         if self._lease_duration > timedelta(0):
             reg.lease_duration.CopyFrom(_pb.td_to_dur(self._lease_duration))
@@ -571,7 +700,11 @@ class Worker:
             if self._force_cancel.is_set():
                 return
             # Initial credit is sent only after WorkerRegistered was observed.
-            await self._send_credit(self._concurrency)
+            if self._controller is None:
+                await self._send_credit(self._concurrency)
+            else:
+                self._controller._check_loop()
+                await self._controller._start_member(self._member)
             if self._on_idle is not None and not self._close_started.is_set():
                 # Eligible only from the first/current WorkerRegistered plus
                 # a successful initial credit of this session. The candidate
@@ -628,6 +761,8 @@ class Worker:
             if not future.done():
                 future.set_exception(UnavailableError("Work stream closed"))
         self._pending.clear()
+        if self._controller is not None and self._member is not None:
+            await self._controller._stream_lost(self._member)
 
     async def _enqueue_frame(self, frame: queue_pb2.WorkRequest) -> None:
         outbox = self._outbox
@@ -666,7 +801,11 @@ class Worker:
                             _pb.dur_to_td(registered.lease_duration)
                             or self._lease_duration
                         )
-                    if registered.max_concurrency:
+                    if (
+                        self._controller is None
+                        and not self._fixed_dynamic_target
+                        and registered.max_concurrency
+                    ):
                         self._concurrency = registered.max_concurrency
                     if self._registered is not None:
                         self._registered.set()
@@ -722,13 +861,24 @@ class Worker:
         # existing assignment if a malformed peer sends a duplicate.
         if task.id in self._inflight:
             return
+        if self._controller is not None and self._member is not None:
+            # Consuming the reserved probe is the linearization point at which
+            # an assignment becomes a running controller permit.
+            if not await self._controller._probe_consumed(self._member):
+                # A withdrawn probe must not become a new handler.  The lease
+                # remains server-owned and will naturally expire if this was a
+                # transport race.
+                return
         self._inflight[task.id] = assignment
+        if self._controller is None:
+            self._fixed_running += 1
         assignment.renew_task = asyncio.create_task(
             self._renew_loop(assignment), name=f"ringo-renew-{task.id}"
         )
         assignment.handler_task = asyncio.create_task(
             self._handle_assignment(assignment), name=f"ringo-handler-{task.id}"
         )
+        self._handler_tasks.add(assignment.handler_task)
         assignment.handler_task.add_done_callback(_consume_task_exception)
 
     async def _handle_assignment(self, assignment: _Assignment) -> None:
@@ -753,6 +903,13 @@ class Worker:
         except Exception as exc:
             if assignment.alive and not self._force_cancel.is_set():
                 await self._finish(assignment, "nack", exc)
+        finally:
+            current = asyncio.current_task()
+            if current is not None:
+                self._handler_tasks.discard(current)
+            await self._release_capacity_if_settled(
+                assignment, handler_settling=True
+            )
 
     async def _finish(
         self,
@@ -811,15 +968,26 @@ class Worker:
             # protocol; leave the assignment fenced until the stream/server does.
             assignment.completion_started = False
             return
-        await self._release_assignment(assignment)
-        await self._send_credit(self._concurrency - len(self._inflight))
+        congestion = isinstance(error, RetryTaskError) and error.congestion
+        await self._release_assignment(
+            assignment,
+            success=operation == "ack",
+            congestion=congestion,
+        )
 
     async def _release_assignment(
-        self, assignment: _Assignment, *, cancel_context: bool = False
+        self,
+        assignment: _Assignment,
+        *,
+        cancel_context: bool = False,
+        success: bool = False,
+        congestion: bool = False,
     ) -> None:
         if not assignment.alive:
             return
         assignment.alive = False
+        assignment.completion_success = success
+        assignment.completion_congestion = congestion
         self._inflight.pop(assignment.task.id, None)
         if cancel_context and assignment.context is not None:
             assignment.context._cancel()
@@ -829,21 +997,22 @@ class Worker:
         ):
             assignment.renew_task.cancel()
             await asyncio.gather(assignment.renew_task, return_exceptions=True)
+        await self._release_capacity_if_settled(assignment)
 
     async def _lease_lost(self, assignment: _Assignment) -> None:
         if not assignment.alive:
             return
         await self._release_assignment(assignment, cancel_context=True)
-        current = assignment.handler_task
+        handler = assignment.handler_task
         if (
-            current is not None
-            and current is not asyncio.current_task()
-            and not current.done()
+            handler is not None
+            and handler is not asyncio.current_task()
+            and not handler.done()
         ):
-            current.cancel()
-            await asyncio.gather(current, return_exceptions=True)
-        if not self._close_started.is_set():
-            await self._send_credit(self._concurrency - len(self._inflight))
+            # Do not await an application coroutine that may intentionally
+            # swallow cancellation. Its capacity is released by its finally.
+            handler.cancel()
+        await self._release_capacity_if_settled(assignment)
 
     async def _renew_loop(self, assignment: _Assignment) -> None:
         while assignment.alive and not self._force_cancel.is_set():
@@ -905,35 +1074,169 @@ class Worker:
             raise error_from_detail(detail)
 
     async def _send_credit(self, available: int) -> None:
-        if (
-            self._close_started.is_set()
-            or self._session_done is None
-            or self._session_done.is_set()
-        ):
-            return
-        available = max(0, available)
-        # The server can have an assignment queued in the transport before the
-        # single reader materializes it locally. A very fast handler may then
-        # briefly advertise too much capacity. The stable metadata reason lets
-        # clients retry that exact non-mutating absolute update without parsing
-        # human-readable error text.
+        # ``available`` is intentionally not trusted across retries.  A task
+        # may arrive between a rejected snapshot and its retry, so every
+        # credit-capacity retry samples local in-flight state again.
+        await self._publish_fixed_capacity()
+
+    async def _publish_fixed_capacity(self) -> None:
+        async with self._publisher_lock:
+            epoch = self._session_epoch
+            for attempt in range(20):
+                if (
+                    self._publisher_closed
+                    or self._close_started.is_set()
+                    or epoch != self._session_epoch
+                    or self._session_done is None
+                    or self._session_done.is_set()
+                ):
+                    return
+                available = max(0, self._concurrency - self._fixed_running)
+                update = queue_pb2.CreditUpdate(available=available)
+                if self._fixed_dynamic_target:
+                    update.target_concurrency = self._concurrency
+                result = await self._send_command(
+                    queue_pb2.WorkRequest(credit=update), wait=True
+                )
+                if (
+                    result is None
+                    or not result.HasField("error")
+                    or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
+                ):
+                    return
+                if (
+                    result.error.metadata.get("reason") != "credit_capacity"
+                    or attempt == 19
+                ):
+                    raise error_from_detail(result.error)
+                await asyncio.sleep(0.01)
+
+    async def _publish_controller_initial(self, capacity: int) -> None:
+        async with self._publisher_lock:
+            if (
+                self._publisher_closed
+                or self._close_started.is_set()
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                raise UnavailableError("Work stream is not connected")
+            ok = await self._publish_dynamic_locked(capacity, 0)
+            if not ok:
+                raise UnavailableError("initial dynamic credit was not accepted")
+
+    async def _publish_controller_probe(
+        self, member: _Member, capacity: int, generation: int
+    ) -> None:
+        async with self._publisher_lock:
+            if (
+                self._publisher_closed
+                or self._close_started.is_set()
+                or self._member is not member
+                or not member.probe
+                or member.probe_generation != generation
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                return
+            await self._publish_dynamic_locked(
+                capacity,
+                1,
+                probe_member=member,
+                probe_generation=generation,
+            )
+
+    async def _publish_controller_target(self, capacity: int) -> bool:
+        """Publish target/zero credit for a live shared-controller stream."""
+        async with self._publisher_lock:
+            if self._publisher_closed:
+                return True
+            if self._session_done is None or self._session_done.is_set():
+                return True
+            return await self._publish_dynamic_locked(capacity, 0)
+
+    async def _withdraw_controller_probe(
+        self,
+        member: _Member,
+        capacity: int,
+        *,
+        force_target: bool = False,
+    ) -> bool:
+        async with self._publisher_lock:
+            if self._publisher_closed:
+                return True
+            if self._session_done is None or self._session_done.is_set():
+                # The old stream is definitively gone; its credit no longer
+                # exists server-side and the controller may reclaim the token.
+                return True
+            return await self._publish_dynamic_locked(
+                capacity,
+                0,
+                allow_closing=True,
+                explicit_target=force_target,
+            )
+
+    async def _publish_dynamic_locked(
+        self,
+        target: int,
+        available: int,
+        *,
+        allow_closing: bool = False,
+        explicit_target: bool = False,
+        probe_member: _Member | None = None,
+        probe_generation: int = 0,
+    ) -> bool:
+        epoch = self._session_epoch
         for attempt in range(20):
-            request = queue_pb2.WorkRequest(
-                credit=queue_pb2.CreditUpdate(available=available)
+            if (
+                self._publisher_closed
+                or (self._close_started.is_set() and not allow_closing)
+                or epoch != self._session_epoch
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                return self._session_done is None or self._session_done.is_set()
+            if probe_member is not None and (
+                not probe_member.probe
+                or probe_member.probe_generation != probe_generation
+            ):
+                # The assignment consumed this probe while a previous
+                # capacity update was in flight; never replay available=1.
+                return True
+            # target and available are one atomic protocol snapshot.  A
+            # credit_capacity rejection must be retried with a fresh target /
+            # in-flight sample, never by replaying the old snapshot.
+            current_target = (
+                self._controller.capacity
+                if self._controller is not None and not explicit_target
+                else target
+            )
+            current_available = available
+            if self._controller is None:
+                # An assignment may have been sent by the server before the
+                # rejected credit result reached this reader. Re-sample the
+                # now-materialized local assignments instead of replaying the
+                # stale fixed-capacity snapshot.
+                current_available = max(0, current_target - self._fixed_running)
+            update = queue_pb2.CreditUpdate(
+                available=max(0, current_available),
+                target_concurrency=current_target,
+            )
+            result = await self._send_command(
+                queue_pb2.WorkRequest(credit=update), wait=True
             )
-            result = await self._send_command(request, wait=True)
             if (
                 result is None
                 or not result.HasField("error")
                 or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
             ):
-                return
+                return True
             if (
                 result.error.metadata.get("reason") != "credit_capacity"
                 or attempt == 19
             ):
                 raise error_from_detail(result.error)
             await asyncio.sleep(0.01)
+        return False
 
     def _allocate_request_id(self) -> int:
         rid = self._next_request_id