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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -9,15 +9,17 @@ from __future__ import annotations
9
9
  import asyncio
10
10
  import inspect
11
11
  import math
12
+ import sys
12
13
  from collections import deque
13
- from dataclasses import dataclass
14
+ from dataclasses import dataclass, field
14
15
  from datetime import datetime, timedelta, timezone
15
- from typing import Any, Awaitable, Callable, Iterable
16
+ from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterable
16
17
 
17
18
  import grpc
18
19
 
19
20
  from . import _pb
20
21
  from ._proto.ringo.v1 import queue_pb2
22
+ from ._supervisor_policy import TaskTimeouts
21
23
  from .concurrency import AdaptiveConcurrency, ResourceConcurrency, _Member
22
24
  from .errors import (
23
25
  IncompatibleVersionError,
@@ -32,6 +34,16 @@ from .errors import (
32
34
  )
33
35
  from .models import Task
34
36
 
37
+ if TYPE_CHECKING:
38
+ from ._supervisor_watchdog import _WorkerWatchdog
39
+
40
+
41
+ def _current_runtime() -> Any:
42
+ # A normal SDK install does not import the process-supervision runtime.
43
+ module = sys.modules.get("ringo_task_queue._supervisor_runtime")
44
+ return None if module is None else module.current_runtime()
45
+
46
+
35
47
  _MAX_ERROR_BYTES = 256 * 1024
36
48
 
37
49
  #: Fixed on_idle monitor cadence (design §10.6). These are deliberately not
@@ -41,9 +53,15 @@ _IDLE_POLL_SECONDS = 0.25
41
53
  _IDLE_QUIET_SECONDS = 1.0
42
54
  #: GetQueueActivity was introduced in protocol minor 1.
43
55
  _IDLE_MIN_PROTOCOL_MINOR = 1
56
+ _ORDERED_REQUEST_IDS_MIN_PROTOCOL_MINOR = 6
57
+ _MAX_REQUEST_ID = 2**64 - 1
58
+
59
+
60
+ class _WatchdogFenced(Exception):
61
+ """A stopped deadline invocation abandoned during graceful drain."""
44
62
 
45
63
 
46
- def _consume_task_exception(task: asyncio.Task[Any]) -> None:
64
+ def _consume_task_exception(task: asyncio.Future[Any]) -> None:
47
65
  if not task.cancelled():
48
66
  task.exception()
49
67
 
@@ -125,12 +143,19 @@ class TaskContext:
125
143
  async def wait_cancelled(self) -> None:
126
144
  await self._cancel_event.wait()
127
145
 
146
+ def mark_progress(self) -> None:
147
+ """Mark local business progress without sending any RPC."""
148
+ watchdog = self._worker._watchdog
149
+ if watchdog is not None:
150
+ watchdog.mark_progress(self._assignment)
151
+
128
152
  async def report_progress(
129
153
  self, current: float, total: float, message: str = ""
130
154
  ) -> None:
131
- """Report progress for this attempt and await its correlated result."""
155
+ """Mark local progress, report it and await its correlated result."""
132
156
  if self.cancelled:
133
157
  raise LeaseLostError("task lease is no longer valid")
158
+ self.mark_progress()
134
159
  await self._worker._report_progress(self._assignment, current, total, message)
135
160
 
136
161
  def _cancel(self) -> None:
@@ -156,6 +181,8 @@ class _Assignment:
156
181
  started: bool = False
157
182
  session_epoch: int = 0
158
183
  session: asyncio.Event | None = None
184
+ renew_deadline: float | None = None
185
+ watchdog_expired: bool = False
159
186
 
160
187
  @property
161
188
  def identity(self) -> tuple[str, int, str]:
@@ -170,6 +197,39 @@ class _Assignment:
170
197
  )
171
198
 
172
199
 
200
+ @dataclass
201
+ class _WorkSource:
202
+ """One Work call owns its correlation and transport lifetime.
203
+
204
+ Reader/writer tasks keep this object even after a replacement becomes the
205
+ current source. No receipt, EOF or timeout from an old call can touch the
206
+ replacement's pending Futures, registration or transport.
207
+ """
208
+
209
+ epoch: int
210
+ call: Any
211
+ minor: int = 0
212
+ token: str = ""
213
+ ordered_request_ids: bool = False
214
+ done: asyncio.Event = field(default_factory=asyncio.Event)
215
+ registered: asyncio.Event = field(default_factory=asyncio.Event)
216
+ outbox: asyncio.Queue[queue_pb2.WorkRequest | None] = field(
217
+ default_factory=asyncio.Queue
218
+ )
219
+ pending: dict[int, asyncio.Future[queue_pb2.WorkResult]] = field(
220
+ default_factory=dict
221
+ )
222
+ next_request_id: int = 1
223
+ error: BaseException | None = None
224
+ reader: asyncio.Task[Any] | None = None
225
+ writer: asyncio.Task[Any] | None = None
226
+ publisher: asyncio.Task[None] | None = None
227
+ cleanup: asyncio.Task[None] | None = None
228
+ # Diagnostics only: accepted credit never participates in admission itself.
229
+ accepted_credit: int = 0
230
+ credit_generation: int = -1
231
+
232
+
173
233
  class Worker:
174
234
  """A credit-controlled asynchronous task worker.
175
235
 
@@ -190,6 +250,7 @@ class Worker:
190
250
  grace_period: timedelta,
191
251
  on_idle: Callable[[], Any] | None = None,
192
252
  weight: int = 1,
253
+ task_timeouts: TaskTimeouts | None = None,
193
254
  ) -> None:
194
255
  if isinstance(concurrency, bool) or not isinstance(
195
256
  concurrency, (int, ResourceConcurrency, AdaptiveConcurrency)
@@ -214,7 +275,6 @@ class Worker:
214
275
  if isinstance(concurrency, int)
215
276
  else max(1, concurrency.capacity)
216
277
  )
217
- self._configured_concurrency = self._concurrency
218
278
  self._weight = weight
219
279
  self._lease_duration = lease_duration
220
280
  self._grace_period = max(0.0, grace_period.total_seconds())
@@ -246,15 +306,10 @@ class Worker:
246
306
  self._command_timeout = 30.0
247
307
  self._close_lock = asyncio.Lock()
248
308
  self._close_task: asyncio.Task[None] | None = None
249
- self._session_done: asyncio.Event | None = None
250
- self._call: Any = None
251
- self._outbox: asyncio.Queue[Any] | None = None
252
- self._reader_task: asyncio.Task[Any] | None = None
253
- self._writer_task: asyncio.Task[Any] | None = None
254
- self._registered: asyncio.Event | None = None
255
- self._stream_error: BaseException | None = None
256
- self._next_request_id = 1
257
- self._pending: dict[int, asyncio.Future[queue_pb2.WorkResult]] = {}
309
+ self._source: _WorkSource | None = None
310
+ # Recovery history is distinct from a live source's error. A clean EOF
311
+ # remains clean, while exhausted recovery can report its earlier cause.
312
+ self._last_stream_error: BaseException | None = None
258
313
  # A successor lease can arrive before the predecessor's NACK receipt
259
314
  # is consumed. Only an identical complete LeaseRef is a duplicate.
260
315
  self._inflight: dict[tuple[str, int, str], _Assignment] = {}
@@ -271,13 +326,32 @@ class Worker:
271
326
  self._publisher_lock = asyncio.Lock()
272
327
  self._session_epoch = 0
273
328
  self._publisher_closed = False
274
- self._probe: tuple[_Member, int] | None = None
275
- self._detached = False
276
- self._controller_registered = False
329
+ self._fixed_capacity_generation = 0
330
+ self._fixed_publication_task: asyncio.Task[None] | None = None
331
+ self._fixed_publication_source: tuple[int, asyncio.Event] | None = None
332
+ self._fixed_publication_waiters: list[tuple[int, asyncio.Future[None]]] = []
277
333
  self._fixed_dynamic_target = False
334
+ if task_timeouts is not None and not isinstance(task_timeouts, TaskTimeouts):
335
+ raise ValueError("task_timeouts must be TaskTimeouts or None")
336
+ self._runtime = _current_runtime()
337
+ self._task_timeouts = (
338
+ self._runtime.task_timeouts
339
+ if task_timeouts is None and self._runtime is not None
340
+ else task_timeouts
341
+ )
342
+ self._watchdog: _WorkerWatchdog | None = None
343
+ self._watchdog_error: RuntimeError | None = None
344
+ if (
345
+ self._runtime is not None
346
+ and self._runtime.policy.no_start_timeout is not None
347
+ ) or (self._task_timeouts is not None and self._task_timeouts.enabled):
348
+ from ._supervisor_watchdog import _WorkerWatchdog
349
+
350
+ self._watchdog = _WorkerWatchdog(
351
+ self, self._task_timeouts or TaskTimeouts(), self._runtime
352
+ )
278
353
  if self._controller is not None:
279
354
  self._member = self._controller._attach(self, weight)
280
- self._controller_registered = True
281
355
 
282
356
  @staticmethod
283
357
  def _new_worker_id() -> str:
@@ -286,6 +360,42 @@ class Worker:
286
360
 
287
361
  return new_worker_id()
288
362
 
363
+ # Read-only private views keep diagnostics and existing benchmark probes
364
+ # useful without maintaining a second set of mutable transport state.
365
+ @property
366
+ def _session_done(self) -> asyncio.Event | None:
367
+ return self._source.done if self._source is not None else None
368
+
369
+ @property
370
+ def _call(self) -> Any:
371
+ return self._source.call if self._source is not None else None
372
+
373
+ @property
374
+ def _outbox(self) -> asyncio.Queue[queue_pb2.WorkRequest | None] | None:
375
+ return self._source.outbox if self._source is not None else None
376
+
377
+ @property
378
+ def _pending(self) -> dict[int, asyncio.Future[queue_pb2.WorkResult]]:
379
+ return self._source.pending if self._source is not None else {}
380
+
381
+ @property
382
+ def _registered(self) -> asyncio.Event | None:
383
+ return self._source.registered if self._source is not None else None
384
+
385
+ @property
386
+ def _reader_task(self) -> asyncio.Task[Any] | None:
387
+ return self._source.reader if self._source is not None else None
388
+
389
+ @property
390
+ def _writer_task(self) -> asyncio.Task[Any] | None:
391
+ return self._source.writer if self._source is not None else None
392
+
393
+ @property
394
+ def _stream_error(self) -> BaseException | None:
395
+ return (
396
+ self._source.error if self._source is not None else self._last_stream_error
397
+ )
398
+
289
399
  @property
290
400
  def worker_id(self) -> str:
291
401
  return self._worker_id
@@ -331,6 +441,10 @@ class Worker:
331
441
  self._run_started = True
332
442
  failures = 0
333
443
  try:
444
+ if self._runtime is not None:
445
+ self._runtime.register_worker(self)
446
+ if self._watchdog is not None:
447
+ self._watchdog.start()
334
448
  conn = self._client._require_conn()
335
449
  if (
336
450
  self._controller is not None or self._fixed_dynamic_target
@@ -367,7 +481,7 @@ class Worker:
367
481
  # downgrade that makes on_idle unservable; surface
368
482
  # it instead of burning the reconnect budget.
369
483
  raise
370
- self._stream_error = exc
484
+ self._last_stream_error = exc
371
485
  failures += 1
372
486
  if self._close_started.is_set():
373
487
  break
@@ -389,11 +503,19 @@ class Worker:
389
503
  )
390
504
  if not self._close_started.is_set():
391
505
  await asyncio.sleep(0)
506
+ if self._watchdog_error is not None:
507
+ raise self._watchdog_error
392
508
  finally:
393
509
  try:
394
510
  await self._cancel_assignments()
395
511
  finally:
396
- self._run_finished.set()
512
+ try:
513
+ if self._watchdog is not None:
514
+ await self._watchdog.stop()
515
+ finally:
516
+ if self._runtime is not None:
517
+ self._runtime.unregister_worker(self)
518
+ self._run_finished.set()
397
519
 
398
520
  async def close(self) -> None:
399
521
  """Drain and stop once; caller cancellation cannot cancel cleanup."""
@@ -417,7 +539,6 @@ class Worker:
417
539
  try:
418
540
  if self._controller is not None and self._member is not None:
419
541
  await self._controller._detach(self._member)
420
- self._detached = True
421
542
  self._publisher_closed = True
422
543
  # Drain keeps this Work session alive for handler ACK/renewal.
423
544
  # Close is not a session replacement: changing its epoch here
@@ -430,10 +551,10 @@ class Worker:
430
551
  await asyncio.wait_for(
431
552
  self._send_command(
432
553
  queue_pb2.WorkRequest(
433
- request_id=self._allocate_request_id(),
434
554
  drain=queue_pb2.WorkerDrain(begin=True),
435
555
  ),
436
- wait=True,
556
+ expected_epoch=self._session_epoch,
557
+ expected_session=session,
437
558
  ),
438
559
  timeout=max(
439
560
  0.1,
@@ -690,9 +811,36 @@ class Worker:
690
811
  idle candidate) under the fresh token. It is neither a callback
691
812
  failure nor a reconnect-worthy stream error.
692
813
  """
814
+ source = self._source
815
+ if source is None or source.done is not session:
816
+ return
693
817
  self._idle_candidate = False
694
818
  self._idle_zero_since = None
695
- session.set()
819
+ self._end_source(source)
820
+
821
+ def _end_source(
822
+ self, source: _WorkSource, error: BaseException | None = None
823
+ ) -> None:
824
+ """Synchronously end exactly one source before any cleanup wait.
825
+
826
+ Failing pending Futures here replaces per-command session-wait tasks.
827
+ A graceful Worker.close merely drains; it calls this only when the
828
+ transport actually ends or the grace period has expired.
829
+ """
830
+ if error is not None and source.error is None:
831
+ source.error = error
832
+ if self._source is source:
833
+ self._last_stream_error = error
834
+ source.done.set()
835
+ failure = source.error or UnavailableError("Work stream closed")
836
+ for future in source.pending.values():
837
+ if not future.done():
838
+ future.set_exception(failure)
839
+ source.pending.clear()
840
+ try:
841
+ source.call.cancel()
842
+ except Exception:
843
+ pass
696
844
 
697
845
  async def _run_session(self) -> None:
698
846
  conn = self._client._require_conn()
@@ -721,27 +869,48 @@ class Worker:
721
869
  async with conn._lock:
722
870
  stub = conn._require_stub()
723
871
  reg.context.CopyFrom(conn.request_context())
872
+ # Mode belongs to this registration snapshot, not to a later unary
873
+ # negotiation that may rotate the connection's token and minor.
874
+ minor = reg.context.protocol_minor
875
+ if (
876
+ self._controller is not None or self._fixed_dynamic_target
877
+ ) and minor < 2:
878
+ raise IncompatibleVersionError(
879
+ "dynamic concurrency requires protocol minor >= 2; "
880
+ f"negotiated minor is {minor}"
881
+ )
882
+ if self._on_idle is not None and minor < _IDLE_MIN_PROTOCOL_MINOR:
883
+ raise IncompatibleVersionError(
884
+ "on_idle requires a server speaking protocol minor"
885
+ f" >= {_IDLE_MIN_PROTOCOL_MINOR}; negotiated minor is {minor}"
886
+ )
887
+ ordered = minor >= _ORDERED_REQUEST_IDS_MIN_PROTOCOL_MINOR
888
+ if ordered:
889
+ reg.ordered_request_ids = True
724
890
  call = stub.Work(metadata=conn._metadata())
725
- self._call = call
726
- self._session_done = asyncio.Event()
727
- self._registered = asyncio.Event()
728
- self._stream_error = None
729
- self._outbox = asyncio.Queue()
730
- self._pending.clear()
731
- self._writer_task = asyncio.create_task(
732
- self._write_loop(call), name="ringo-worker-writer"
891
+ source = _WorkSource(
892
+ self._session_epoch,
893
+ call,
894
+ minor=minor,
895
+ token=reg.context.session_token,
896
+ ordered_request_ids=ordered,
733
897
  )
734
- self._reader_task = asyncio.create_task(
735
- self._read_loop(call), name="ringo-worker-reader"
898
+ self._source = source
899
+ self._last_stream_error = None
900
+ source.writer = asyncio.create_task(
901
+ self._write_loop(source), name="ringo-worker-writer"
902
+ )
903
+ source.reader = asyncio.create_task(
904
+ self._read_loop(source), name="ringo-worker-reader"
736
905
  )
737
906
  try:
738
- await self._enqueue_frame(queue_pb2.WorkRequest(register=reg))
739
- await self._wait_event_or_session(self._registered)
907
+ self._publish_frame(source, queue_pb2.WorkRequest(register=reg))
908
+ await self._wait_event_or_session(source.registered, source)
740
909
  if self._force_cancel.is_set():
741
910
  return
742
911
  # Initial credit is sent only after WorkerRegistered was observed.
743
912
  if self._controller is None:
744
- await self._send_credit(self._concurrency)
913
+ await self._publish_fixed_capacity()
745
914
  else:
746
915
  self._controller._check_loop()
747
916
  await self._controller._start_member(self._member)
@@ -751,108 +920,111 @@ class Worker:
751
920
  # is bound to the token this stream was registered with: a
752
921
  # later unary renegotiation that rotates the token disarms
753
922
  # the monitor until the next session.
754
- self._idle_session_token = reg.context.session_token
923
+ self._idle_session_token = source.token
755
924
  self._idle_candidate = True
756
925
  self._idle_zero_since = None
757
- await self._wait_event_or_session(self._session_done)
758
- if self._stream_error is not None and not self._close_started.is_set():
759
- raise self._stream_error
926
+ await source.done.wait()
927
+ if source.error is not None and not self._close_started.is_set():
928
+ raise source.error
760
929
  finally:
761
- await self._stop_session()
762
- self._call = None
763
- self._session_done = None
764
- self._registered = None
930
+ await self._stop_session(source)
931
+ if self._source is source:
932
+ self._source = None
765
933
 
766
- async def _wait_event_or_session(self, event: asyncio.Event) -> None:
767
- if self._session_done is None:
768
- return
934
+ async def _wait_event_or_session(
935
+ self, event: asyncio.Event, source: _WorkSource
936
+ ) -> None:
937
+ # Registration is a cold-path Event race. Commands await only their
938
+ # private Future; every source-end path fails those Futures directly.
769
939
  event_task = asyncio.create_task(event.wait())
770
- done_task = asyncio.create_task(self._session_done.wait())
771
- done, pending = await asyncio.wait(
772
- (event_task, done_task), return_when=asyncio.FIRST_COMPLETED
773
- )
774
- for task in pending:
775
- task.cancel()
776
- await asyncio.gather(*pending, return_exceptions=True)
777
- if done_task in done and not event.is_set():
778
- if self._stream_error is not None:
779
- raise self._stream_error
780
- raise UnavailableError("Work stream closed")
940
+ done_task = asyncio.create_task(source.done.wait())
941
+ try:
942
+ done, _ = await asyncio.wait(
943
+ (event_task, done_task), return_when=asyncio.FIRST_COMPLETED
944
+ )
945
+ if done_task in done and not event.is_set():
946
+ raise source.error or UnavailableError("Work stream closed")
947
+ finally:
948
+ for task in (event_task, done_task):
949
+ if not task.done():
950
+ task.cancel()
951
+ await asyncio.gather(event_task, done_task, return_exceptions=True)
952
+
953
+ async def _stop_session(self, source: _WorkSource | None = None) -> None:
954
+ source = self._source if source is None else source
955
+ if source is None:
956
+ return
957
+ # Fail command Futures before any lock/join can block teardown. One
958
+ # shielded owner survives repeated caller cancellation and concurrent
959
+ # close/run cleanup, including cancellation before its first step.
960
+ self._end_source(source)
961
+ if source.cleanup is None:
962
+ task = asyncio.create_task(
963
+ self._stop_source(source),
964
+ name=f"ringo-stop-{self._worker_id}-{source.epoch}",
965
+ )
966
+ source.cleanup = task
967
+ self._capacity_tasks.add(task)
968
+ task.add_done_callback(self._capacity_tasks.discard)
969
+ task.add_done_callback(_consume_task_exception)
970
+ await asyncio.shield(source.cleanup)
781
971
 
782
- async def _stop_session(self) -> None:
783
- session = self._session_done
784
- if session is not None:
785
- session.set()
786
- call, self._call = self._call, None
787
- writer, self._writer_task = self._writer_task, None
788
- reader, self._reader_task = self._reader_task, None
789
- outbox, self._outbox = self._outbox, None
790
- if outbox is not None and writer is not None and not writer.done():
791
- await outbox.put(None)
792
- if call is not None:
793
- try:
794
- call.cancel()
795
- except Exception:
796
- pass
797
- # Unblock publishers before waiting for a reader that may itself be
798
- # awaiting a controller publication boundary. Cancel Python tasks as
799
- # well as the transport; call.cancel alone cannot interrupt that await.
800
- for future in self._pending.values():
801
- if not future.done():
802
- future.set_exception(UnavailableError("Work stream closed"))
803
- self._pending.clear()
972
+ async def _stop_source(self, source: _WorkSource) -> None:
973
+ source.outbox.put_nowait(None)
804
974
  tasks = [
805
- t
806
- for t in (reader, writer)
807
- if t is not None and t is not asyncio.current_task()
975
+ task
976
+ for task in (source.reader, source.writer, source.publisher)
977
+ if task is not None and task is not asyncio.current_task()
808
978
  ]
809
979
  for task in tasks:
810
980
  task.cancel()
811
981
  controller, member = self._controller, self._member
812
- if controller is not None and member is not None:
813
- # Schedule controller fencing before a reader join can be cancelled.
814
- # It owns only stream/probe liveness, never a still-running handler's
815
- # permit, and must survive a second cancellation of run() cleanup.
816
- lost = asyncio.create_task(controller._stream_lost(member))
817
- self._capacity_tasks.add(lost)
818
- lost.add_done_callback(self._capacity_tasks.discard)
819
- lost.add_done_callback(_consume_task_exception)
820
- await asyncio.shield(lost)
982
+ if controller is not None and member is not None and self._source is source:
983
+ # The run loop joins this cleanup before admitting its replacement.
984
+ # Only unused old-stream credit is reclaimed; running handlers keep
985
+ # their permits until their real coroutine settles.
986
+ await controller._stream_lost(member)
821
987
  if tasks:
822
988
  await asyncio.gather(*tasks, return_exceptions=True)
823
989
 
824
- async def _enqueue_frame(self, frame: queue_pb2.WorkRequest) -> None:
825
- outbox = self._outbox
826
- if outbox is None:
827
- raise UnavailableError("Work stream is not connected")
828
- await outbox.put(frame)
990
+ def _publish_frame(self, source: _WorkSource, frame: queue_pb2.WorkRequest) -> None:
991
+ """Synchronous FIFO insertion; callers own validation and correlation."""
992
+ source.outbox.put_nowait(frame)
829
993
 
830
- async def _write_loop(self, call: Any) -> None:
994
+ async def _write_loop(self, source: _WorkSource) -> None:
831
995
  try:
832
- while True:
833
- frame = await self._outbox.get() if self._outbox is not None else None
834
- if frame is None:
996
+ while not source.done.is_set():
997
+ frame = await source.outbox.get()
998
+ if frame is None or source.done.is_set():
835
999
  try:
836
- done_writing = getattr(call, "done_writing", None)
1000
+ done_writing = getattr(source.call, "done_writing", None)
837
1001
  if done_writing is not None:
838
1002
  await done_writing()
839
1003
  except Exception:
840
1004
  pass
841
1005
  return
842
- await call.write(frame)
1006
+ await source.call.write(frame)
843
1007
  except asyncio.CancelledError:
844
1008
  raise
845
1009
  except Exception as exc:
846
- self._stream_error = self._map_transport(exc)
847
- if self._session_done is not None:
848
- self._session_done.set()
1010
+ self._end_source(source, self._map_transport(exc))
1011
+ finally:
1012
+ self._end_source(source)
849
1013
 
850
- async def _read_loop(self, call: Any) -> None:
1014
+ async def _read_loop(self, source: _WorkSource) -> None:
851
1015
  try:
852
- async for response in call:
1016
+ async for response in source.call:
1017
+ if source.done.is_set():
1018
+ return
853
1019
  body = response.WhichOneof("body")
854
1020
  if body == "registered":
1021
+ if self._source is not source:
1022
+ continue
855
1023
  registered = response.registered
1024
+ if registered.ordered_request_ids != source.ordered_request_ids:
1025
+ raise IncompatibleVersionError(
1026
+ "Work registration did not confirm the requested ID mode"
1027
+ )
856
1028
  if registered.HasField("lease_duration"):
857
1029
  self._lease_duration = (
858
1030
  _pb.dur_to_td(registered.lease_duration)
@@ -864,23 +1036,22 @@ class Worker:
864
1036
  and registered.max_concurrency
865
1037
  ):
866
1038
  self._concurrency = registered.max_concurrency
867
- if self._registered is not None:
868
- self._registered.set()
1039
+ source.registered.set()
869
1040
  elif body == "task":
870
- await self._on_assignment(response.task)
1041
+ await self._on_assignment(response.task, source=source)
871
1042
  elif body == "result":
872
1043
  result = response.result
873
- future = self._pending.pop(result.request_id, None)
1044
+ future = source.pending.pop(result.request_id, None)
874
1045
  if future is not None and not future.done():
875
1046
  future.set_result(result)
876
- # Heartbeats are unsolicited liveness frames and need no action.
1047
+ # Unknown/late results and heartbeats keep their existing
1048
+ # non-fatal behavior; they never correlate into another source.
877
1049
  except asyncio.CancelledError:
878
1050
  raise
879
1051
  except Exception as exc:
880
- self._stream_error = self._map_transport(exc)
1052
+ self._end_source(source, self._map_transport(exc))
881
1053
  finally:
882
- if self._session_done is not None:
883
- self._session_done.set()
1054
+ self._end_source(source)
884
1055
 
885
1056
  @staticmethod
886
1057
  def _map_transport(exc: BaseException) -> BaseException:
@@ -890,9 +1061,15 @@ class Worker:
890
1061
  return map_rpc_error(exc)
891
1062
  return UnavailableError(f"Work stream failed: {_safe_error(exc)}")
892
1063
 
893
- async def _on_assignment(self, leased: queue_pb2.LeasedTask) -> None:
1064
+ async def _on_assignment(
1065
+ self, leased: queue_pb2.LeasedTask, *, source: _WorkSource | None = None
1066
+ ) -> None:
1067
+ source = self._source if source is None else source
894
1068
  if (
895
- self._force_cancel.is_set()
1069
+ source is None
1070
+ or self._source is not source
1071
+ or source.done.is_set()
1072
+ or self._force_cancel.is_set()
896
1073
  or not leased.HasField("task")
897
1074
  or not leased.HasField("lease")
898
1075
  ):
@@ -904,6 +1081,8 @@ class Worker:
904
1081
  attempt=leased.lease.attempt,
905
1082
  token=leased.lease.lease_token,
906
1083
  until=until,
1084
+ renew_deadline=asyncio.get_running_loop().time()
1085
+ + (until - datetime.now(timezone.utc)).total_seconds(),
907
1086
  )
908
1087
  context = TaskContext(
909
1088
  self,
@@ -919,8 +1098,8 @@ class Worker:
919
1098
  # a duplicate, and never overwrite either lease's cleanup ownership.
920
1099
  if assignment.identity in self._inflight:
921
1100
  return
922
- assignment.session_epoch = self._session_epoch
923
- assignment.session = self._session_done
1101
+ assignment.session_epoch = source.epoch
1102
+ assignment.session = source.done
924
1103
  if self._controller is not None and self._member is not None:
925
1104
  # Consuming the reserved probe is the linearization point at which
926
1105
  # an assignment becomes a running controller permit.
@@ -929,6 +1108,15 @@ class Worker:
929
1108
  # remains server-owned and will naturally expire if this was a
930
1109
  # transport race.
931
1110
  return
1111
+ if self._source is not source or source.done.is_set():
1112
+ assignment.alive = False
1113
+ context._cancel()
1114
+ await self._release_controller_permit(assignment)
1115
+ return
1116
+ # A complete short task can arrive and ACK between two activity polls.
1117
+ # Acceptance itself restarts quiet time; old/duplicate/rejected frames
1118
+ # returned above must never reset the current source's observation.
1119
+ self._idle_zero_since = None
932
1120
  self._inflight[assignment.identity] = assignment
933
1121
  assignment.renew_task = asyncio.create_task(
934
1122
  self._renew_loop(assignment),
@@ -938,6 +1126,7 @@ class Worker:
938
1126
  # Absolute credit updates can race assignments already in transit.
939
1127
  # Retain and renew such leases, but gate actual handler entry locally.
940
1128
  self._waiting_assignments.append(assignment)
1129
+ self._fixed_capacity_generation += 1
941
1130
  self._start_waiting_assignments()
942
1131
  else:
943
1132
  self._start_handler(assignment)
@@ -955,7 +1144,12 @@ class Worker:
955
1144
  assignment = self._waiting_assignments.popleft()
956
1145
  if not assignment.alive or assignment.session_epoch != self._session_epoch:
957
1146
  continue
958
- if assignment.until <= datetime.now(timezone.utc):
1147
+ expired = (
1148
+ assignment.renew_deadline <= asyncio.get_running_loop().time()
1149
+ if assignment.renew_deadline is not None
1150
+ else assignment.until <= datetime.now(timezone.utc)
1151
+ )
1152
+ if expired:
959
1153
  self._fail_session(
960
1154
  UnavailableError("queued lease reached its local deadline"),
961
1155
  epoch=assignment.session_epoch,
@@ -981,6 +1175,21 @@ class Worker:
981
1175
  def _handler_done(self, assignment: _Assignment, task: asyncio.Task[Any]) -> None:
982
1176
  _consume_task_exception(task)
983
1177
  self._handler_tasks.discard(task)
1178
+ if task.cancelled() and assignment.alive:
1179
+ # Task.cancel before its first step never enters the coroutine's
1180
+ # cancellation handler. Stop that handlerless lease synchronously;
1181
+ # the existing release owner still accounts for actual Task.done.
1182
+ self._fail_session(
1183
+ UnavailableError("task handler cancelled before settlement"),
1184
+ epoch=assignment.session_epoch,
1185
+ session=assignment.session,
1186
+ )
1187
+ self._invalidate_assignment(assignment)
1188
+ if self._inflight.get(assignment.identity) is assignment:
1189
+ self._inflight.pop(assignment.identity)
1190
+ renewal = assignment.renew_task
1191
+ if renewal is not None and not renewal.done():
1192
+ renewal.cancel()
984
1193
  self._schedule_capacity_release(assignment)
985
1194
 
986
1195
  def _schedule_capacity_release(self, assignment: _Assignment) -> None:
@@ -992,6 +1201,11 @@ class Worker:
992
1201
  async def _release_finished_capacity(self, assignment: _Assignment) -> None:
993
1202
  epoch, session = self._session_epoch, self._session_done
994
1203
  try:
1204
+ renewal = assignment.renew_task
1205
+ if not assignment.alive and renewal is not None and not renewal.done():
1206
+ # A done callback can cancel renewal before the handler ever
1207
+ # ran. Join that SDK task without releasing another assignment.
1208
+ await asyncio.gather(renewal, return_exceptions=True)
995
1209
  await self._release_capacity_if_settled(assignment)
996
1210
  except asyncio.CancelledError:
997
1211
  raise
@@ -1007,27 +1221,43 @@ class Worker:
1007
1221
  epoch: int | None = None,
1008
1222
  session: asyncio.Event | None = None,
1009
1223
  ) -> None:
1010
- if epoch is not None and (
1011
- epoch != self._session_epoch or session is not self._session_done
1012
- ):
1224
+ source = self._source
1225
+ if source is None:
1226
+ return
1227
+ if epoch is not None and (epoch != source.epoch or session is not source.done):
1013
1228
  return
1014
- self._stream_error = error
1015
- if self._session_done is not None:
1016
- self._session_done.set()
1017
- if self._call is not None:
1018
- self._call.cancel()
1229
+ self._end_source(source, error)
1019
1230
 
1020
1231
  async def _handle_assignment(self, assignment: _Assignment) -> None:
1021
1232
  try:
1022
1233
  context = assignment.context
1023
1234
  assert context is not None
1024
- if not assignment.alive or self._force_cancel.is_set():
1235
+ if not assignment.alive:
1025
1236
  return
1026
- result = self._handler(assignment.task, context)
1027
- if inspect.isawaitable(result):
1028
- await result
1237
+ if self._force_cancel.is_set() or (
1238
+ assignment.session is not None and assignment.session.is_set()
1239
+ ):
1240
+ # Reader bursts can schedule a handler and observe EOF before
1241
+ # its first step. The immutable origin Event, not a replacement
1242
+ # source or close_started, gates entry into application code.
1243
+ await self._release_assignment(assignment, cancel_context=True)
1244
+ return
1245
+ # One disabled-path check; no monitor task, clock read or RPC for
1246
+ # ordinary workers. The enabled helper enters immediately before
1247
+ # the real handler call and exits before any terminal receipt.
1248
+ watchdog = self._watchdog
1249
+ if watchdog is None:
1250
+ result = self._handler(assignment.task, context)
1251
+ if inspect.isawaitable(result):
1252
+ await result
1253
+ else:
1254
+ await watchdog.invoke(assignment, context)
1029
1255
  if assignment.alive and not self._force_cancel.is_set():
1030
1256
  await self._finish(assignment, "ack")
1257
+ except _WatchdogFenced:
1258
+ # A deadline racing graceful close abandons only its own stopped
1259
+ # invocation; other handlers keep their live drain/ACK/renewal lane.
1260
+ await self._release_assignment(assignment, cancel_context=True)
1031
1261
  except asyncio.CancelledError:
1032
1262
  # Cancellation is not a NACK, but an application-cancelled handler
1033
1263
  # must not retain a live, perpetually renewed lease and capacity.
@@ -1051,6 +1281,38 @@ class Worker:
1051
1281
  # Capacity belongs to the complete handler task, including receipt
1052
1282
  # consumption; its done callback releases the slot exactly once.
1053
1283
 
1284
+ def _watchdog_unresponsive(self, assignment: _Assignment, reason: str) -> None:
1285
+ """Fence before notifying: never NACK while old business still runs."""
1286
+ self._invalidate_assignment(assignment)
1287
+ renewal = assignment.renew_task
1288
+ if renewal is not None:
1289
+ renewal.cancel()
1290
+ # Do not recycle local capacity; the real handler's done callback owns
1291
+ # it even if run() returns while the application swallows cancellation.
1292
+ self._close_started.set()
1293
+ self._publisher_closed = True
1294
+ source = self._source
1295
+ if source is not None:
1296
+ # This is an intentional Worker-wide hard fence, not a stale RPC
1297
+ # result: even a replacement source must stop admitting new work.
1298
+ self._end_source(
1299
+ source, UnavailableError("unresponsive timed-out handler; Worker fenced")
1300
+ )
1301
+ detail = {
1302
+ "worker_id": self._worker_id,
1303
+ "queue": self._queue.name,
1304
+ "task_id": assignment.task.id,
1305
+ "attempt": assignment.attempt,
1306
+ "reason": reason,
1307
+ }
1308
+ if self._runtime is not None:
1309
+ self._runtime.request_restart("task cancellation grace exceeded", detail)
1310
+ else:
1311
+ self._watchdog_error = RuntimeError(
1312
+ "task cancellation grace exceeded: Worker stopped accepting tasks; "
1313
+ "a handler is still running; use process supervision for hard isolation"
1314
+ )
1315
+
1054
1316
  async def _finish(
1055
1317
  self,
1056
1318
  assignment: _Assignment,
@@ -1092,9 +1354,22 @@ class Worker:
1092
1354
  reject=queue_pb2.RejectRequest(lease=assignment.ref, error=text)
1093
1355
  )
1094
1356
  try:
1357
+ # Task.cancel() can be pending after a handler returns/throws without
1358
+ # suspending. Queue.put() on the unbounded outbox is a fast path, so
1359
+ # deliver cancellation before any terminal frame reaches the writer.
1360
+ # Keep this inside our cancellation cleanup, including NACK/reject
1361
+ # callers in handler except blocks. A consumed cancellation remains
1362
+ # a valid normal result; its Task.cancelling() count is not a gate.
1363
+ await asyncio.sleep(0)
1364
+ if assignment.watchdog_expired and (
1365
+ not assignment.alive
1366
+ or self._close_started.is_set()
1367
+ or self._force_cancel.is_set()
1368
+ ):
1369
+ await self._release_assignment(assignment, cancel_context=True)
1370
+ return
1095
1371
  result = await self._send_command(
1096
1372
  request,
1097
- wait=True,
1098
1373
  expected_epoch=assignment.session_epoch,
1099
1374
  expected_session=assignment.session,
1100
1375
  )
@@ -1117,14 +1392,6 @@ class Worker:
1117
1392
  )
1118
1393
  await self._release_assignment(assignment, cancel_context=True)
1119
1394
  return
1120
- if result is None:
1121
- self._fail_session(
1122
- UnavailableError("terminal command returned no receipt"),
1123
- epoch=assignment.session_epoch,
1124
- session=assignment.session,
1125
- )
1126
- await self._release_assignment(assignment, cancel_context=True)
1127
- return
1128
1395
  detail = result.error if result.HasField("error") else None
1129
1396
  if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
1130
1397
  await self._lease_lost(assignment)
@@ -1191,8 +1458,14 @@ class Worker:
1191
1458
  await self._release_capacity_if_settled(assignment)
1192
1459
 
1193
1460
  async def _renew_loop(self, assignment: _Assignment) -> None:
1461
+ loop = asyncio.get_running_loop()
1462
+ if assignment.renew_deadline is None:
1463
+ assignment.renew_deadline = (
1464
+ loop.time()
1465
+ + (assignment.until - datetime.now(timezone.utc)).total_seconds()
1466
+ )
1194
1467
  while assignment.alive and not self._force_cancel.is_set():
1195
- remaining = (assignment.until - datetime.now(timezone.utc)).total_seconds()
1468
+ remaining = assignment.renew_deadline - loop.time()
1196
1469
  delay = max(0.01, remaining - self._lease_duration.total_seconds() / 3)
1197
1470
  try:
1198
1471
  await asyncio.wait_for(self._force_cancel.wait(), timeout=delay)
@@ -1207,10 +1480,11 @@ class Worker:
1207
1480
  extension=_pb.td_to_dur(self._lease_duration),
1208
1481
  )
1209
1482
  )
1483
+ sent_at = loop.time()
1484
+ sent_wall = datetime.now(timezone.utc)
1210
1485
  try:
1211
1486
  result = await self._send_command(
1212
1487
  request,
1213
- wait=True,
1214
1488
  expected_epoch=assignment.session_epoch,
1215
1489
  expected_session=assignment.session,
1216
1490
  )
@@ -1222,14 +1496,6 @@ class Worker:
1222
1496
  )
1223
1497
  await self._lease_lost(assignment)
1224
1498
  return
1225
- if result is None:
1226
- self._fail_session(
1227
- UnavailableError("lease renewal returned no receipt"),
1228
- epoch=assignment.session_epoch,
1229
- session=assignment.session,
1230
- )
1231
- await self._lease_lost(assignment)
1232
- return
1233
1499
  detail = result.error if result.HasField("error") else None
1234
1500
  if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
1235
1501
  await self._lease_lost(assignment)
@@ -1242,10 +1508,22 @@ class Worker:
1242
1508
  )
1243
1509
  await self._lease_lost(assignment)
1244
1510
  return
1245
- # ExtendLeaseResponse is not carried in WorkResult. The protocol
1246
- # defines the new deadline as server_now + extension; this local
1247
- # estimate is the only safe deadline available to the SDK.
1248
- assignment.until = datetime.now(timezone.utc) + self._lease_duration
1511
+ if (
1512
+ not assignment.alive
1513
+ or self._force_cancel.is_set()
1514
+ or assignment.session_epoch != self._session_epoch
1515
+ or assignment.session is not self._session_done
1516
+ or assignment.session is None
1517
+ or assignment.session.is_set()
1518
+ ):
1519
+ return
1520
+ # WorkResult has no renewed deadline. Server processing cannot
1521
+ # precede publication, so send-time + extension is a conservative
1522
+ # lower bound, including queued writes/backpressure. Receipt time
1523
+ # must not grant an extra lease. Monotonic scheduling survives wall
1524
+ # clock changes; keep the private wall-clock field conservative too.
1525
+ assignment.renew_deadline = sent_at + self._lease_duration.total_seconds()
1526
+ assignment.until = sent_wall + self._lease_duration
1249
1527
 
1250
1528
  async def _report_progress(
1251
1529
  self, assignment: _Assignment, current: float, total: float, message: str
@@ -1262,12 +1540,9 @@ class Worker:
1262
1540
  )
1263
1541
  result = await self._send_command(
1264
1542
  request,
1265
- wait=True,
1266
1543
  expected_epoch=assignment.session_epoch,
1267
1544
  expected_session=assignment.session,
1268
1545
  )
1269
- if result is None:
1270
- return
1271
1546
  detail = result.error if result.HasField("error") else None
1272
1547
  if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
1273
1548
  await self._lease_lost(assignment)
@@ -1275,53 +1550,142 @@ class Worker:
1275
1550
  if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
1276
1551
  raise error_from_detail(detail)
1277
1552
 
1278
- async def _send_credit(self, available: int) -> None:
1279
- # ``available`` is intentionally not trusted across retries. A task
1280
- # may arrive between a rejected snapshot and its retry, so every
1281
- # credit-capacity retry samples local in-flight state again.
1282
- await self._publish_fixed_capacity()
1283
-
1284
1553
  async def _publish_fixed_capacity(self) -> None:
1285
1554
  epoch, session = self._session_epoch, self._session_done
1286
- async with self._publisher_lock:
1287
- for attempt in range(20):
1288
- if (
1289
- self._publisher_closed
1290
- or self._close_started.is_set()
1291
- or epoch != self._session_epoch
1292
- or self._session_done is None
1293
- or self._session_done.is_set()
1294
- ):
1295
- return
1296
- available = max(
1297
- 0,
1298
- self._concurrency
1299
- - self._fixed_running
1300
- - len(self._waiting_assignments),
1301
- )
1302
- update = queue_pb2.CreditUpdate(available=available)
1303
- if self._fixed_dynamic_target:
1304
- update.target_concurrency = self._concurrency
1305
- result = await self._send_command(
1306
- queue_pb2.WorkRequest(credit=update),
1307
- wait=True,
1308
- expected_epoch=epoch,
1309
- expected_session=session,
1310
- )
1311
- if (
1312
- result is None
1313
- or not result.HasField("error")
1314
- or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
1315
- ):
1316
- return
1317
- if (
1318
- result.error.metadata.get("reason") != "credit_capacity"
1319
- or attempt == 19
1320
- ):
1321
- error = error_from_detail(result.error)
1322
- self._fail_session(error, epoch=epoch, session=session)
1323
- raise error
1324
- await asyncio.sleep(0.01)
1555
+ if (
1556
+ self._publisher_closed
1557
+ or self._close_started.is_set()
1558
+ or session is None
1559
+ or session.is_set()
1560
+ ):
1561
+ return
1562
+ self._fixed_capacity_generation += 1
1563
+ generation = self._fixed_capacity_generation
1564
+ task = self._fixed_publication_task
1565
+ new_publisher = (
1566
+ task is None
1567
+ or task.done()
1568
+ or self._fixed_publication_source != (epoch, session)
1569
+ )
1570
+ if new_publisher:
1571
+ self._fixed_publication_waiters = []
1572
+ self._fixed_publication_source = (epoch, session)
1573
+ waiters = self._fixed_publication_waiters
1574
+ accepted: asyncio.Future[None] = asyncio.get_running_loop().create_future()
1575
+ accepted.add_done_callback(_consume_task_exception)
1576
+ waiters.append((generation, accepted))
1577
+ if new_publisher:
1578
+ task = asyncio.create_task(
1579
+ self._drain_fixed_capacity(epoch, session, waiters),
1580
+ name=f"ringo-credit-{self._worker_id}-{epoch}",
1581
+ )
1582
+ self._fixed_publication_task = task
1583
+ assert self._source is not None
1584
+ self._source.publisher = task
1585
+ self._capacity_tasks.add(task)
1586
+ task.add_done_callback(
1587
+ lambda completed: self._fixed_publication_done(completed, waiters)
1588
+ )
1589
+ # A resize waits for its generation's accepted snapshot, not for the
1590
+ # publisher to become idle: a continuously busy stream may never idle.
1591
+ # Caller/grace cancellation cannot abandon the shared publication.
1592
+ await asyncio.shield(accepted)
1593
+
1594
+ def _fixed_publication_done(
1595
+ self,
1596
+ task: asyncio.Task[None],
1597
+ waiters: list[tuple[int, asyncio.Future[None]]],
1598
+ ) -> None:
1599
+ self._capacity_tasks.discard(task)
1600
+ error = asyncio.CancelledError() if task.cancelled() else task.exception()
1601
+ # Task.cancel() before the coroutine's first step never enters its
1602
+ # finally. Complete this source's waiters even at that boundary.
1603
+ for _, future in waiters:
1604
+ if not future.done():
1605
+ if error is None:
1606
+ future.set_result(None)
1607
+ else:
1608
+ future.set_exception(error)
1609
+ waiters.clear()
1610
+
1611
+ async def _drain_fixed_capacity(
1612
+ self,
1613
+ epoch: int,
1614
+ session: asyncio.Event,
1615
+ waiters: list[tuple[int, asyncio.Future[None]]],
1616
+ ) -> None:
1617
+ error: BaseException | None = None
1618
+ try:
1619
+ async with self._publisher_lock:
1620
+ while True:
1621
+ for attempt in range(20):
1622
+ if (
1623
+ self._publisher_closed
1624
+ or self._close_started.is_set()
1625
+ or epoch != self._session_epoch
1626
+ or session is not self._session_done
1627
+ or session.is_set()
1628
+ ):
1629
+ return
1630
+ # Absolute credit is consumed by Assign frames. Coalesce
1631
+ # requests, never equal numeric values; an in-transit
1632
+ # assignment dirties the generation even if a completion
1633
+ # leaves occupancy unchanged. Retry re-samples all state.
1634
+ generation = self._fixed_capacity_generation
1635
+ available = max(
1636
+ 0,
1637
+ self._concurrency
1638
+ - self._fixed_running
1639
+ - len(self._waiting_assignments),
1640
+ )
1641
+ update = queue_pb2.CreditUpdate(available=available)
1642
+ if self._fixed_dynamic_target:
1643
+ update.target_concurrency = self._concurrency
1644
+ result = await self._send_command(
1645
+ queue_pb2.WorkRequest(credit=update),
1646
+ expected_epoch=epoch,
1647
+ expected_session=session,
1648
+ )
1649
+ if (
1650
+ not result.HasField("error")
1651
+ or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
1652
+ ):
1653
+ source = self._source
1654
+ if source is not None and source.done is session:
1655
+ source.accepted_credit = available
1656
+ source.credit_generation = generation
1657
+ break
1658
+ if (
1659
+ result.error.metadata.get("reason") != "credit_capacity"
1660
+ or attempt == 19
1661
+ ):
1662
+ raise error_from_detail(result.error)
1663
+ await asyncio.sleep(0.01)
1664
+ remaining = []
1665
+ for requested, future in waiters:
1666
+ if requested <= generation:
1667
+ if not future.done():
1668
+ future.set_result(None)
1669
+ else:
1670
+ remaining.append((requested, future))
1671
+ waiters[:] = remaining
1672
+ if generation == self._fixed_capacity_generation:
1673
+ return
1674
+ except BaseException as exc:
1675
+ error = exc
1676
+ if not isinstance(exc, asyncio.CancelledError):
1677
+ self._fail_session(exc, epoch=epoch, session=session)
1678
+ raise
1679
+ finally:
1680
+ # Session replacement/close releases only this publisher's waiters.
1681
+ # A late old failure never completes or cancels a replacement batch.
1682
+ for _, future in waiters:
1683
+ if not future.done():
1684
+ if error is None:
1685
+ future.set_result(None)
1686
+ else:
1687
+ future.set_exception(error)
1688
+ waiters.clear()
1325
1689
 
1326
1690
  async def _publish_controller_initial(self, capacity: int) -> None:
1327
1691
  async with self._publisher_lock:
@@ -1397,6 +1761,8 @@ class Worker:
1397
1761
  probe_member: _Member | None = None,
1398
1762
  probe_generation: int = 0,
1399
1763
  ) -> bool:
1764
+ controller = self._controller
1765
+ assert controller is not None
1400
1766
  epoch, session = self._session_epoch, self._session_done
1401
1767
  for attempt in range(20):
1402
1768
  if (
@@ -1417,38 +1783,24 @@ class Worker:
1417
1783
  # target and available are one atomic protocol snapshot. A
1418
1784
  # credit_capacity rejection must be retried with a fresh target /
1419
1785
  # in-flight sample, never by replaying the old snapshot.
1420
- current_target = (
1421
- self._controller.capacity
1422
- if self._controller is not None and not explicit_target
1423
- else target
1424
- )
1425
- current_available = available
1426
- if self._controller is None:
1427
- # An assignment may have been sent by the server before the
1428
- # rejected credit result reached this reader. Re-sample the
1429
- # now-materialized local assignments instead of replaying the
1430
- # stale fixed-capacity snapshot.
1431
- current_available = max(
1432
- 0,
1433
- current_target
1434
- - self._fixed_running
1435
- - len(self._waiting_assignments),
1436
- )
1786
+ current_target = controller.capacity if not explicit_target else target
1437
1787
  update = queue_pb2.CreditUpdate(
1438
- available=max(0, current_available),
1788
+ available=max(0, available),
1439
1789
  target_concurrency=current_target,
1440
1790
  )
1441
1791
  result = await self._send_command(
1442
1792
  queue_pb2.WorkRequest(credit=update),
1443
- wait=True,
1444
1793
  expected_epoch=epoch,
1445
1794
  expected_session=session,
1446
1795
  )
1447
1796
  if (
1448
- result is None
1449
- or not result.HasField("error")
1797
+ not result.HasField("error")
1450
1798
  or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
1451
1799
  ):
1800
+ source = self._source
1801
+ if source is not None and source.done is session:
1802
+ source.accepted_credit = available
1803
+ source.credit_generation = probe_generation
1452
1804
  return True
1453
1805
  if (
1454
1806
  result.error.metadata.get("reason") != "credit_capacity"
@@ -1460,72 +1812,57 @@ class Worker:
1460
1812
  await asyncio.sleep(0.01)
1461
1813
  return False
1462
1814
 
1463
- def _allocate_request_id(self) -> int:
1464
- rid = self._next_request_id
1465
- self._next_request_id = 1 if rid == 2**64 - 1 else rid + 1
1466
- return rid
1467
-
1468
1815
  async def _send_command(
1469
1816
  self,
1470
1817
  request: queue_pb2.WorkRequest,
1471
1818
  *,
1472
- wait: bool,
1473
1819
  expected_epoch: int | None = None,
1474
1820
  expected_session: asyncio.Event | None = None,
1475
- ) -> queue_pb2.WorkResult | None:
1821
+ ) -> queue_pb2.WorkResult:
1822
+ source = self._source
1823
+ if source is None or source.done.is_set():
1824
+ raise UnavailableError("Work stream is not connected")
1476
1825
  if expected_epoch is not None and (
1477
- expected_epoch != self._session_epoch
1478
- or expected_session is not self._session_done
1826
+ expected_epoch != source.epoch or expected_session is not source.done
1479
1827
  ):
1480
1828
  raise UnavailableError("command belongs to a replaced Work session")
1481
- if self._session_done is None or self._session_done.is_set():
1482
- raise UnavailableError("Work stream is not connected")
1483
- epoch, session = self._session_epoch, self._session_done
1484
- rid = request.request_id or self._allocate_request_id()
1829
+ if source.next_request_id > _MAX_REQUEST_ID:
1830
+ error = UnavailableError(
1831
+ "Work request ID space exhausted; reconnect required"
1832
+ )
1833
+ self._end_source(source, error)
1834
+ raise error
1835
+ # This is the actual FIFO publication point, shared by every producer
1836
+ # including drain. No await separates source validation, numbering,
1837
+ # pending registration and insertion. Even an immediate receipt can
1838
+ # find its Future, and cancellation never rolls a consumed ID back.
1839
+ rid = source.next_request_id
1840
+ source.next_request_id += 1
1485
1841
  request.request_id = rid
1486
- loop = asyncio.get_running_loop()
1487
- future: asyncio.Future[queue_pb2.WorkResult] = loop.create_future()
1488
- self._pending[rid] = future
1842
+ future: asyncio.Future[queue_pb2.WorkResult] = (
1843
+ asyncio.get_running_loop().create_future()
1844
+ )
1845
+ source.pending[rid] = future
1489
1846
  try:
1490
- await self._enqueue_frame(request)
1491
- if not wait:
1492
- return None
1493
- return await asyncio.wait_for(
1494
- self._wait_command_or_session(future), timeout=self._command_timeout
1495
- )
1847
+ self._publish_frame(source, request)
1848
+ async with asyncio.timeout(self._command_timeout):
1849
+ return await future
1496
1850
  except asyncio.CancelledError:
1497
- if self._pending.get(rid) is future:
1498
- self._pending.pop(rid)
1499
1851
  if future.done() and not future.cancelled():
1500
1852
  future.exception()
1501
1853
  raise
1502
1854
  except Exception as exc:
1503
- if self._pending.get(rid) is future:
1504
- self._pending.pop(rid)
1855
+ # Source ownership, rather than the current Work call, fences a
1856
+ # timeout or publication failure. Unknown mutations are not replayed.
1857
+ self._end_source(source, exc)
1505
1858
  if future.done() and not future.cancelled():
1859
+ # Publication can fail before we ever await this Future. End
1860
+ # the source first, then consume its newly installed exception.
1506
1861
  future.exception()
1507
- self._fail_session(exc, epoch=epoch, session=session)
1508
1862
  raise
1509
-
1510
- async def _wait_command_or_session(
1511
- self, future: asyncio.Future[queue_pb2.WorkResult]
1512
- ) -> queue_pb2.WorkResult:
1513
- session = self._session_done
1514
- if session is None:
1515
- raise UnavailableError("Work stream is not connected")
1516
- session_wait = asyncio.create_task(session.wait())
1517
- try:
1518
- done, _ = await asyncio.wait(
1519
- (future, session_wait), return_when=asyncio.FIRST_COMPLETED
1520
- )
1521
- if future in done:
1522
- return future.result()
1523
- if future.done():
1524
- return future.result()
1525
- raise self._stream_error or UnavailableError("Work stream closed")
1526
1863
  finally:
1527
- session_wait.cancel()
1528
- await asyncio.gather(session_wait, return_exceptions=True)
1864
+ if source.pending.get(rid) is future:
1865
+ source.pending.pop(rid)
1529
1866
 
1530
1867
 
1531
1868
  __all__ = ["TaskContext", "Worker"]