Fix a deprecation in the RabbitMQ backend, and install aio-pika automatically in temp venvs that use it

Author: codeSamuraii

offwork/__main__.py

@@ -54,6 +54,8 @@ async def _run_in_tmp_venv(args: argparse.Namespace) -> None:
     extras: list[str] = []
     if args.backend and args.backend.startswith(("redis://", "rediss://")):
         extras.append("redis")
+    if args.backend and args.backend.startswith(("amqp://", "amqps://")):
+        extras.append("rabbitmq")
 
     async with temp_venv(install_offwork=True, extras=extras) as venv:
         cmd = _build_worker_cmd(str(venv.python), args)
@@ -315,6 +317,8 @@ def _collect_extras(args: argparse.Namespace, script: Path) -> list[str]:
     backend = os.environ.get("OFFWORK_BACKEND", "")
     if backend.startswith(("redis://", "rediss://")) and "redis" not in extras:
         extras.append("redis")
+    if backend.startswith(("amqp://", "amqps://")) and "rabbitmq" not in extras:
+        extras.append("rabbitmq")
     for extra in _detect_offwork_extras(str(script)):
         if extra not in extras:
             extras.append(extra)

offwork/worker/backends/http.py

@@ -22,8 +22,8 @@ class HttpBackend(Backend):
     The base URL can point either at the broker root or at the service root;
     when no path is provided, ``/api/v1/broker`` is assumed.
 
-    Authentication is intentionally simple for the proof-of-concept:
-    include ``?api_key=...`` in the URL and the backend will move it into the
+    Authentication is currently supported via an API key, which can be provided by
+    including ``?api_key=...`` in the URL and the backend will move it into the
     ``X-Offwork-API-Key`` request header.
     """
 

offwork/worker/backends/rabbitmq.py

@@ -67,13 +67,22 @@ def __init__(
         *,
         task_queue: str | None = None,
         result_ttl: int | None = None,
+        queue_namespace: str | None = None,
     ) -> None:
         self._url = url
         self._task_queue_name = task_queue or self.TASK_QUEUE
         self._result_ttl = result_ttl or self.DEFAULT_RESULT_TTL
         self._connection: Any = None
         self._channel: Any = None
         self._lock = asyncio.Lock()
+        ns = f"{queue_namespace}." if queue_namespace else ""
+        self._result_prefix = f"{ns}{self.RESULT_PREFIX}"
+        self._heartbeat_prefix = f"{ns}{self.HEARTBEAT_PREFIX}"
+        self._cancel_prefix = f"{ns}{self.CANCEL_PREFIX}"
+        self._progress_prefix = f"{ns}{self.PROGRESS_PREFIX}"
+        self._schedule_prefix = f"{ns}{self.SCHEDULE_PREFIX}"
+        self._throttle_prefix = f"{ns}{self.THROTTLE_PREFIX}"
+        self._notify_exchange = f"{ns}{self.NOTIFY_EXCHANGE}"
 
     # -- connection management --------------------------------------------------
 
@@ -146,7 +155,7 @@ async def _declare_queue_robust(
         redeclares it with the new arguments.
         """
         try:
-            return await channel.declare_queue(name, arguments=arguments)
+            return await channel.declare_queue(name, durable=True, arguments=arguments)
         except Exception:
             # Channel is closed by RabbitMQ after PRECONDITION_FAILED.
             # Reopen a fresh channel, purge the stale queue, and retry.
@@ -157,7 +166,7 @@ async def _declare_queue_robust(
             except Exception:
                 self._channel = None
                 channel = await self._ensure_channel()
-            return await channel.declare_queue(name, arguments=arguments)
+            return await channel.declare_queue(name, durable=True, arguments=arguments)
 
     async def _kv_put(
         self, prefix: str, task_id: str, value: str, ttl_s: int,
@@ -229,8 +238,8 @@ async def listen(self) -> AsyncIterator[str]:
     async def send_result(self, task_id: str, result_json: str) -> None:
         async with self._lock:
             channel = await self._ensure_channel()
-            name = f"{self.RESULT_PREFIX}{task_id}"
-            await channel.declare_queue(name, arguments=self._result_args())
+            name = f"{self._result_prefix}{task_id}"
+            await channel.declare_queue(name, durable=True, arguments=self._result_args())
             await channel.default_exchange.publish(
                 aio_pika.Message(result_json.encode()),
                 routing_key=name,
@@ -239,9 +248,9 @@ async def send_result(self, task_id: str, result_json: str) -> None:
     async def get_result(self, task_id: str, timeout: float | None = None) -> str:
         channel = await self._new_channel()
         try:
-            name = f"{self.RESULT_PREFIX}{task_id}"
+            name = f"{self._result_prefix}{task_id}"
             queue = await channel.declare_queue(
-                name, arguments=self._result_args(),
+                name, durable=True, arguments=self._result_args(),
             )
             future: asyncio.Future[str] = asyncio.get_running_loop().create_future()
 
@@ -270,9 +279,9 @@ async def _on_message(msg: Any) -> None:
     async def try_get_result(self, task_id: str) -> str | None:
         async with self._lock:
             channel = await self._ensure_channel()
-            name = f"{self.RESULT_PREFIX}{task_id}"
+            name = f"{self._result_prefix}{task_id}"
             queue = await channel.declare_queue(
-                name, arguments=self._result_args(),
+                name, durable=True, arguments=self._result_args(),
             )
             msg = await queue.get(fail=False)
             if msg is None:
@@ -285,55 +294,51 @@ async def try_get_result(self, task_id: str) -> str | None:
 
     async def send_heartbeat(self, task_id: str) -> None:
         await self._kv_put(
-            self.HEARTBEAT_PREFIX, task_id,
+            self._heartbeat_prefix, task_id,
             str(time.time()), self.HEARTBEAT_TTL,
         )
 
     async def get_heartbeat(self, task_id: str) -> float | None:
-        # Consume (ack) the heartbeat rather than peeking. This avoids a
-        # RabbitMQ race where nack+requeue bypasses x-max-length=1 and causes
-        # the stale heartbeat to be returned on every subsequent poll, making
-        # stall detection fire spuriously.
         raw = await self._kv_get(
-            self.HEARTBEAT_PREFIX, task_id, self.HEARTBEAT_TTL, peek=False,
+            self._heartbeat_prefix, task_id, self.HEARTBEAT_TTL, peek=True,
         )
         return float(raw) if raw is not None else None
 
     # -- Cancellation ----------------------------------------------------------
 
     async def cancel_task(self, task_id: str) -> None:
         await self._kv_put(
-            self.CANCEL_PREFIX, task_id, "1", self.CANCEL_TTL,
+            self._cancel_prefix, task_id, "1", self.CANCEL_TTL,
         )
 
     async def is_cancelled(self, task_id: str) -> bool:
         raw = await self._kv_get(
-            self.CANCEL_PREFIX, task_id, self.CANCEL_TTL, peek=True,
+            self._cancel_prefix, task_id, self.CANCEL_TTL, peek=True,
         )
         return raw is not None
 
     # -- Progress --------------------------------------------------------------
 
     async def send_progress(self, task_id: str, progress_json: str) -> None:
         await self._kv_put(
-            self.PROGRESS_PREFIX, task_id, progress_json, self.PROGRESS_TTL,
+            self._progress_prefix, task_id, progress_json, self.PROGRESS_TTL,
         )
 
     async def get_progress(self, task_id: str) -> str | None:
         return await self._kv_get(
-            self.PROGRESS_PREFIX, task_id, self.PROGRESS_TTL, peek=True,
+            self._progress_prefix, task_id, self.PROGRESS_TTL, peek=True,
         )
 
     # -- Schedule cancellation -------------------------------------------------
 
     async def cancel_schedule(self, schedule_id: str) -> None:
         await self._kv_put(
-            self.SCHEDULE_PREFIX, schedule_id, "1", self.SCHEDULE_TTL,
+            self._schedule_prefix, schedule_id, "1", self.SCHEDULE_TTL,
         )
 
     async def is_schedule_cancelled(self, schedule_id: str) -> bool:
         raw = await self._kv_get(
-            self.SCHEDULE_PREFIX, schedule_id, self.SCHEDULE_TTL, peek=True,
+            self._schedule_prefix, schedule_id, self.SCHEDULE_TTL, peek=True,
         )
         return raw is not None
 
@@ -352,7 +357,7 @@ async def check_throttle(self, function_name: str) -> bool:
         # message body.  Function names can contain characters rejected by
         # the AMQP queue-name grammar (e.g. ``<locals>``), so we hash them.
         raw = await self._kv_get(
-            self.THROTTLE_PREFIX, self._safe_suffix(function_name),
+            self._throttle_prefix, self._safe_suffix(function_name),
             self.THROTTLE_QUEUE_TTL, peek=True,
         )
         if raw is None:
@@ -364,7 +369,7 @@ async def record_throttle(
     ) -> None:
         expiry = time.time() + throttle_seconds
         await self._kv_put(
-            self.THROTTLE_PREFIX, self._safe_suffix(function_name),
+            self._throttle_prefix, self._safe_suffix(function_name),
             str(expiry), self.THROTTLE_QUEUE_TTL,
         )
 
@@ -374,7 +379,7 @@ async def notify_result(self, task_id: str) -> None:
         async with self._lock:
             channel = await self._ensure_channel()
             exchange = await channel.declare_exchange(
-                self.NOTIFY_EXCHANGE, aio_pika.ExchangeType.FANOUT,
+                self._notify_exchange, aio_pika.ExchangeType.FANOUT,
             )
             await exchange.publish(
                 aio_pika.Message(task_id.encode()),
@@ -385,7 +390,7 @@ async def subscribe_results(self) -> AsyncIterator[str]:
         channel = await self._new_channel()
         try:
             exchange = await channel.declare_exchange(
-                self.NOTIFY_EXCHANGE, aio_pika.ExchangeType.FANOUT,
+                self._notify_exchange, aio_pika.ExchangeType.FANOUT,
             )
             queue = await channel.declare_queue(exclusive=True)
             await queue.bind(exchange)

tests/test_rabbitmq_backend.py

@@ -56,6 +56,7 @@ async def backend() -> AsyncIterator["RabbitMQBackend"]:
     b = RabbitMQBackend(
         RABBITMQ_URL,
         task_queue=f"offwork.test.tasks.{suffix}",
+        queue_namespace=f"test.{suffix}",
     )
     yield b
     await b.close()