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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: redis-message-queue
-Version: 8.0.0
+Version: 8.0.2
 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.2}/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.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redis-message-queue"
-version = "8.0.0"
+version = "8.0.2"
 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.2}/redis_message_queue/__init__.py

@@ -5,6 +5,7 @@ from redis_message_queue._exceptions import (
     ConfigurationError,
     GatewayContractError,
     LuaScriptError,
+    MalformedStoredMessageError,
     QueueBackpressureError,
     QueueDrainedError,
     RedisMessageQueueError,
@@ -33,6 +34,7 @@ __all__ = [
     "ConfigurationError",
     "GatewayContractError",
     "LuaScriptError",
+    "MalformedStoredMessageError",
     "QueueBackpressureError",
     "QueueDrainedError",
     "CleanupFailedError",

{redis_message_queue-8.0.0 → redis_message_queue-8.0.2}/redis_message_queue/_config.py

@@ -739,6 +739,26 @@ local function redis_message_queue_decode_claim(cached_claim)
     return nil
 end
 
+local function redis_message_queue_decode_envelope(stored)
+    local prefix = string.char(30) .. 'RMQ1:'
+    if type(stored) ~= 'string' or string.sub(stored, 1, string.len(prefix)) ~= prefix then
+        return nil
+    end
+    local ok, envelope = pcall(cjson.decode, string.sub(stored, string.len(prefix) + 1))
+    if ok and type(envelope) == 'table' and type(envelope['id']) == 'string' then
+        return envelope
+    end
+    return nil
+end
+
+local function redis_message_queue_message_id(stored)
+    local envelope = redis_message_queue_decode_envelope(stored)
+    if envelope then
+        return envelope['id']
+    end
+    return ''
+end
+
 local time = redis.call('TIME')
 local now_ms = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000)
 
@@ -779,6 +799,7 @@ end
 -- With a single consumer polling at default interval, 1000 expired leases drain in ~2.5s.
 local expired = redis.call('ZRANGEBYSCORE', KEYS[3], '-inf', now_ms, 'LIMIT', 0, 100)
 local to_requeue = {}
+local reclaimed_events = {}
 for i = #expired, 1, -1 do
     local expired_lease_token = redis.call('HGET', KEYS[4], expired[i])
     redis.call('ZREM', KEYS[3], expired[i])
@@ -800,13 +821,14 @@ for i = #expired, 1, -1 do
     end
     if redis.call('LREM', KEYS[2], 1, expired[i]) == 1 then
         table.insert(to_requeue, expired[i])
+        local delivery_count = redis.call('HGET', KEYS[6], expired[i])
+        table.insert(reclaimed_events, {redis_message_queue_message_id(expired[i]), tostring(delivery_count or '0')})
     end
 end
 if #to_requeue > 0 then
     redis.call('RPUSH', KEYS[1], unpack(to_requeue))
 end
-local reclaimed_count = #to_requeue
-local dead_lettered_count = 0
+local dead_lettered_events = {}
 
 local function store_claim_and_return(stored)
     -- pcall guards against OOM mid-write: compensate by returning message to pending
@@ -819,7 +841,7 @@ local function store_claim_and_return(stored)
         redis.call('HSET', KEYS[9], lease_token, KEYS[8])
         redis.call('HSET', KEYS[10], ARGV[4], claim_payload)
         redis.call('HSET', KEYS[11], lease_token, ARGV[4])
-        return {stored, lease_token, tostring(reclaimed_count), tostring(dead_lettered_count)}
+        return {stored, lease_token, reclaimed_events, dead_lettered_events}
     end)
     if not ok then
         redis.call('LREM', KEYS[2], 1, stored)
@@ -835,36 +857,28 @@ while claim_attempts < 100 do
 
     local stored = redis.call('LMOVE', KEYS[1], KEYS[2], 'RIGHT', 'LEFT')
     if not stored then
-        return {'', '', tostring(reclaimed_count), tostring(dead_lettered_count)}
+        return {'', '', reclaimed_events, dead_lettered_events}
     end
 
-    if max_delivery_count > 0 then
-        local count = redis.call('HINCRBY', KEYS[6], stored, 1)
-        if count > max_delivery_count then
-            redis.call('LREM', KEYS[2], 1, stored)
-            redis.call('HDEL', KEYS[6], stored)
-            -- Strip envelope to store raw payload in DLQ, consistent with completed/failed queues.
-            -- The per-delivery UUID in the envelope is lost; see README dead-letter notes.
-            local dead_letter_value = stored
-            local prefix = string.char(30) .. 'RMQ1:'
-            if string.sub(stored, 1, string.len(prefix)) == prefix then
-                local ok, envelope = pcall(cjson.decode, string.sub(stored, string.len(prefix) + 1))
-                if ok and type(envelope) == 'table' and type(envelope['id']) == 'string'
-                        and type(envelope['payload']) == 'string' then
-                    dead_letter_value = envelope['payload']
-                end
-            end
-            redis.call('LPUSH', KEYS[7], dead_letter_value)
-            dead_lettered_count = dead_lettered_count + 1
-        else
-            return store_claim_and_return(stored)
+    local count = redis.call('HINCRBY', KEYS[6], stored, 1)
+    if max_delivery_count > 0 and count > max_delivery_count then
+        redis.call('LREM', KEYS[2], 1, stored)
+        redis.call('HDEL', KEYS[6], stored)
+        -- Strip envelope to store raw payload in DLQ, consistent with completed/failed queues.
+        -- The per-delivery UUID in the envelope is lost; see README dead-letter notes.
+        local dead_letter_value = stored
+        local envelope = redis_message_queue_decode_envelope(stored)
+        if envelope and type(envelope['payload']) == 'string' then
+            dead_letter_value = envelope['payload']
         end
+        redis.call('LPUSH', KEYS[7], dead_letter_value)
+        table.insert(dead_lettered_events, {redis_message_queue_message_id(stored), tostring(count)})
     else
         return store_claim_and_return(stored)
     end
 end
 
-return {'', '', tostring(reclaimed_count), tostring(dead_lettered_count)}
+return {'', '', reclaimed_events, dead_lettered_events}
 """
 )
 
@@ -1037,6 +1051,46 @@ return removed
 """
 )
 
+CLEANUP_DRAINED_LEASE_TOKEN_COUNTER_LUA_SCRIPT = (
+    _LUA_KEY_TYPE_GUARD
+    + """
+local err = redis_message_queue_require_type(KEYS[1], 'list')
+if err then
+    return err
+end
+
+local err = redis_message_queue_require_type(KEYS[2], 'zset')
+if err then
+    return err
+end
+
+local err = redis_message_queue_require_type(KEYS[3], 'hash')
+if err then
+    return err
+end
+
+local err = redis_message_queue_require_type(KEYS[4], 'hash')
+if err then
+    return err
+end
+
+local err = redis_message_queue_require_type(KEYS[5], 'string')
+if err then
+    return err
+end
+
+if redis.call('LLEN', KEYS[1]) == 0
+    and redis.call('ZCARD', KEYS[2]) == 0
+    and redis.call('HLEN', KEYS[3]) == 0
+    and redis.call('HLEN', KEYS[4]) == 0 then
+    redis.call('DEL', KEYS[5])
+    return 1
+end
+
+return 0
+"""
+)
+
 RENEW_MESSAGE_LEASE_LUA_SCRIPT = (
     _LUA_KEY_TYPE_GUARD
     + """

{redis_message_queue-8.0.0 → redis_message_queue-8.0.2}/redis_message_queue/_event.py

@@ -52,6 +52,10 @@ class QueueEvent:
     """a diagnostic hash of the lease token when visibility timeout is enabled"""
     destination_queue: str | None = None
     """the queue a message was moved to, when applicable"""
+    delivery_count: int | None = None
+    """the number of delivery attempts recorded for this message, when applicable"""
+    max_delivery_count: int | None = None
+    """the configured delivery-attempt threshold, when applicable"""
     exception_type: str | None = None
     """
     type name of the raised exception for metrics labels (e.g., 'TimeoutError');

{redis_message_queue-8.0.0 → redis_message_queue-8.0.2}/redis_message_queue/_exceptions.py

@@ -21,6 +21,10 @@ class CleanupFailedError(RedisMessageQueueError):
     """Cleanup after handler completion failed."""
 
 
+class MalformedStoredMessageError(RedisMessageQueueError):
+    """Stored value starts with the RMQ envelope prefix but is not a valid envelope."""
+
+
 class QueueBackpressureError(RedisMessageQueueError):
     """Publish rejected because the pending queue is at its configured limit."""
 

{redis_message_queue-8.0.0 → redis_message_queue-8.0.2}/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
@@ -15,6 +16,7 @@ from redis_message_queue._config import (
     ADD_MESSAGE_LUA_SCRIPT,
     CLAIM_MESSAGE_LUA_SCRIPT,
     CLAIM_MESSAGE_WITH_VISIBILITY_TIMEOUT_LUA_SCRIPT,
+    CLEANUP_DRAINED_LEASE_TOKEN_COUNTER_LUA_SCRIPT,
     DEFAULT_MESSAGE_DEDUPLICATION_LOG_TTL,
     DEFAULT_MESSAGE_WAIT_INTERVAL_SECONDS,
     DEFAULT_PENDING_OVERLOAD_BLOCK_TIMEOUT_SECONDS,
@@ -54,6 +56,8 @@ from redis_message_queue.interrupt_handler._interface import (
 
 logger = logging.getLogger(__name__)
 _TClaim = TypeVar("_TClaim", bound=ClaimedMessage | MessageData)
+_TRedisCall = TypeVar("_TRedisCall")
+_MessageAttemptEvent = tuple[str | None, int]
 
 _LEASE_DEADLINES_SUFFIX = ":lease_deadlines"
 _LEASE_TOKENS_SUFFIX = ":lease_tokens"
@@ -71,6 +75,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")
@@ -80,6 +127,26 @@ def _coerce_lua_count(value: object) -> int:
         return 0
 
 
+def _decode_lua_text(value: object) -> str | None:
+    if isinstance(value, bytes):
+        value = value.decode("utf-8")
+    if isinstance(value, str) and value:
+        return value
+    return None
+
+
+def _coerce_lua_message_attempts(value: object) -> list[_MessageAttemptEvent]:
+    if not isinstance(value, list | tuple):
+        return []
+
+    attempts: list[_MessageAttemptEvent] = []
+    for item in value:
+        if not isinstance(item, list | tuple) or len(item) < 2:
+            continue
+        attempts.append((_decode_lua_text(item[0]), _coerce_lua_count(item[1])))
+    return attempts
+
+
 def _pending_overload_max_backoff_seconds(block_timeout_seconds: float) -> float:
     return min(_PENDING_OVERLOAD_MAX_BACKOFF_SECONDS, block_timeout_seconds / 10)
 
@@ -214,8 +281,11 @@ class RedisGateway(AbstractRedisGateway):
         operation: EventOperation,
         outcome: EventOutcome,
         *,
+        message_id: str | None = None,
         claim_id: str | None = None,
         destination_queue: str | None = None,
+        delivery_count: int | None = None,
+        max_delivery_count: int | None = None,
         exception_type: str | None = None,
         error: BaseException | None = None,
     ) -> None:
@@ -224,8 +294,11 @@ class RedisGateway(AbstractRedisGateway):
         self._event_emitter(
             operation,
             outcome,
+            message_id=message_id,
             claim_id=claim_id,
             destination_queue=destination_queue,
+            delivery_count=delivery_count,
+            max_delivery_count=max_delivery_count,
             exception_type=exception_type,
             error=error,
         )
@@ -233,12 +306,20 @@ class RedisGateway(AbstractRedisGateway):
     def _emit_repeated_event(
         self,
         operation: EventOperation,
-        count: int,
+        attempts: list[_MessageAttemptEvent],
         *,
         destination_queue: str | None = None,
+        max_delivery_count: int | None = None,
     ) -> None:
-        for _ in range(max(count, 0)):
-            self._emit_event(operation, "success", destination_queue=destination_queue)
+        for message_id, delivery_count in attempts:
+            self._emit_event(
+                operation,
+                "success",
+                message_id=message_id,
+                destination_queue=destination_queue,
+                delivery_count=delivery_count,
+                max_delivery_count=max_delivery_count,
+            )
 
     def _eval(self, *args: object) -> object:
         try:
@@ -525,6 +606,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 +634,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 +657,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 +689,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 +722,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 +757,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 +800,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 +823,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(
@@ -760,10 +875,15 @@ class RedisGateway(AbstractRedisGateway):
             return None
 
         stored_message, lease_token = result[0], result[1]
-        reclaimed_count = _coerce_lua_count(result[2]) if len(result) > 2 else 0
-        dead_lettered_count = _coerce_lua_count(result[3]) if len(result) > 3 else 0
-        self._emit_repeated_event("claim_reclaim", reclaimed_count)
-        self._emit_repeated_event("dlq", dead_lettered_count, destination_queue=self._dead_letter_queue)
+        reclaimed_attempts = _coerce_lua_message_attempts(result[2]) if len(result) > 2 else []
+        dead_lettered_attempts = _coerce_lua_message_attempts(result[3]) if len(result) > 3 else []
+        self._emit_repeated_event("claim_reclaim", reclaimed_attempts)
+        self._emit_repeated_event(
+            "dlq",
+            dead_lettered_attempts,
+            destination_queue=self._dead_letter_queue,
+            max_delivery_count=self._max_delivery_count,
+        )
         if stored_message in ("", b"") and lease_token in ("", b""):
             return None
         if isinstance(lease_token, bytes):
@@ -833,16 +953,54 @@ 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 _cleanup_drained_lease_token_counter(self, processing_queue: str) -> bool:
+        if self._message_visibility_timeout_seconds is None:
+            return False
+        return bool(
+            _coerce_lua_count(
+                self._eval(
+                    CLEANUP_DRAINED_LEASE_TOKEN_COUNTER_LUA_SCRIPT,
+                    5,
+                    processing_queue,
+                    self._lease_deadlines_key(processing_queue),
+                    self._lease_tokens_key(processing_queue),
+                    self._delivery_counts_key(processing_queue),
+                    self._lease_token_counter_key(processing_queue),
+                )
+            )
+        )
+
+    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 +1064,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 +1113,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 +1122,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 +1130,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 +1198,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