redis-message-queue 8.2.5__tar.gz → 8.2.7__tar.gz

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: redis-message-queue
-Version: 8.2.5
+Version: 8.2.7
 Summary: Python message queuing with Redis and message deduplication
 Project-URL: Homepage, https://github.com/Elijas/redis-message-queue
 Project-URL: Repository, https://github.com/Elijas/redis-message-queue
@@ -34,7 +34,7 @@ Description-Content-Type: text/markdown
 **Lightweight Python message queuing with Redis and built-in publish-side deduplication.** Deduplicate publishes within a TTL window, with optional crash recovery — across any number of producers and consumers.
 
 ```bash
-pip install "redis-message-queue>=8.2.5,<9.0.0"
+pip install "redis-message-queue>=8.2.7,<9.0.0"
 ```
 
 Requires Redis server >= 6.2.
@@ -46,6 +46,7 @@ Redis must be running locally first: use `redis-server` or
 
 ```python
 import json
+from uuid import uuid4
 from redis import Redis
 from redis_message_queue import RedisMessageQueue
 
@@ -56,7 +57,8 @@ queue = RedisMessageQueue(
     deduplication=True,
     get_deduplication_key=lambda msg: msg["id"],
 )
-queue.publish({"id": "msg-1", "text": "hello"})
+message = {"id": f"msg-{uuid4().hex}", "text": "hello"}
+queue.publish(message)
 with queue.process_message() as message:
     if message is not None:
         payload = json.loads(message)
@@ -80,6 +82,7 @@ with queue.process_message() as message:
 ```python
 import asyncio
 import json
+from uuid import uuid4
 from redis.asyncio import Redis
 from redis_message_queue.asyncio import RedisMessageQueue
 
@@ -91,10 +94,12 @@ async def main():
         deduplication=True,
         get_deduplication_key=lambda msg: msg["id"],
     )
-    await queue.publish({"id": "msg-1", "text": "hello"})
+    message = {"id": f"msg-{uuid4().hex}", "text": "hello"}
+    await queue.publish(message)
     async with queue.process_message() as message:
-        payload = json.loads(message)
-        print(f"got {payload['text']}")
+        if message is not None:
+            payload = json.loads(message)
+            print(f"got {payload['text']}")
     await client.aclose()
 
 asyncio.run(main())  # Expected output: got hello
@@ -1044,7 +1049,7 @@ created:
   dead-letter handling is required.`
 
 **AC-16: redis-py is capped below 8.0.0.** The package dependency is
-`redis>=5.0.0,<8.0.0` until redis-py 8 RESP3-default behavior is verified.
+`redis>=5.0.1,<8.0.0` until redis-py 8 RESP3-default behavior is verified.
 Users on redis-py 7.x and earlier are unaffected. If you installed a redis-py
 8.0.0 beta explicitly, downgrade with `pip install "redis<8.0.0"`.
 

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/README.md

@@ -11,7 +11,7 @@
 **Lightweight Python message queuing with Redis and built-in publish-side deduplication.** Deduplicate publishes within a TTL window, with optional crash recovery — across any number of producers and consumers.
 
 ```bash
-pip install "redis-message-queue>=8.2.5,<9.0.0"
+pip install "redis-message-queue>=8.2.7,<9.0.0"
 ```
 
 Requires Redis server >= 6.2.
@@ -23,6 +23,7 @@ Redis must be running locally first: use `redis-server` or
 
 ```python
 import json
+from uuid import uuid4
 from redis import Redis
 from redis_message_queue import RedisMessageQueue
 
@@ -33,7 +34,8 @@ queue = RedisMessageQueue(
     deduplication=True,
     get_deduplication_key=lambda msg: msg["id"],
 )
-queue.publish({"id": "msg-1", "text": "hello"})
+message = {"id": f"msg-{uuid4().hex}", "text": "hello"}
+queue.publish(message)
 with queue.process_message() as message:
     if message is not None:
         payload = json.loads(message)
@@ -57,6 +59,7 @@ with queue.process_message() as message:
 ```python
 import asyncio
 import json
+from uuid import uuid4
 from redis.asyncio import Redis
 from redis_message_queue.asyncio import RedisMessageQueue
 
@@ -68,10 +71,12 @@ async def main():
         deduplication=True,
         get_deduplication_key=lambda msg: msg["id"],
     )
-    await queue.publish({"id": "msg-1", "text": "hello"})
+    message = {"id": f"msg-{uuid4().hex}", "text": "hello"}
+    await queue.publish(message)
     async with queue.process_message() as message:
-        payload = json.loads(message)
-        print(f"got {payload['text']}")
+        if message is not None:
+            payload = json.loads(message)
+            print(f"got {payload['text']}")
     await client.aclose()
 
 asyncio.run(main())  # Expected output: got hello
@@ -1021,7 +1026,7 @@ created:
   dead-letter handling is required.`
 
 **AC-16: redis-py is capped below 8.0.0.** The package dependency is
-`redis>=5.0.0,<8.0.0` until redis-py 8 RESP3-default behavior is verified.
+`redis>=5.0.1,<8.0.0` until redis-py 8 RESP3-default behavior is verified.
 Users on redis-py 7.x and earlier are unaffected. If you installed a redis-py
 8.0.0 beta explicitly, downgrade with `pip install "redis<8.0.0"`.
 

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "redis-message-queue"
-version = "8.2.5"
+version = "8.2.7"
 description = "Python message queuing with Redis and message deduplication"
 authors = [{ name = "Elijas", email = "4084885+Elijas@users.noreply.github.com" }]
 readme = "README.md"
@@ -48,7 +48,7 @@ default-groups = ["dev", "test"]
 ##############################
 
 [tool.bumpversion]
-current_version = "8.2.5"
+current_version = "8.2.7"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/redis_message_queue/_redis_gateway.py

@@ -291,6 +291,7 @@ class RedisGateway(AbstractRedisGateway):
         self._pending_overload_policy = pending_overload_policy
         self._pending_overload_block_timeout_seconds = pending_overload_block_timeout_seconds
         self._pending_claim_ids: dict[str, list[str]] = {}
+        self._in_flight_claim_ids: dict[str, set[str]] = {}
         self._recovering_claim_ids: dict[str, set[str]] = {}
         self._pending_claim_ids_lock = threading.Lock()
         self._drain_pending_claim_ids_lock = threading.Lock()
@@ -715,10 +716,31 @@ class RedisGateway(AbstractRedisGateway):
             return None
 
         pending_claim_id_to_share: str | None = None
+        active_claim_id: str | None = None
+
+        def begin_active_claim(claim_id: str) -> None:
+            nonlocal active_claim_id
+            if active_claim_id == claim_id:
+                return
+            if active_claim_id is not None:
+                self._finish_in_flight_claim_id(to_queue, active_claim_id)
+            self._begin_in_flight_claim_id(to_queue, claim_id)
+            active_claim_id = claim_id
+
+        def finish_active_claim() -> None:
+            nonlocal active_claim_id
+            if active_claim_id is None:
+                return
+            self._finish_in_flight_claim_id(to_queue, active_claim_id)
+            active_claim_id = None
+
         try:
             if self._message_wait_interval_seconds == 0:
                 claim_id = uuid.uuid4().hex
                 claim_may_need_recovery = False
+                begin_active_claim(claim_id)
+                if self._is_interrupted(is_interrupted):
+                    return None
                 try:
                     claimed_message = claim_message(from_queue, to_queue, claim_id)
                 except Exception as exc:
@@ -768,6 +790,8 @@ class RedisGateway(AbstractRedisGateway):
             claim_may_need_recovery = False
             last_retryable_exception: Exception | None = None
             while True:
+                if active_claim_id is None:
+                    begin_active_claim(claim_id)
                 if self._is_interrupted(is_interrupted):
                     if claim_may_need_recovery:
                         pending_claim_id_to_share = claim_id
@@ -796,6 +820,7 @@ class RedisGateway(AbstractRedisGateway):
                         return claimed_message
                     claim_may_need_recovery = False
                     last_retryable_exception = None
+                    finish_active_claim()
                     claim_id = uuid.uuid4().hex
 
                 remaining = deadline - time.monotonic()
@@ -836,6 +861,7 @@ class RedisGateway(AbstractRedisGateway):
         finally:
             if pending_claim_id_to_share is not None:
                 self._set_pending_claim_id(to_queue, pending_claim_id_to_share)
+            finish_active_claim()
 
     def wait_for_message_and_move(self, from_queue: str, to_queue: str) -> ClaimedMessage | MessageData | None:
         if self._is_interrupted():
@@ -1127,6 +1153,19 @@ class RedisGateway(AbstractRedisGateway):
             if claim_id not in pending_claim_ids:
                 pending_claim_ids.append(claim_id)
 
+    def _begin_in_flight_claim_id(self, processing_queue: str, claim_id: str) -> None:
+        with self._pending_claim_ids_lock:
+            self._in_flight_claim_ids.setdefault(processing_queue, set()).add(claim_id)
+
+    def _finish_in_flight_claim_id(self, processing_queue: str, claim_id: str) -> None:
+        with self._pending_claim_ids_lock:
+            in_flight_claim_ids = self._in_flight_claim_ids.get(processing_queue)
+            if in_flight_claim_ids is None:
+                return
+            in_flight_claim_ids.discard(claim_id)
+            if not in_flight_claim_ids:
+                self._in_flight_claim_ids.pop(processing_queue, None)
+
     def _finish_pending_claim_recovery(
         self,
         processing_queue: str,
@@ -1328,15 +1367,29 @@ class RedisGateway(AbstractRedisGateway):
             with self._pending_claim_ids_lock:
                 pending = self._pending_claim_ids.get(processing_queue)
                 if not pending:
-                    return True
-                recovering = self._recovering_claim_ids.setdefault(processing_queue, set())
-                claim_id = next(
-                    (cid for cid in pending if cid not in recovering and cid not in skipped_unresolved),
-                    None,
-                )
-                if claim_id is None:
-                    break
-                recovering.add(claim_id)
+                    in_flight = self._in_flight_claim_ids.get(processing_queue)
+                    if not in_flight:
+                        return True
+                    claim_id = None
+                else:
+                    recovering = self._recovering_claim_ids.setdefault(processing_queue, set())
+                    claim_id = next(
+                        (cid for cid in pending if cid not in recovering and cid not in skipped_unresolved),
+                        None,
+                    )
+                    if claim_id is None:
+                        break
+                    recovering.add(claim_id)
+            if claim_id is None:
+                if deadline_monotonic is None:
+                    time.sleep(_VISIBILITY_TIMEOUT_POLL_INTERVAL_SECONDS)
+                else:
+                    remaining = deadline_monotonic - time.monotonic()
+                    if remaining <= 0:
+                        last_error = TimeoutError("drain pending-claim recovery deadline expired")
+                        break
+                    time.sleep(min(_VISIBILITY_TIMEOUT_POLL_INTERVAL_SECONDS, remaining))
+                continue
             clear = False
             try:
                 try:
@@ -1385,6 +1438,7 @@ class RedisGateway(AbstractRedisGateway):
                 self._finish_pending_claim_recovery(processing_queue, claim_id, clear=clear)
         with self._pending_claim_ids_lock:
             pending = self._pending_claim_ids.get(processing_queue)
-            if pending:
+            in_flight = self._in_flight_claim_ids.get(processing_queue)
+            if pending or in_flight:
                 self._last_drain_error = last_error
-            return not pending
+            return not pending and not in_flight

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/redis_message_queue/asyncio/_redis_gateway.py

@@ -265,6 +265,7 @@ class RedisGateway(AbstractRedisGateway):
         self._pending_overload_policy = pending_overload_policy
         self._pending_overload_block_timeout_seconds = pending_overload_block_timeout_seconds
         self._pending_claim_ids: dict[str, list[str]] = {}
+        self._in_flight_claim_ids: dict[str, set[str]] = {}
         self._recovering_claim_ids: dict[str, set[str]] = {}
         self._pending_claim_ids_lock = threading.Lock()
         self._drain_pending_claim_ids_lock = asyncio.Lock()
@@ -694,10 +695,31 @@ class RedisGateway(AbstractRedisGateway):
             return None
 
         pending_claim_id_to_share: str | None = None
+        active_claim_id: str | None = None
+
+        def begin_active_claim(claim_id: str) -> None:
+            nonlocal active_claim_id
+            if active_claim_id == claim_id:
+                return
+            if active_claim_id is not None:
+                self._finish_in_flight_claim_id(to_queue, active_claim_id)
+            self._begin_in_flight_claim_id(to_queue, claim_id)
+            active_claim_id = claim_id
+
+        def finish_active_claim() -> None:
+            nonlocal active_claim_id
+            if active_claim_id is None:
+                return
+            self._finish_in_flight_claim_id(to_queue, active_claim_id)
+            active_claim_id = None
+
         try:
             if self._message_wait_interval_seconds == 0:
                 claim_id = uuid.uuid4().hex
                 claim_may_need_recovery = False
+                begin_active_claim(claim_id)
+                if self._is_interrupted(is_interrupted):
+                    return None
                 try:
                     claimed_message = await claim_message(from_queue, to_queue, claim_id)
                 except Exception as exc:
@@ -748,6 +770,8 @@ class RedisGateway(AbstractRedisGateway):
             claim_may_need_recovery = False
             last_retryable_exception: Exception | None = None
             while True:
+                if active_claim_id is None:
+                    begin_active_claim(claim_id)
                 if self._is_interrupted(is_interrupted):
                     if claim_may_need_recovery:
                         pending_claim_id_to_share = claim_id
@@ -776,6 +800,7 @@ class RedisGateway(AbstractRedisGateway):
                         return claimed_message
                     claim_may_need_recovery = False
                     last_retryable_exception = None
+                    finish_active_claim()
                     claim_id = uuid.uuid4().hex
 
                 remaining = deadline - loop.time()
@@ -816,6 +841,7 @@ class RedisGateway(AbstractRedisGateway):
         finally:
             if pending_claim_id_to_share is not None:
                 self._set_pending_claim_id(to_queue, pending_claim_id_to_share)
+            finish_active_claim()
 
     async def wait_for_message_and_move(self, from_queue: str, to_queue: str) -> ClaimedMessage | MessageData | None:
         if self._is_interrupted():
@@ -1107,6 +1133,19 @@ class RedisGateway(AbstractRedisGateway):
             if claim_id not in pending_claim_ids:
                 pending_claim_ids.append(claim_id)
 
+    def _begin_in_flight_claim_id(self, processing_queue: str, claim_id: str) -> None:
+        with self._pending_claim_ids_lock:
+            self._in_flight_claim_ids.setdefault(processing_queue, set()).add(claim_id)
+
+    def _finish_in_flight_claim_id(self, processing_queue: str, claim_id: str) -> None:
+        with self._pending_claim_ids_lock:
+            in_flight_claim_ids = self._in_flight_claim_ids.get(processing_queue)
+            if in_flight_claim_ids is None:
+                return
+            in_flight_claim_ids.discard(claim_id)
+            if not in_flight_claim_ids:
+                self._in_flight_claim_ids.pop(processing_queue, None)
+
     def _finish_pending_claim_recovery(
         self,
         processing_queue: str,
@@ -1308,15 +1347,29 @@ class RedisGateway(AbstractRedisGateway):
             with self._pending_claim_ids_lock:
                 pending = self._pending_claim_ids.get(processing_queue)
                 if not pending:
-                    return True
-                recovering = self._recovering_claim_ids.setdefault(processing_queue, set())
-                claim_id = next(
-                    (cid for cid in pending if cid not in recovering and cid not in skipped_unresolved),
-                    None,
-                )
-                if claim_id is None:
-                    break
-                recovering.add(claim_id)
+                    in_flight = self._in_flight_claim_ids.get(processing_queue)
+                    if not in_flight:
+                        return True
+                    claim_id = None
+                else:
+                    recovering = self._recovering_claim_ids.setdefault(processing_queue, set())
+                    claim_id = next(
+                        (cid for cid in pending if cid not in recovering and cid not in skipped_unresolved),
+                        None,
+                    )
+                    if claim_id is None:
+                        break
+                    recovering.add(claim_id)
+            if claim_id is None:
+                if deadline_monotonic is None:
+                    await asyncio.sleep(_VISIBILITY_TIMEOUT_POLL_INTERVAL_SECONDS)
+                else:
+                    remaining = deadline_monotonic - loop.time()
+                    if remaining <= 0:
+                        last_error = TimeoutError("drain pending-claim recovery deadline expired")
+                        break
+                    await asyncio.sleep(min(_VISIBILITY_TIMEOUT_POLL_INTERVAL_SECONDS, remaining))
+                continue
             clear = False
             try:
                 try:
@@ -1365,6 +1418,7 @@ class RedisGateway(AbstractRedisGateway):
                 self._finish_pending_claim_recovery(processing_queue, claim_id, clear=clear)
         with self._pending_claim_ids_lock:
             pending = self._pending_claim_ids.get(processing_queue)
-            if pending:
+            in_flight = self._in_flight_claim_ids.get(processing_queue)
+            if pending or in_flight:
                 self._last_drain_error = last_error
-            return not pending
+            return not pending and not in_flight

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/redis_message_queue/asyncio/redis_message_queue.py

@@ -488,18 +488,22 @@ class _LeaseHeartbeat:
             if current_task is not None and current_task.cancelling() > 0:
                 raise
             logger.exception("on_heartbeat_failure callback raised an exception")
-            warnings.warn(
-                f"on_heartbeat_failure callback raised {type(exc).__name__}",
-                RuntimeWarning,
-                stacklevel=1,
-            )
+            with warnings.catch_warnings():
+                warnings.simplefilter("always", RuntimeWarning)
+                warnings.warn(
+                    f"on_heartbeat_failure callback raised {type(exc).__name__}",
+                    RuntimeWarning,
+                    stacklevel=1,
+                )
         except Exception as exc:
             logger.exception("on_heartbeat_failure callback raised an exception")
-            warnings.warn(
-                f"on_heartbeat_failure callback raised {type(exc).__name__}",
-                RuntimeWarning,
-                stacklevel=1,
-            )
+            with warnings.catch_warnings():
+                warnings.simplefilter("always", RuntimeWarning)
+                warnings.warn(
+                    f"on_heartbeat_failure callback raised {type(exc).__name__}",
+                    RuntimeWarning,
+                    stacklevel=1,
+                )
 
     async def _run(self) -> None:
         try:
@@ -535,13 +539,15 @@ class _LeaseHeartbeat:
                         exception_type=type(exc).__name__,
                         error=exc,
                     )
-                    warnings.warn(
-                        "Failed to renew message lease "
-                        f"({_warning_exception_name(exc)}); message will be reclaimed by another consumer "
-                        "when the visibility timeout expires",
-                        RuntimeWarning,
-                        stacklevel=1,
-                    )
+                    with warnings.catch_warnings():
+                        warnings.simplefilter("always", RuntimeWarning)
+                        warnings.warn(
+                            "Failed to renew message lease "
+                            f"({_warning_exception_name(exc)}); message will be reclaimed by another consumer "
+                            "when the visibility timeout expires",
+                            RuntimeWarning,
+                            stacklevel=1,
+                        )
                     await self._invoke_failure_callback()
                     return
                 if not renewed:
@@ -956,15 +962,22 @@ class RedisMessageQueue:
         pending_claim_ids = getattr(self._redis, "_pending_claim_ids", None)
         if not isinstance(pending_claim_ids, dict):
             return None
+        in_flight_claim_ids = getattr(self._redis, "_in_flight_claim_ids", None)
 
         lock = getattr(self._redis, "_pending_claim_ids_lock", None)
         if lock is None:
             pending = pending_claim_ids.get(self.key.processing)
-            return len(pending) if pending is not None else 0
+            in_flight = in_flight_claim_ids.get(self.key.processing) if isinstance(in_flight_claim_ids, dict) else None
+            pending_count = len(pending) if pending is not None else 0
+            in_flight_count = len(in_flight) if in_flight is not None else 0
+            return pending_count + in_flight_count
 
         with lock:
             pending = pending_claim_ids.get(self.key.processing)
-            return len(pending) if pending is not None else 0
+            in_flight = in_flight_claim_ids.get(self.key.processing) if isinstance(in_flight_claim_ids, dict) else None
+            pending_count = len(pending) if pending is not None else 0
+            in_flight_count = len(in_flight) if in_flight is not None else 0
+            return pending_count + in_flight_count
 
     def _drain_failure_error(self, timeout_seconds: float | None, pending_claim_ids: int | None) -> BaseException:
         last_drain_error = getattr(self._redis, "_last_drain_error", None)
@@ -1059,9 +1072,23 @@ class RedisMessageQueue:
                         "See AbstractRedisGateway.add_message for the full contract."
                     )
             else:
-                dedup_key = self._get_deduplication_key(message)
-                if inspect.isawaitable(dedup_key):
-                    dedup_key = await dedup_key
+                try:
+                    dedup_key = self._get_deduplication_key(message)
+                    if inspect.isawaitable(dedup_key):
+                        dedup_key = await dedup_key
+                except asyncio.CancelledError as exc:
+                    current_task = asyncio.current_task()
+                    if current_task is not None and current_task.cancelling() > 0:
+                        raise
+                    _set_exception_context(exc, queue=self._queue_name, operation="publish")
+                    await self._emit_event(
+                        "publish",
+                        "failure",
+                        exception_type=type(exc).__name__,
+                        error=exc,
+                        duration_ms=_duration_ms(started_at),
+                    )
+                    raise
                 dedup_key = validate_callable_deduplication_key(dedup_key, message)
                 dedup_key = self.key.deduplication(dedup_key)
                 result = await self._redis.publish_message(self.key.pending, message_str, dedup_key)
@@ -1466,16 +1493,20 @@ class RedisMessageQueue:
         async with self._aclose_lock:
             cleanup_lease_counter = getattr(self._redis, "_cleanup_drained_lease_token_counter", None)
             if self._aclose_result is not None:
-                if cleanup_lease_counter is not None:
-                    await _await_preserving_cancellation(cleanup_lease_counter(self.key.processing))
-                await self._emit_event(
-                    "drain",
-                    "skipped",
-                    duration_ms=_duration_ms(started_at),
-                    timeout_seconds=timeout_seconds,
-                    pending_claim_ids=self._pending_claim_ids_count(),
-                )
-                return self._aclose_result
+                pending_claim_ids = self._pending_claim_ids_count()
+                if pending_claim_ids:
+                    self._aclose_result = None
+                else:
+                    if cleanup_lease_counter is not None:
+                        await _await_preserving_cancellation(cleanup_lease_counter(self.key.processing))
+                    await self._emit_event(
+                        "drain",
+                        "skipped",
+                        duration_ms=_duration_ms(started_at),
+                        timeout_seconds=timeout_seconds,
+                        pending_claim_ids=pending_claim_ids,
+                    )
+                    return self._aclose_result
 
             async with self._publish_lock:
                 self._draining = True

{redis_message_queue-8.2.5 → redis_message_queue-8.2.7}/redis_message_queue/redis_message_queue.py

@@ -80,6 +80,14 @@ def _duration_ms(started_at: float) -> float:
     return (time.perf_counter() - started_at) * 1000
 
 
+def _current_async_task_is_cancelling() -> bool:
+    try:
+        current_task = asyncio.current_task()
+    except RuntimeError:
+        return False
+    return current_task is not None and current_task.cancelling() > 0
+
+
 def _hash_lease_token(lease_token: str | None) -> str | None:
     if lease_token is None:
         return None
@@ -447,18 +455,22 @@ class _LeaseHeartbeat:
                 )
         except asyncio.CancelledError as exc:
             logger.exception("on_heartbeat_failure callback raised an exception")
-            warnings.warn(
-                f"on_heartbeat_failure callback raised {type(exc).__name__}",
-                RuntimeWarning,
-                stacklevel=1,
-            )
+            with warnings.catch_warnings():
+                warnings.simplefilter("always", RuntimeWarning)
+                warnings.warn(
+                    f"on_heartbeat_failure callback raised {type(exc).__name__}",
+                    RuntimeWarning,
+                    stacklevel=1,
+                )
         except Exception as exc:
             logger.exception("on_heartbeat_failure callback raised an exception")
-            warnings.warn(
-                f"on_heartbeat_failure callback raised {type(exc).__name__}",
-                RuntimeWarning,
-                stacklevel=1,
-            )
+            with warnings.catch_warnings():
+                warnings.simplefilter("always", RuntimeWarning)
+                warnings.warn(
+                    f"on_heartbeat_failure callback raised {type(exc).__name__}",
+                    RuntimeWarning,
+                    stacklevel=1,
+                )
 
     def _run(self) -> None:
         # No explicit _is_interrupted() check here. Heartbeat lifetime is owned
@@ -488,13 +500,15 @@ class _LeaseHeartbeat:
                     exception_type=type(exc).__name__,
                     error=exc,
                 )
-                warnings.warn(
-                    "Failed to renew message lease "
-                    f"({_warning_exception_name(exc)}); message will be reclaimed by another consumer "
-                    "when the visibility timeout expires",
-                    RuntimeWarning,
-                    stacklevel=1,
-                )
+                with warnings.catch_warnings():
+                    warnings.simplefilter("always", RuntimeWarning)
+                    warnings.warn(
+                        "Failed to renew message lease "
+                        f"({_warning_exception_name(exc)}); message will be reclaimed by another consumer "
+                        "when the visibility timeout expires",
+                        RuntimeWarning,
+                        stacklevel=1,
+                    )
                 self._invoke_failure_callback()
                 return
             if not renewed:
@@ -914,15 +928,22 @@ class RedisMessageQueue:
         pending_claim_ids = getattr(self._redis, "_pending_claim_ids", None)
         if not isinstance(pending_claim_ids, dict):
             return None
+        in_flight_claim_ids = getattr(self._redis, "_in_flight_claim_ids", None)
 
         lock = getattr(self._redis, "_pending_claim_ids_lock", None)
         if lock is None:
             pending = pending_claim_ids.get(self.key.processing)
-            return len(pending) if pending is not None else 0
+            in_flight = in_flight_claim_ids.get(self.key.processing) if isinstance(in_flight_claim_ids, dict) else None
+            pending_count = len(pending) if pending is not None else 0
+            in_flight_count = len(in_flight) if in_flight is not None else 0
+            return pending_count + in_flight_count
 
         with lock:
             pending = pending_claim_ids.get(self.key.processing)
-            return len(pending) if pending is not None else 0
+            in_flight = in_flight_claim_ids.get(self.key.processing) if isinstance(in_flight_claim_ids, dict) else None
+            pending_count = len(pending) if pending is not None else 0
+            in_flight_count = len(in_flight) if in_flight is not None else 0
+            return pending_count + in_flight_count
 
     def _drain_failure_error(self, timeout_seconds: float | None, pending_claim_ids: int | None) -> BaseException:
         last_drain_error = getattr(self._redis, "_last_drain_error", None)
@@ -1004,7 +1025,20 @@ class RedisMessageQueue:
                         "See AbstractRedisGateway.add_message for the full contract."
                     )
             else:
-                dedup_key = self._get_deduplication_key(message)
+                try:
+                    dedup_key = self._get_deduplication_key(message)
+                except asyncio.CancelledError as exc:
+                    if _current_async_task_is_cancelling():
+                        raise
+                    _set_exception_context(exc, queue=self._queue_name, operation="publish")
+                    self._emit_event(
+                        "publish",
+                        "failure",
+                        exception_type=type(exc).__name__,
+                        error=exc,
+                        duration_ms=_duration_ms(started_at),
+                    )
+                    raise
                 if inspect.isawaitable(dedup_key):
                     is_coroutine = inspect.iscoroutine(dedup_key)
                     _close_or_cancel_awaitable(dedup_key)
@@ -1426,16 +1460,20 @@ class RedisMessageQueue:
         with self._drain_lock:
             cleanup_lease_counter = getattr(self._redis, "_cleanup_drained_lease_token_counter", None)
             if self._drain_result is True:
-                if cleanup_lease_counter is not None:
-                    cleanup_lease_counter(self.key.processing)
-                self._emit_event(
-                    "drain",
-                    "skipped",
-                    duration_ms=_duration_ms(started_at),
-                    timeout_seconds=timeout_seconds,
-                    pending_claim_ids=self._pending_claim_ids_count(),
-                )
-                return True
+                pending_claim_ids = self._pending_claim_ids_count()
+                if pending_claim_ids:
+                    self._drain_result = None
+                else:
+                    if cleanup_lease_counter is not None:
+                        cleanup_lease_counter(self.key.processing)
+                    self._emit_event(
+                        "drain",
+                        "skipped",
+                        duration_ms=_duration_ms(started_at),
+                        timeout_seconds=timeout_seconds,
+                        pending_claim_ids=pending_claim_ids,
+                    )
+                    return True
 
             with self._publish_lock:
                 self._draining = True