ringo-task-queue 0.1.0.dev0__py3-none-macosx_11_0_arm64.whl

ringo_task_queue/worker.py

@@ -0,0 +1,785 @@
+"""Asyncio worker implementation for the credit-controlled Work stream.
+
+The worker deliberately owns one reader and one writer for every Work call.
+Handlers only receive :class:`TaskContext`; lease tokens stay in this module.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import math
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Any, Awaitable, Callable, Iterable
+
+import grpc
+
+from . import _pb
+from ._proto.ringo.v1 import queue_pb2
+from .errors import (
+    LeaseLostError,
+    PermanentTaskError,
+    RingoError,
+    RetryTaskError,
+    UnavailableError,
+    error_from_detail,
+    map_rpc_error,
+)
+from .models import Task
+
+
+_MAX_ERROR_BYTES = 256 * 1024
+
+
+def _consume_task_exception(task: asyncio.Task[Any]) -> None:
+    if not task.cancelled():
+        task.exception()
+
+
+def _safe_error(value: object) -> str:
+    """Return a valid UTF-8 error string no larger than 256 KiB."""
+    try:
+        text = str(value)
+    except Exception:
+        text = type(value).__name__
+    encoded = text.encode("utf-8", errors="replace")
+    if len(encoded) <= _MAX_ERROR_BYTES:
+        return encoded.decode("utf-8", errors="replace")
+    # Decode with replacement after truncating bytes; this is UTF-8 safe and
+    # never emits a broken surrogate or an over-sized protobuf string.
+    return encoded[:_MAX_ERROR_BYTES].decode("utf-8", errors="ignore")
+
+
+class TaskContext:
+    """Context supplied to a worker handler.
+
+    ``cancelled`` becomes true when the local lease is lost, the stream drops,
+    or the worker is closed.  ``report_progress`` is correlated on the Work
+    stream and never exposes the lease token to application code.
+    """
+
+    __slots__ = (
+        "_worker",
+        "_task",
+        "_attempt",
+        "_max_attempts",
+        "_worker_id",
+        "_assignment",
+        "_cancel_event",
+    )
+
+    def __init__(
+        self,
+        worker: "Worker",
+        task: Task,
+        attempt: int,
+        max_attempts: int,
+        worker_id: str,
+        assignment: "_Assignment",
+    ) -> None:
+        self._worker = worker
+        self._task = task
+        self._attempt = attempt
+        self._max_attempts = max_attempts
+        self._worker_id = worker_id
+        self._assignment = assignment
+        self._cancel_event = asyncio.Event()
+
+    @property
+    def task(self) -> Task:
+        return self._task
+
+    @property
+    def attempt(self) -> int:
+        return self._attempt
+
+    @property
+    def max_attempts(self) -> int:
+        return self._max_attempts
+
+    @property
+    def worker_id(self) -> str:
+        return self._worker_id
+
+    @property
+    def cancelled(self) -> bool:
+        return self._cancel_event.is_set()
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Alias useful to code shared with SDKs using a verb-like name."""
+        return self.cancelled
+
+    async def wait_cancelled(self) -> None:
+        await self._cancel_event.wait()
+
+    async def report_progress(
+        self, current: float, total: float, message: str = ""
+    ) -> None:
+        """Report progress for this attempt and await its correlated result."""
+        if self.cancelled:
+            raise LeaseLostError("task lease is no longer valid")
+        await self._worker._report_progress(self._assignment, current, total, message)
+
+    def _cancel(self) -> None:
+        self._cancel_event.set()
+
+
+@dataclass
+class _Assignment:
+    task: Task
+    attempt: int
+    token: str
+    until: datetime
+    context: TaskContext | None = None
+    handler_task: asyncio.Task[Any] | None = None
+    renew_task: asyncio.Task[Any] | None = None
+    alive: bool = True
+    completion_started: bool = False
+
+    @property
+    def ref(self) -> queue_pb2.LeaseRef:
+        return queue_pb2.LeaseRef(
+            task_id=self.task.id,
+            attempt=self.attempt,
+            lease_token=self.token,
+        )
+
+
+class Worker:
+    """A credit-controlled asynchronous task worker.
+
+    ``run`` is intentionally a blocking coroutine: it remains active until
+    ``close`` is called.  The client registers workers and drains them before
+    closing its channel.
+    """
+
+    def __init__(
+        self,
+        queue: Any,
+        handler: Callable[[Task, TaskContext], Awaitable[Any]],
+        *,
+        task_types: Iterable[str],
+        concurrency: int,
+        lease_duration: timedelta,
+        grace_period: timedelta,
+    ) -> None:
+        if concurrency < 1:
+            raise ValueError("concurrency must be >= 1")
+        self._queue = queue
+        self._client = queue._client
+        self._handler = handler
+        self._task_types = list(task_types)
+        self._concurrency = concurrency
+        self._lease_duration = lease_duration
+        self._grace_period = max(0.0, grace_period.total_seconds())
+        self._worker_id = self._new_worker_id()
+
+        # close_started means no new stream may be opened, while force_cancel
+        # is raised only after the graceful drain window expires.  Existing
+        # handlers and renewals remain fully live between those two points.
+        self._close_started = asyncio.Event()
+        self._force_cancel = asyncio.Event()
+        self._run_started = False
+        self._run_finished = asyncio.Event()
+        self._close_finished = asyncio.Event()
+        self._command_timeout = 30.0
+        self._close_lock = asyncio.Lock()
+        self._close_task: asyncio.Task[None] | None = None
+        self._session_done: asyncio.Event | None = None
+        self._call: Any = None
+        self._outbox: asyncio.Queue[Any] | None = None
+        self._reader_task: asyncio.Task[Any] | None = None
+        self._writer_task: asyncio.Task[Any] | None = None
+        self._registered: asyncio.Event | None = None
+        self._stream_error: BaseException | None = None
+        self._next_request_id = 1
+        self._pending: dict[int, asyncio.Future[queue_pb2.WorkResult]] = {}
+        self._inflight: dict[str, _Assignment] = {}
+
+    @staticmethod
+    def _new_worker_id() -> str:
+        # Importing here avoids coupling worker construction to client internals.
+        from .client import new_worker_id
+
+        return new_worker_id()
+
+    @property
+    def worker_id(self) -> str:
+        return self._worker_id
+
+    @property
+    def concurrency(self) -> int:
+        return self._concurrency
+
+    async def run(self) -> None:
+        """Run until :meth:`close` is called, reconnecting boundedly."""
+        if self._run_started:
+            raise RuntimeError("worker.run() may only be called once")
+        self._run_started = True
+        failures = 0
+        try:
+            while not self._close_started.is_set():
+                try:
+                    await self._run_session()
+                except asyncio.CancelledError:
+                    if not self._close_started.is_set():
+                        raise
+                    break
+                except Exception as exc:
+                    self._stream_error = exc
+                    failures += 1
+                    if self._close_started.is_set():
+                        break
+                if self._close_started.is_set():
+                    break
+                # A stream loss invalidates local completion attempts, but does
+                # not send terminal commands for leases that may still exist.
+                await self._cancel_assignments()
+                conn = self._client._require_conn()
+                before_token = conn._session_token
+                await conn._recover_best_effort()
+                if conn._session_token == before_token:
+                    failures += 1
+                else:
+                    failures = 0
+                if failures >= 3:
+                    raise self._stream_error or UnavailableError(
+                        "bounded Work stream recovery exhausted"
+                    )
+                if not self._close_started.is_set():
+                    await asyncio.sleep(0)
+        finally:
+            await self._cancel_assignments()
+            self._run_finished.set()
+
+    async def close(self) -> None:
+        """Drain and stop once; caller cancellation cannot cancel cleanup."""
+        async with self._close_lock:
+            if self._close_task is None:
+                self._close_started.set()
+                self._close_task = asyncio.create_task(
+                    self._close_impl(), name=f"ringo-close-{self._worker_id}"
+                )
+                self._close_task.add_done_callback(_consume_task_exception)
+            close_task = self._close_task
+        await asyncio.shield(close_task)
+
+    async def _close_impl(self) -> None:
+        try:
+            session = self._session_done
+            # Drain is best effort. A disconnected stream cannot safely report
+            # any terminal result, and leases are deliberately left to expire.
+            if session is not None and not session.is_set():
+                try:
+                    await asyncio.wait_for(
+                        self._send_command(
+                            queue_pb2.WorkRequest(
+                                request_id=self._allocate_request_id(),
+                                drain=queue_pb2.WorkerDrain(begin=True),
+                            ),
+                            wait=True,
+                        ),
+                        timeout=max(
+                            0.1,
+                            min(self._command_timeout, self._grace_period or 1.0),
+                        ),
+                    )
+                except Exception:
+                    pass
+            if self._grace_period:
+                try:
+                    await asyncio.wait_for(
+                        self._wait_handlers(), timeout=self._grace_period
+                    )
+                except asyncio.TimeoutError:
+                    pass
+        finally:
+            # From this point handlers and renewals are cancelled; no terminal
+            # frame is generated by their cancellation paths.
+            self._force_cancel.set()
+            try:
+                await self._cancel_assignments()
+                await self._stop_session()
+                self._client._unregister_worker(self)
+                if self._run_started:
+                    await self._run_finished.wait()
+            finally:
+                self._close_finished.set()
+
+    async def _wait_handlers(self) -> None:
+        tasks = [a.handler_task for a in self._inflight.values() if a.handler_task]
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+    async def _cancel_assignments(self) -> None:
+        assignments = list(self._inflight.values())
+        self._inflight.clear()
+        for assignment in assignments:
+            self._invalidate_assignment(assignment)
+        tasks: list[asyncio.Task[Any]] = []
+        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)
+
+    def _invalidate_assignment(self, assignment: _Assignment) -> None:
+        if not assignment.alive:
+            return
+        assignment.alive = False
+        if assignment.context is not None:
+            assignment.context._cancel()
+
+    async def _run_session(self) -> None:
+        conn = self._client._require_conn()
+        reg = queue_pb2.WorkerRegistration(
+            queue=self._queue.name,
+            worker_id=self._worker_id,
+            task_types=self._task_types,
+            max_concurrency=self._concurrency,
+        )
+        if self._lease_duration > timedelta(0):
+            reg.lease_duration.CopyFrom(_pb.td_to_dur(self._lease_duration))
+        async with conn._lock:
+            stub = conn._require_stub()
+            reg.context.CopyFrom(conn.request_context())
+            call = stub.Work(metadata=conn._metadata())
+        self._call = call
+        self._session_done = asyncio.Event()
+        self._registered = asyncio.Event()
+        self._stream_error = None
+        self._outbox = asyncio.Queue()
+        self._pending.clear()
+        self._writer_task = asyncio.create_task(
+            self._write_loop(call), name="ringo-worker-writer"
+        )
+        self._reader_task = asyncio.create_task(
+            self._read_loop(call), name="ringo-worker-reader"
+        )
+        try:
+            await self._enqueue_frame(queue_pb2.WorkRequest(register=reg))
+            await self._wait_event_or_session(self._registered)
+            if self._force_cancel.is_set():
+                return
+            # Initial credit is sent only after WorkerRegistered was observed.
+            await self._send_credit(self._concurrency)
+            await self._wait_event_or_session(self._session_done)
+            if self._stream_error is not None and not self._close_started.is_set():
+                raise self._stream_error
+        finally:
+            await self._stop_session()
+            self._call = None
+            self._session_done = None
+            self._registered = None
+
+    async def _wait_event_or_session(self, event: asyncio.Event) -> None:
+        if self._session_done is None:
+            return
+        event_task = asyncio.create_task(event.wait())
+        done_task = asyncio.create_task(self._session_done.wait())
+        done, pending = await asyncio.wait(
+            (event_task, done_task), return_when=asyncio.FIRST_COMPLETED
+        )
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        if done_task in done and not event.is_set():
+            if self._stream_error is not None:
+                raise self._stream_error
+            raise UnavailableError("Work stream closed")
+
+    async def _stop_session(self) -> None:
+        session = self._session_done
+        if session is not None:
+            session.set()
+        call, self._call = self._call, None
+        writer, self._writer_task = self._writer_task, None
+        reader, self._reader_task = self._reader_task, None
+        outbox, self._outbox = self._outbox, None
+        if outbox is not None and writer is not None and not writer.done():
+            await outbox.put(None)
+        if call is not None:
+            try:
+                call.cancel()
+            except Exception:
+                pass
+        tasks = [t for t in (reader, writer) if t is not None]
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        for future in self._pending.values():
+            if not future.done():
+                future.set_exception(UnavailableError("Work stream closed"))
+        self._pending.clear()
+
+    async def _enqueue_frame(self, frame: queue_pb2.WorkRequest) -> None:
+        outbox = self._outbox
+        if outbox is None:
+            raise UnavailableError("Work stream is not connected")
+        await outbox.put(frame)
+
+    async def _write_loop(self, call: Any) -> None:
+        try:
+            while True:
+                frame = await self._outbox.get() if self._outbox is not None else None
+                if frame is None:
+                    try:
+                        done_writing = getattr(call, "done_writing", None)
+                        if done_writing is not None:
+                            await done_writing()
+                    except Exception:
+                        pass
+                    return
+                await call.write(frame)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._stream_error = self._map_transport(exc)
+            if self._session_done is not None:
+                self._session_done.set()
+
+    async def _read_loop(self, call: Any) -> None:
+        try:
+            async for response in call:
+                body = response.WhichOneof("body")
+                if body == "registered":
+                    registered = response.registered
+                    if registered.HasField("lease_duration"):
+                        self._lease_duration = (
+                            _pb.dur_to_td(registered.lease_duration)
+                            or self._lease_duration
+                        )
+                    if registered.max_concurrency:
+                        self._concurrency = registered.max_concurrency
+                    if self._registered is not None:
+                        self._registered.set()
+                elif body == "task":
+                    await self._on_assignment(response.task)
+                elif body == "result":
+                    result = response.result
+                    future = self._pending.pop(result.request_id, None)
+                    if future is not None and not future.done():
+                        future.set_result(result)
+                # Heartbeats are unsolicited liveness frames and need no action.
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._stream_error = self._map_transport(exc)
+        finally:
+            if self._session_done is not None:
+                self._session_done.set()
+
+    @staticmethod
+    def _map_transport(exc: BaseException) -> BaseException:
+        if isinstance(exc, RingoError):
+            return exc
+        if isinstance(exc, grpc.aio.AioRpcError):
+            return map_rpc_error(exc)
+        return UnavailableError(f"Work stream failed: {_safe_error(exc)}")
+
+    async def _on_assignment(self, leased: queue_pb2.LeasedTask) -> None:
+        if (
+            self._force_cancel.is_set()
+            or not leased.HasField("task")
+            or not leased.HasField("lease")
+        ):
+            return
+        task = Task.from_proto(leased.task)
+        until = _pb.ts_to_dt(leased.lease.until)
+        assignment = _Assignment(
+            task=task,
+            attempt=leased.lease.attempt,
+            token=leased.lease.lease_token,
+            until=until,
+        )
+        context = TaskContext(
+            self,
+            task,
+            assignment.attempt,
+            task.max_attempts,
+            self._worker_id,
+            assignment,
+        )
+        assignment.context = context
+        # Task IDs are unique among valid in-flight leases. Avoid replacing an
+        # existing assignment if a malformed peer sends a duplicate.
+        if task.id in self._inflight:
+            return
+        self._inflight[task.id] = assignment
+        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}"
+        )
+        assignment.handler_task.add_done_callback(_consume_task_exception)
+
+    async def _handle_assignment(self, assignment: _Assignment) -> None:
+        try:
+            context = assignment.context
+            assert context is not None
+            result = self._handler(assignment.task, context)
+            if inspect.isawaitable(result):
+                await result
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "ack")
+        except asyncio.CancelledError:
+            # Cancellation is intentionally not a task outcome. In particular,
+            # do not emit NACK while draining or after a lease loss.
+            raise
+        except RetryTaskError as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "nack", exc)
+        except PermanentTaskError as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "reject", exc)
+        except Exception as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "nack", exc)
+
+    async def _finish(
+        self,
+        assignment: _Assignment,
+        operation: str,
+        error: BaseException | None = None,
+    ) -> None:
+        if not assignment.alive or assignment.completion_started:
+            return
+        assignment.completion_started = True
+        request: queue_pb2.WorkRequest
+        text = _safe_error(error) if error is not None else ""
+        if operation == "ack":
+            request = queue_pb2.WorkRequest(
+                ack=queue_pb2.AckRequest(lease=assignment.ref)
+            )
+        elif operation == "nack":
+            nack = queue_pb2.NackRequest(lease=assignment.ref, error=text)
+            if isinstance(error, RetryTaskError) and error.delay is not None:
+                if (
+                    not isinstance(error.delay, (int, float))
+                    or not math.isfinite(error.delay)
+                    or error.delay < 0
+                ):
+                    # Treat an invalid override as an ordinary failure rather
+                    # than allowing a negative Duration onto the wire.
+                    nack.error = _safe_error(error)
+                else:
+                    try:
+                        nack.retry_delay_override.CopyFrom(
+                            _pb.td_to_dur(timedelta(seconds=error.delay))
+                        )
+                    except Exception:
+                        # An unrepresentable override is treated as a normal
+                        # NACK; never strand the assignment on serialization.
+                        nack.error = _safe_error(error)
+            request = queue_pb2.WorkRequest(nack=nack)
+        else:
+            request = queue_pb2.WorkRequest(
+                reject=queue_pb2.RejectRequest(lease=assignment.ref, error=text)
+            )
+        try:
+            result = await self._send_command(request, wait=True)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            return
+        if result is None:
+            return
+        detail = result.error if result.HasField("error") else None
+        if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+            await self._lease_lost(assignment)
+            return
+        if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+            # A non-lease terminal failure is not terminal according to the
+            # 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))
+
+    async def _release_assignment(
+        self, assignment: _Assignment, *, cancel_context: bool = False
+    ) -> None:
+        if not assignment.alive:
+            return
+        assignment.alive = False
+        self._inflight.pop(assignment.task.id, None)
+        if cancel_context and assignment.context is not None:
+            assignment.context._cancel()
+        if (
+            assignment.renew_task is not None
+            and assignment.renew_task is not asyncio.current_task()
+        ):
+            assignment.renew_task.cancel()
+            await asyncio.gather(assignment.renew_task, return_exceptions=True)
+
+    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
+        if (
+            current is not None
+            and current is not asyncio.current_task()
+            and not current.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))
+
+    async def _renew_loop(self, assignment: _Assignment) -> None:
+        while assignment.alive and not self._force_cancel.is_set():
+            remaining = (assignment.until - datetime.now(timezone.utc)).total_seconds()
+            delay = max(0.01, remaining - self._lease_duration.total_seconds() / 3)
+            try:
+                await asyncio.wait_for(self._force_cancel.wait(), timeout=delay)
+                return
+            except asyncio.TimeoutError:
+                pass
+            if not assignment.alive or self._force_cancel.is_set():
+                return
+            request = queue_pb2.WorkRequest(
+                extend_lease=queue_pb2.ExtendLeaseRequest(
+                    lease=assignment.ref,
+                    extension=_pb.td_to_dur(self._lease_duration),
+                )
+            )
+            try:
+                result = await self._send_command(request, wait=True)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                return
+            if result is None:
+                return
+            detail = result.error if result.HasField("error") else None
+            if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+                await self._lease_lost(assignment)
+                return
+            if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+                return
+            # ExtendLeaseResponse is not carried in WorkResult. The protocol
+            # defines the new deadline as server_now + extension; this local
+            # estimate is the only safe deadline available to the SDK.
+            assignment.until = datetime.now(timezone.utc) + self._lease_duration
+
+    async def _report_progress(
+        self, assignment: _Assignment, current: float, total: float, message: str
+    ) -> None:
+        if not assignment.alive:
+            raise LeaseLostError("task lease is no longer valid")
+        request = queue_pb2.WorkRequest(
+            progress=queue_pb2.ReportProgressRequest(
+                lease=assignment.ref,
+                current=current,
+                total=total,
+                message=_safe_error(message),
+            )
+        )
+        result = await self._send_command(request, wait=True)
+        if result is None:
+            return
+        detail = result.error if result.HasField("error") else None
+        if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+            await self._lease_lost(assignment)
+            raise LeaseLostError(detail.message)
+        if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+            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.
+        for attempt in range(20):
+            request = queue_pb2.WorkRequest(
+                credit=queue_pb2.CreditUpdate(available=available)
+            )
+            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
+            if (
+                result.error.metadata.get("reason") != "credit_capacity"
+                or attempt == 19
+            ):
+                raise error_from_detail(result.error)
+            await asyncio.sleep(0.01)
+
+    def _allocate_request_id(self) -> int:
+        rid = self._next_request_id
+        self._next_request_id += 1
+        if self._next_request_id == 0:
+            self._next_request_id = 1
+        return rid
+
+    async def _send_command(
+        self, request: queue_pb2.WorkRequest, *, wait: bool
+    ) -> queue_pb2.WorkResult | None:
+        if self._session_done is None or self._session_done.is_set():
+            raise UnavailableError("Work stream is not connected")
+        rid = request.request_id or self._allocate_request_id()
+        request.request_id = rid
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[queue_pb2.WorkResult] = loop.create_future()
+        self._pending[rid] = future
+        try:
+            await self._enqueue_frame(request)
+            if not wait:
+                return None
+            return await asyncio.wait_for(
+                self._wait_command_or_session(future), timeout=self._command_timeout
+            )
+        except asyncio.CancelledError:
+            self._pending.pop(rid, None)
+            if future.done() and not future.cancelled():
+                future.exception()
+            raise
+        except Exception:
+            self._pending.pop(rid, None)
+            if future.done() and not future.cancelled():
+                future.exception()
+            raise
+
+    async def _wait_command_or_session(
+        self, future: asyncio.Future[queue_pb2.WorkResult]
+    ) -> queue_pb2.WorkResult:
+        session = self._session_done
+        if session is None:
+            raise UnavailableError("Work stream is not connected")
+        session_wait = asyncio.create_task(session.wait())
+        try:
+            done, _ = await asyncio.wait(
+                (future, session_wait), return_when=asyncio.FIRST_COMPLETED
+            )
+            if future in done:
+                return future.result()
+            if future.done():
+                return future.result()
+            raise self._stream_error or UnavailableError("Work stream closed")
+        finally:
+            session_wait.cancel()
+            await asyncio.gather(session_wait, return_exceptions=True)
+
+
+__all__ = ["TaskContext", "Worker"]