redis-message-queue 8.0.0__tar.gz → 8.0.1__tar.gz

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: redis-message-queue
-Version: 8.0.0
+Version: 8.0.1
 Summary: Python message queuing with Redis and message deduplication
 License: MIT
 License-File: LICENSE
@@ -26,7 +26,7 @@ Description-Content-Type: text/markdown
 
 # redis-message-queue
 
-[![PyPI Version](https://img.shields.io/badge/v8.0.0-version?color=43cd0f&style=flat&label=pypi)](https://pypi.org/project/redis-message-queue)
+[![PyPI Version](https://img.shields.io/badge/v8.0.1-version?color=43cd0f&style=flat&label=pypi)](https://pypi.org/project/redis-message-queue)
 [![PyPI Downloads](https://img.shields.io/pypi/dm/redis-message-queue?color=43cd0f&style=flat&label=downloads)](https://pypistats.org/packages/redis-message-queue)
 [![License: MIT](https://img.shields.io/badge/License-MIT-43cd0f.svg?style=flat&label=license)](LICENSE)
 [![Maintained: yes](https://img.shields.io/badge/yes-43cd0f.svg?style=flat&label=maintained)](https://github.com/Elijas/redis-message-queue/issues)

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/README.md

@@ -1,6 +1,6 @@
 # redis-message-queue
 
-[![PyPI Version](https://img.shields.io/badge/v8.0.0-version?color=43cd0f&style=flat&label=pypi)](https://pypi.org/project/redis-message-queue)
+[![PyPI Version](https://img.shields.io/badge/v8.0.1-version?color=43cd0f&style=flat&label=pypi)](https://pypi.org/project/redis-message-queue)
 [![PyPI Downloads](https://img.shields.io/pypi/dm/redis-message-queue?color=43cd0f&style=flat&label=downloads)](https://pypistats.org/packages/redis-message-queue)
 [![License: MIT](https://img.shields.io/badge/License-MIT-43cd0f.svg?style=flat&label=license)](LICENSE)
 [![Maintained: yes](https://img.shields.io/badge/yes-43cd0f.svg?style=flat&label=maintained)](https://github.com/Elijas/redis-message-queue/issues)

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redis-message-queue"
-version = "8.0.0"
+version = "8.0.1"
 description = "Python message queuing with Redis and message deduplication"
 authors = ["Elijas <4084885+Elijas@users.noreply.github.com>"]
 readme = "README.md"

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/redis_message_queue/_redis_gateway.py

@@ -1,10 +1,11 @@
 import json
 import logging
+import queue
 import random
 import threading
 import time
 import uuid
-from typing import Callable, Optional, TypeVar
+from typing import Callable, Optional, TypeVar, cast
 
 import redis
 import redis.asyncio
@@ -54,6 +55,7 @@ from redis_message_queue.interrupt_handler._interface import (
 
 logger = logging.getLogger(__name__)
 _TClaim = TypeVar("_TClaim", bound=ClaimedMessage | MessageData)
+_TRedisCall = TypeVar("_TRedisCall")
 
 _LEASE_DEADLINES_SUFFIX = ":lease_deadlines"
 _LEASE_TOKENS_SUFFIX = ":lease_tokens"
@@ -71,6 +73,49 @@ _PENDING_OVERLOAD_INITIAL_BACKOFF_SECONDS = 0.010
 _PENDING_OVERLOAD_MAX_BACKOFF_SECONDS = 0.500
 
 
+class _DrainDeadlineExceeded(Exception):
+    """Internal sentinel for drain-only pending-claim recovery deadlines."""
+
+
+def _raise_if_drain_deadline_expired(deadline_monotonic: float | None) -> None:
+    if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
+        raise _DrainDeadlineExceeded
+
+
+def _call_with_drain_deadline(
+    call: Callable[[], _TRedisCall],
+    *,
+    deadline_monotonic: float | None,
+) -> _TRedisCall:
+    if deadline_monotonic is None:
+        return call()
+
+    remaining_seconds = deadline_monotonic - time.monotonic()
+    if remaining_seconds <= 0:
+        raise _DrainDeadlineExceeded
+
+    result_queue: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1)
+
+    def run_call() -> None:
+        try:
+            result = call()
+        except Exception as exc:
+            result_queue.put_nowait((False, exc))
+        else:
+            result_queue.put_nowait((True, result))
+
+    thread = threading.Thread(target=run_call, daemon=True)
+    thread.start()
+    thread.join(timeout=remaining_seconds)
+    if thread.is_alive():
+        raise _DrainDeadlineExceeded
+
+    succeeded, value = result_queue.get_nowait()
+    if succeeded:
+        return cast(_TRedisCall, value)
+    raise cast(BaseException, value)
+
+
 def _coerce_lua_count(value: object) -> int:
     if isinstance(value, bytes):
         value = value.decode("utf-8")
@@ -525,6 +570,7 @@ class RedisGateway(AbstractRedisGateway):
         claim_message: Callable[[str, str, str], _TClaim | None],
         non_blocking_retry_log: str,
         polling_retry_log: str,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
     ) -> _TClaim | None:
         while True:
             pending_claim_id = self._acquire_pending_claim_id(to_queue)
@@ -552,7 +598,7 @@ class RedisGateway(AbstractRedisGateway):
                 self._emit_event("claim_reclaim", "success", claim_id=pending_claim_id)
                 return recovered_claim
 
-        if self._is_interrupted():
+        if self._is_interrupted(is_interrupted):
             return None
 
         pending_claim_id_to_share: str | None = None
@@ -575,7 +621,7 @@ class RedisGateway(AbstractRedisGateway):
                         error=exc,
                     )
                     logger.warning(non_blocking_retry_log, type(exc).__name__)
-                    if self._is_interrupted():
+                    if self._is_interrupted(is_interrupted):
                         pending_claim_id_to_share = claim_id
                         return None
                     try:
@@ -607,7 +653,7 @@ class RedisGateway(AbstractRedisGateway):
             claim_may_need_recovery = False
             last_retryable_exception: Exception | None = None
             while True:
-                if self._is_interrupted():
+                if self._is_interrupted(is_interrupted):
                     if claim_may_need_recovery:
                         pending_claim_id_to_share = claim_id
                     return None
@@ -640,7 +686,7 @@ class RedisGateway(AbstractRedisGateway):
                 remaining = deadline - time.monotonic()
                 if remaining <= 0:
                     if last_retryable_exception is not None:
-                        if self._is_interrupted():
+                        if self._is_interrupted(is_interrupted):
                             if claim_may_need_recovery:
                                 pending_claim_id_to_share = claim_id
                             return None
@@ -675,11 +721,36 @@ class RedisGateway(AbstractRedisGateway):
     def wait_for_message_and_move(self, from_queue: str, to_queue: str) -> ClaimedMessage | MessageData | None:
         if self._is_interrupted():
             return None
+        return self._wait_for_message_and_move_interruptible(from_queue, to_queue)
+
+    def _wait_for_message_and_move_interruptible(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> ClaimedMessage | MessageData | None:
+        if self._is_interrupted(is_interrupted):
+            return None
         if self._message_visibility_timeout_seconds is not None:
-            return self._wait_for_message_with_visibility_timeout(from_queue, to_queue)
-        return self._wait_for_message_without_visibility_timeout(from_queue, to_queue)
+            return self._wait_for_message_with_visibility_timeout(
+                from_queue,
+                to_queue,
+                is_interrupted=is_interrupted,
+            )
+        return self._wait_for_message_without_visibility_timeout(
+            from_queue,
+            to_queue,
+            is_interrupted=is_interrupted,
+        )
 
-    def _wait_for_message_without_visibility_timeout(self, from_queue: str, to_queue: str) -> MessageData | None:
+    def _wait_for_message_without_visibility_timeout(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> MessageData | None:
         return self._wait_for_claim(
             from_queue,
             to_queue,
@@ -693,9 +764,16 @@ class RedisGateway(AbstractRedisGateway):
                 "Transient error during non-visibility-timeout non-blocking claim, retrying once to recover claim: %s"
             ),
             polling_retry_log="Transient error during non-visibility-timeout claim poll, will retry: %s",
+            is_interrupted=is_interrupted,
         )
 
-    def _wait_for_message_with_visibility_timeout(self, from_queue: str, to_queue: str) -> ClaimedMessage | None:
+    def _wait_for_message_with_visibility_timeout(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> ClaimedMessage | None:
         return self._wait_for_claim(
             from_queue,
             to_queue,
@@ -709,6 +787,7 @@ class RedisGateway(AbstractRedisGateway):
                 "Transient error during visibility-timeout non-blocking claim, retrying once to recover claim: %s"
             ),
             polling_retry_log="Transient error during visibility-timeout claim poll, will retry: %s",
+            is_interrupted=is_interrupted,
         )
 
     def _claim_message_without_visibility_timeout(
@@ -833,16 +912,37 @@ class RedisGateway(AbstractRedisGateway):
     def _claim_result_ttl_ms(self) -> str:
         return str(max(self._message_wait_interval_seconds, 120) * 1000)
 
-    def _delete_claim_result_key(self, claim_result_key: str) -> None:
+    def _delete_claim_result_key(
+        self,
+        claim_result_key: str,
+        *,
+        deadline_monotonic: float | None = None,
+    ) -> None:
         try:
-            self._redis_client.delete(claim_result_key)
+            _call_with_drain_deadline(
+                lambda: self._redis_client.delete(claim_result_key),
+                deadline_monotonic=deadline_monotonic,
+            )
+        except _DrainDeadlineExceeded:
+            raise
         except Exception:
             # Claim-result keys have bounded TTLs; this cleanup is intentionally best-effort.
             logger.warning("Failed to delete claim result key %s", claim_result_key, exc_info=True)
 
-    def _delete_claim_result_ref(self, claim_result_refs_key: str, lease_token: str) -> None:
+    def _delete_claim_result_ref(
+        self,
+        claim_result_refs_key: str,
+        lease_token: str,
+        *,
+        deadline_monotonic: float | None = None,
+    ) -> None:
         try:
-            self._redis_client.hdel(claim_result_refs_key, lease_token)
+            _call_with_drain_deadline(
+                lambda: self._redis_client.hdel(claim_result_refs_key, lease_token),
+                deadline_monotonic=deadline_monotonic,
+            )
+        except _DrainDeadlineExceeded:
+            raise
         except Exception:
             # Claim-result refs have bounded TTLs; this cleanup is intentionally best-effort.
             logger.warning(
@@ -906,25 +1006,48 @@ class RedisGateway(AbstractRedisGateway):
         self,
         processing_queue: str,
         claim_id: str,
+        *,
+        deadline_monotonic: float | None = None,
     ) -> MessageData | None:
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         claim_result_key = self._claim_result_key(processing_queue, claim_id)
-        cached_claim = self._redis_client.get(claim_result_key)
+        cached_claim = _call_with_drain_deadline(
+            lambda: self._redis_client.get(claim_result_key),
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         if cached_claim is None:
-            cached_claim = self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id)
+            cached_claim = _call_with_drain_deadline(
+                lambda: self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id),
+                deadline_monotonic=deadline_monotonic,
+            )
+            _raise_if_drain_deadline_expired(deadline_monotonic)
             if cached_claim is None:
                 return None
-        self._delete_claim_result_key(claim_result_key)
+        self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         return cached_claim
 
     def _recover_pending_visibility_timeout_claim(
         self,
         processing_queue: str,
         claim_id: str,
+        *,
+        deadline_monotonic: float | None = None,
     ) -> ClaimedMessage | None:
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         claim_result_key = self._claim_result_key(processing_queue, claim_id)
-        cached_claim = self._redis_client.get(claim_result_key)
+        cached_claim = _call_with_drain_deadline(
+            lambda: self._redis_client.get(claim_result_key),
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         if cached_claim is None:
-            cached_claim = self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id)
+            cached_claim = _call_with_drain_deadline(
+                lambda: self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id),
+                deadline_monotonic=deadline_monotonic,
+            )
+            _raise_if_drain_deadline_expired(deadline_monotonic)
             if cached_claim is None:
                 return None
 
@@ -932,7 +1055,7 @@ class RedisGateway(AbstractRedisGateway):
         try:
             claim = json.loads(cached_claim_text)
         except json.JSONDecodeError:
-            self._delete_claim_result_key(claim_result_key)
+            self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
             return None
 
         if (
@@ -941,7 +1064,7 @@ class RedisGateway(AbstractRedisGateway):
             or not isinstance(claim[0], str)
             or not isinstance(claim[1], str)
         ):
-            self._delete_claim_result_key(claim_result_key)
+            self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
             return None
 
         stored_message: MessageData = claim[0]
@@ -949,12 +1072,19 @@ class RedisGateway(AbstractRedisGateway):
             stored_message = stored_message.encode("utf-8")
         lease_token = claim[1]
 
-        self._delete_claim_result_key(claim_result_key)
-        self._delete_claim_result_ref(self._claim_result_refs_key(processing_queue), lease_token)
+        self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
+        self._delete_claim_result_ref(
+            self._claim_result_refs_key(processing_queue),
+            lease_token,
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         return ClaimedMessage(stored_message=stored_message, lease_token=lease_token)
 
-    def _is_interrupted(self) -> bool:
-        return self._interrupt is not None and self._interrupt.is_interrupted()
+    def _is_interrupted(self, is_interrupted: BaseGracefulInterruptHandler | None = None) -> bool:
+        return (self._interrupt is not None and self._interrupt.is_interrupted()) or (
+            is_interrupted is not None and is_interrupted.is_interrupted()
+        )
 
     def _drain_pending_claim_ids(
         self,
@@ -1010,8 +1140,10 @@ class RedisGateway(AbstractRedisGateway):
             clear = False
             try:
                 try:
-                    recover(processing_queue, claim_id)
+                    recover(processing_queue, claim_id, deadline_monotonic=deadline_monotonic)
                     clear = True
+                except _DrainDeadlineExceeded:
+                    break
                 except Exception as exc:
                     if not is_redis_retryable_exception(exc):
                         raise

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/redis_message_queue/asyncio/_redis_gateway.py

@@ -54,6 +54,7 @@ from redis_message_queue.interrupt_handler._interface import (
 
 logger = logging.getLogger(__name__)
 _TClaim = TypeVar("_TClaim", bound=ClaimedMessage | MessageData)
+_TRedisCall = TypeVar("_TRedisCall")
 
 _LEASE_DEADLINES_SUFFIX = ":lease_deadlines"
 _LEASE_TOKENS_SUFFIX = ":lease_tokens"
@@ -71,6 +72,33 @@ _PENDING_OVERLOAD_INITIAL_BACKOFF_SECONDS = 0.010
 _PENDING_OVERLOAD_MAX_BACKOFF_SECONDS = 0.500
 
 
+class _DrainDeadlineExceeded(Exception):
+    """Internal sentinel for drain-only pending-claim recovery deadlines."""
+
+
+def _raise_if_drain_deadline_expired(deadline_monotonic: float | None) -> None:
+    if deadline_monotonic is not None and asyncio.get_running_loop().time() >= deadline_monotonic:
+        raise _DrainDeadlineExceeded
+
+
+async def _call_with_drain_deadline(
+    call: Callable[[], Awaitable[_TRedisCall]],
+    *,
+    deadline_monotonic: float | None,
+) -> _TRedisCall:
+    if deadline_monotonic is None:
+        return await call()
+
+    remaining_seconds = deadline_monotonic - asyncio.get_running_loop().time()
+    if remaining_seconds <= 0:
+        raise _DrainDeadlineExceeded
+
+    try:
+        return await asyncio.wait_for(call(), timeout=remaining_seconds)
+    except TimeoutError as exc:
+        raise _DrainDeadlineExceeded from exc
+
+
 def _coerce_lua_count(value: object) -> int:
     if isinstance(value, bytes):
         value = value.decode("utf-8")
@@ -521,6 +549,7 @@ class RedisGateway(AbstractRedisGateway):
         claim_message: Callable[[str, str, str], Awaitable[_TClaim | None]],
         non_blocking_retry_log: str,
         polling_retry_log: str,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
     ) -> _TClaim | None:
         while True:
             pending_claim_id = self._acquire_pending_claim_id(to_queue)
@@ -548,7 +577,7 @@ class RedisGateway(AbstractRedisGateway):
                 await self._emit_event("claim_reclaim", "success", claim_id=pending_claim_id)
                 return recovered_claim
 
-        if self._is_interrupted():
+        if self._is_interrupted(is_interrupted):
             return None
 
         pending_claim_id_to_share: str | None = None
@@ -571,7 +600,7 @@ class RedisGateway(AbstractRedisGateway):
                         error=exc,
                     )
                     logger.warning(non_blocking_retry_log, type(exc).__name__)
-                    if self._is_interrupted():
+                    if self._is_interrupted(is_interrupted):
                         pending_claim_id_to_share = claim_id
                         return None
                     try:
@@ -604,7 +633,7 @@ class RedisGateway(AbstractRedisGateway):
             claim_may_need_recovery = False
             last_retryable_exception: Exception | None = None
             while True:
-                if self._is_interrupted():
+                if self._is_interrupted(is_interrupted):
                     if claim_may_need_recovery:
                         pending_claim_id_to_share = claim_id
                     return None
@@ -637,7 +666,7 @@ class RedisGateway(AbstractRedisGateway):
                 remaining = deadline - loop.time()
                 if remaining <= 0:
                     if last_retryable_exception is not None:
-                        if self._is_interrupted():
+                        if self._is_interrupted(is_interrupted):
                             if claim_may_need_recovery:
                                 pending_claim_id_to_share = claim_id
                             return None
@@ -672,11 +701,36 @@ class RedisGateway(AbstractRedisGateway):
     async def wait_for_message_and_move(self, from_queue: str, to_queue: str) -> ClaimedMessage | MessageData | None:
         if self._is_interrupted():
             return None
+        return await self._wait_for_message_and_move_interruptible(from_queue, to_queue)
+
+    async def _wait_for_message_and_move_interruptible(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> ClaimedMessage | MessageData | None:
+        if self._is_interrupted(is_interrupted):
+            return None
         if self._message_visibility_timeout_seconds is not None:
-            return await self._wait_for_message_with_visibility_timeout(from_queue, to_queue)
-        return await self._wait_for_message_without_visibility_timeout(from_queue, to_queue)
+            return await self._wait_for_message_with_visibility_timeout(
+                from_queue,
+                to_queue,
+                is_interrupted=is_interrupted,
+            )
+        return await self._wait_for_message_without_visibility_timeout(
+            from_queue,
+            to_queue,
+            is_interrupted=is_interrupted,
+        )
 
-    async def _wait_for_message_without_visibility_timeout(self, from_queue: str, to_queue: str) -> MessageData | None:
+    async def _wait_for_message_without_visibility_timeout(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> MessageData | None:
         return await self._wait_for_claim(
             from_queue,
             to_queue,
@@ -690,9 +744,16 @@ class RedisGateway(AbstractRedisGateway):
                 "Transient error during non-visibility-timeout non-blocking claim, retrying once to recover claim: %s"
             ),
             polling_retry_log="Transient error during non-visibility-timeout claim poll, will retry: %s",
+            is_interrupted=is_interrupted,
         )
 
-    async def _wait_for_message_with_visibility_timeout(self, from_queue: str, to_queue: str) -> ClaimedMessage | None:
+    async def _wait_for_message_with_visibility_timeout(
+        self,
+        from_queue: str,
+        to_queue: str,
+        *,
+        is_interrupted: BaseGracefulInterruptHandler | None = None,
+    ) -> ClaimedMessage | None:
         return await self._wait_for_claim(
             from_queue,
             to_queue,
@@ -706,6 +767,7 @@ class RedisGateway(AbstractRedisGateway):
                 "Transient error during visibility-timeout non-blocking claim, retrying once to recover claim: %s"
             ),
             polling_retry_log="Transient error during visibility-timeout claim poll, will retry: %s",
+            is_interrupted=is_interrupted,
         )
 
     async def _claim_message_without_visibility_timeout(
@@ -830,16 +892,37 @@ class RedisGateway(AbstractRedisGateway):
     def _claim_result_ttl_ms(self) -> str:
         return str(max(self._message_wait_interval_seconds, 120) * 1000)
 
-    async def _delete_claim_result_key(self, claim_result_key: str) -> None:
+    async def _delete_claim_result_key(
+        self,
+        claim_result_key: str,
+        *,
+        deadline_monotonic: float | None = None,
+    ) -> None:
         try:
-            await self._redis_client.delete(claim_result_key)
+            await _call_with_drain_deadline(
+                lambda: self._redis_client.delete(claim_result_key),
+                deadline_monotonic=deadline_monotonic,
+            )
+        except _DrainDeadlineExceeded:
+            raise
         except Exception:
             # Claim-result keys have bounded TTLs; this cleanup is intentionally best-effort.
             logger.warning("Failed to delete claim result key %s", claim_result_key, exc_info=True)
 
-    async def _delete_claim_result_ref(self, claim_result_refs_key: str, lease_token: str) -> None:
+    async def _delete_claim_result_ref(
+        self,
+        claim_result_refs_key: str,
+        lease_token: str,
+        *,
+        deadline_monotonic: float | None = None,
+    ) -> None:
         try:
-            await self._redis_client.hdel(claim_result_refs_key, lease_token)
+            await _call_with_drain_deadline(
+                lambda: self._redis_client.hdel(claim_result_refs_key, lease_token),
+                deadline_monotonic=deadline_monotonic,
+            )
+        except _DrainDeadlineExceeded:
+            raise
         except Exception:
             # Claim-result refs have bounded TTLs; this cleanup is intentionally best-effort.
             logger.warning(
@@ -903,25 +986,48 @@ class RedisGateway(AbstractRedisGateway):
         self,
         processing_queue: str,
         claim_id: str,
+        *,
+        deadline_monotonic: float | None = None,
     ) -> MessageData | None:
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         claim_result_key = self._claim_result_key(processing_queue, claim_id)
-        cached_claim = await self._redis_client.get(claim_result_key)
+        cached_claim = await _call_with_drain_deadline(
+            lambda: self._redis_client.get(claim_result_key),
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         if cached_claim is None:
-            cached_claim = await self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id)
+            cached_claim = await _call_with_drain_deadline(
+                lambda: self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id),
+                deadline_monotonic=deadline_monotonic,
+            )
+            _raise_if_drain_deadline_expired(deadline_monotonic)
             if cached_claim is None:
                 return None
-        await self._delete_claim_result_key(claim_result_key)
+        await self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         return cached_claim
 
     async def _recover_pending_visibility_timeout_claim(
         self,
         processing_queue: str,
         claim_id: str,
+        *,
+        deadline_monotonic: float | None = None,
     ) -> ClaimedMessage | None:
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         claim_result_key = self._claim_result_key(processing_queue, claim_id)
-        cached_claim = await self._redis_client.get(claim_result_key)
+        cached_claim = await _call_with_drain_deadline(
+            lambda: self._redis_client.get(claim_result_key),
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         if cached_claim is None:
-            cached_claim = await self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id)
+            cached_claim = await _call_with_drain_deadline(
+                lambda: self._redis_client.hget(self._claim_result_ids_key(processing_queue), claim_id),
+                deadline_monotonic=deadline_monotonic,
+            )
+            _raise_if_drain_deadline_expired(deadline_monotonic)
             if cached_claim is None:
                 return None
 
@@ -929,7 +1035,7 @@ class RedisGateway(AbstractRedisGateway):
         try:
             claim = json.loads(cached_claim_text)
         except json.JSONDecodeError:
-            await self._delete_claim_result_key(claim_result_key)
+            await self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
             return None
 
         if (
@@ -938,7 +1044,7 @@ class RedisGateway(AbstractRedisGateway):
             or not isinstance(claim[0], str)
             or not isinstance(claim[1], str)
         ):
-            await self._delete_claim_result_key(claim_result_key)
+            await self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
             return None
 
         stored_message: MessageData = claim[0]
@@ -946,12 +1052,20 @@ class RedisGateway(AbstractRedisGateway):
             stored_message = stored_message.encode("utf-8")
         lease_token = claim[1]
 
-        await self._delete_claim_result_key(claim_result_key)
-        await self._delete_claim_result_ref(self._claim_result_refs_key(processing_queue), lease_token)
+        await self._delete_claim_result_key(claim_result_key, deadline_monotonic=deadline_monotonic)
+        _raise_if_drain_deadline_expired(deadline_monotonic)
+        await self._delete_claim_result_ref(
+            self._claim_result_refs_key(processing_queue),
+            lease_token,
+            deadline_monotonic=deadline_monotonic,
+        )
+        _raise_if_drain_deadline_expired(deadline_monotonic)
         return ClaimedMessage(stored_message=stored_message, lease_token=lease_token)
 
-    def _is_interrupted(self) -> bool:
-        return self._interrupt is not None and self._interrupt.is_interrupted()
+    def _is_interrupted(self, is_interrupted: BaseGracefulInterruptHandler | None = None) -> bool:
+        return (self._interrupt is not None and self._interrupt.is_interrupted()) or (
+            is_interrupted is not None and is_interrupted.is_interrupted()
+        )
 
     async def _drain_pending_claim_ids(
         self,
@@ -1006,8 +1120,10 @@ class RedisGateway(AbstractRedisGateway):
             clear = False
             try:
                 try:
-                    await recover(processing_queue, claim_id)
+                    await recover(processing_queue, claim_id, deadline_monotonic=deadline_monotonic)
                     clear = True
+                except _DrainDeadlineExceeded:
+                    break
                 except Exception as exc:
                     if not is_redis_retryable_exception(exc):
                         raise

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/redis_message_queue/asyncio/redis_message_queue.py

@@ -74,6 +74,30 @@ def _warning_exception_name(exc: BaseException) -> str:
     return type(exc).__name__
 
 
+def _find_non_string_dict_keys(value: object) -> list[object]:
+    non_str_keys: list[object] = []
+    seen: set[int] = set()
+
+    def visit(current: object) -> None:
+        if not isinstance(current, (dict, list, tuple)):
+            return
+        current_id = id(current)
+        if current_id in seen:
+            return
+        seen.add(current_id)
+        if isinstance(current, dict):
+            for key, child in current.items():
+                if not isinstance(key, str):
+                    non_str_keys.append(key)
+                visit(child)
+            return
+        for child in current:
+            visit(child)
+
+    visit(value)
+    return non_str_keys
+
+
 class _TaskBaseException(Exception):
     def __init__(self, original: BaseException):
         super().__init__(str(original))
@@ -285,6 +309,14 @@ class _StopEventInterrupt(BaseGracefulInterruptHandler):
         return self._stop_event.is_set()
 
 
+class _DrainInterrupt(BaseGracefulInterruptHandler):
+    def __init__(self, is_draining: Callable[[], bool]) -> None:
+        self._is_draining = is_draining
+
+    def is_interrupted(self) -> bool:
+        return self._is_draining()
+
+
 class _LeaseHeartbeat:
     def __init__(
         self,
@@ -758,11 +790,10 @@ class RedisMessageQueue:
         """Publish a message.
 
         Dict messages are serialized via ``json.dumps(message, sort_keys=True)``.
-        All top-level dict keys must be strings; non-string keys raise
+        All dict keys must be strings; non-string keys raise
         ``TypeError`` to avoid silent ``json.dumps`` coercion that would
         collapse distinct keys into the same dedup key (e.g. ``{1: "x"}``
-        vs ``{"1": "x"}``). Only top-level keys are validated; nested
-        dicts follow ``json.dumps`` defaults.
+        vs ``{"1": "x"}``).
         """
         async with self._publish_lock:
             if self._drained:
@@ -773,7 +804,7 @@ class RedisMessageQueue:
         if not isinstance(message, (str, dict)):
             raise TypeError(f"'message' must be a str or dict, got {type(message).__name__}")
         if isinstance(message, dict):
-            non_str_keys = [k for k in message if not isinstance(k, str)]
+            non_str_keys = _find_non_string_dict_keys(message)
             if non_str_keys:
                 raise TypeError(
                     "'message' dict keys must all be strings; "
@@ -846,10 +877,7 @@ class RedisMessageQueue:
             yield None
             return
         try:
-            claimed_message = await self._redis.wait_for_message_and_move(
-                self.key.pending,
-                self.key.processing,
-            )
+            claimed_message = await self._wait_for_message_and_move()
         except Exception as exc:
             await self._emit_event(
                 "claim",
@@ -1027,6 +1055,19 @@ class RedisMessageQueue:
                 else:
                     await _await_suppressing_external_cancellation(lease_heartbeat.stop())
 
+    async def _wait_for_message_and_move(self) -> ClaimedMessage | MessageData | None:
+        interruptible_wait = getattr(self._redis, "_wait_for_message_and_move_interruptible", None)
+        if callable(interruptible_wait):
+            return await interruptible_wait(
+                self.key.pending,
+                self.key.processing,
+                is_interrupted=_DrainInterrupt(lambda: self._draining),
+            )
+        return await self._redis.wait_for_message_and_move(
+            self.key.pending,
+            self.key.processing,
+        )
+
     async def _move_processed_message(
         self,
         destination_queue: str,

{redis_message_queue-8.0.0 → redis_message_queue-8.0.1}/redis_message_queue/redis_message_queue.py

@@ -73,6 +73,30 @@ def _warning_exception_name(exc: BaseException) -> str:
     return type(exc).__name__
 
 
+def _find_non_string_dict_keys(value: object) -> list[object]:
+    non_str_keys: list[object] = []
+    seen: set[int] = set()
+
+    def visit(current: object) -> None:
+        if not isinstance(current, (dict, list, tuple)):
+            return
+        current_id = id(current)
+        if current_id in seen:
+            return
+        seen.add(current_id)
+        if isinstance(current, dict):
+            for key, child in current.items():
+                if not isinstance(key, str):
+                    non_str_keys.append(key)
+                visit(child)
+            return
+        for child in current:
+            visit(child)
+
+    visit(value)
+    return non_str_keys
+
+
 def _validate_heartbeat_interval_seconds(
     heartbeat_interval_seconds: int | float | None,
     visibility_timeout_seconds: int | None,
@@ -235,6 +259,14 @@ class _StopEventInterrupt(BaseGracefulInterruptHandler):
         return self._stop_event.is_set()
 
 
+class _DrainInterrupt(BaseGracefulInterruptHandler):
+    def __init__(self, is_draining: Callable[[], bool]) -> None:
+        self._is_draining = is_draining
+
+    def is_interrupted(self) -> bool:
+        return self._is_draining()
+
+
 class _LeaseHeartbeat:
     def __init__(
         self,
@@ -711,12 +743,10 @@ class RedisMessageQueue:
         """Publish a message.
 
         Dict messages are serialized via ``json.dumps(message, sort_keys=True)``.
-        All top-level dict keys must be strings; non-string keys raise
+        All dict keys must be strings; non-string keys raise
         ``TypeError`` to avoid silent ``json.dumps`` coercion that would
         collapse distinct keys into the same dedup key (e.g. ``{1: "x"}``
-        vs ``{"1": "x"}``). Only top-level keys are validated; nested
-        dicts follow ``json.dumps`` defaults (e.g. nested non-string keys
-        are silently coerced: integer keys become strings).
+        vs ``{"1": "x"}``).
         """
         with self._publish_lock:
             if self._drained.is_set():
@@ -727,7 +757,7 @@ class RedisMessageQueue:
         if not isinstance(message, (str, dict)):
             raise TypeError(f"'message' must be a str or dict, got {type(message).__name__}")
         if isinstance(message, dict):
-            non_str_keys = [k for k in message if not isinstance(k, str)]
+            non_str_keys = _find_non_string_dict_keys(message)
             if non_str_keys:
                 raise TypeError(
                     "'message' dict keys must all be strings; "
@@ -808,10 +838,7 @@ class RedisMessageQueue:
             yield None
             return
         try:
-            claimed_message = self._redis.wait_for_message_and_move(
-                self.key.pending,
-                self.key.processing,
-            )
+            claimed_message = self._wait_for_message_and_move()
         except Exception as exc:
             self._emit_event(
                 "claim",
@@ -979,6 +1006,19 @@ class RedisMessageQueue:
             if lease_heartbeat is not None:
                 lease_heartbeat.stop()
 
+    def _wait_for_message_and_move(self) -> ClaimedMessage | MessageData | None:
+        interruptible_wait = getattr(self._redis, "_wait_for_message_and_move_interruptible", None)
+        if callable(interruptible_wait):
+            return interruptible_wait(
+                self.key.pending,
+                self.key.processing,
+                is_interrupted=_DrainInterrupt(lambda: self._draining),
+            )
+        return self._redis.wait_for_message_and_move(
+            self.key.pending,
+            self.key.processing,
+        )
+
     def _move_processed_message(
         self,
         destination_queue: str,