kiba-core 0.5.3.dev85__tar.gz → 0.5.3.dev86__tar.gz

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/CHANGELOG.md

@@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 
 ### Added
+- [MAJOR] Added abstract `MessageQueue.release_message` to hand a claimed message back to be claimed again straight away without using an attempt (`SqlMessageQueue` and `CosmosMessageQueue` keep `postCount`; `SqsMessageQueue` and `AqsMessageQueue` make the message visible again, though SQS still counts the next receive towards its redrive policy)
+- [MAJOR] Added abstract `MessageQueue.delete_expired_messages`: `SqlMessageQueue` deletes `succeeded` and `deduplicated` messages older than `retentionSeconds` (the other queues remove messages on completion and return 0). `SqlMessageQueue.get_message_count` no longer deletes them and its `shouldSkipRemovingNonRetained` parameter is removed
+- [MINOR] Updated `MessageQueueProcessor` to shut down gracefully: `run`/`run_batches` handle SIGTERM and SIGINT (a second signal falls back to the default handler), as does `request_shutdown()`, by finishing the message being processed, releasing claimed messages not yet started and stopping instead of sleeping. A processed message whose task is cancelled is released rather than left until its lease expires
+- [MINOR] Added `deleteExpiredMessagesIntervalSeconds` (default 600) to `MessageQueueProcessor`, which calls `delete_expired_messages` at most this often while running, logging rather than raising on failure; added `shouldProcessInParallel` to `run_batches`
 - [MINOR] Added optional `proxyUrl` to `Requester` to send all its requests through an HTTP(S) proxy, e.g. to reach an API that allowlists a fixed egress IP
 - [MINOR] Added optional `timeoutSeconds` (default 10) to `RestEthClient` to set the timeout of each request, e.g. a shorter one for an archive node so a degraded node fails fast
 - [MINOR] Added `UNDEFINED`, `Undefined` and `UpdateValue[T]` to `core.util.typing_util` for partial updates: `UpdateValue[T]` is `T` or the `UNDEFINED` sentinel, so an update can tell "not provided" apart from an explicit `None`. Use it as a pydantic field default (`title: UpdateValue[str | None] = UNDEFINED`); request input never validates to `UNDEFINED` (a JSON `"UNDEFINED"` stays a string), and `x is not UNDEFINED` narrows `x` to `T` for type checkers

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev85
+Version: 0.5.3.dev86
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/aqs.py

@@ -113,6 +113,14 @@ class AqsMessageQueue(MessageQueue[AqsMessage]):
         message.popReceipt = updatedMessage.pop_receipt
         return True
 
+    async def release_message(self, message: AqsMessage) -> None:
+        # NOTE(krishan711): a message whose lease was already lost is left alone
+        await self.extend_message_lease(message=message, expectedProcessingSeconds=0)
+
+    async def delete_expired_messages(self) -> int:
+        # NOTE(krishan711): azure deletes messages when they are completed and expires the rest itself
+        return 0
+
     async def get_message_count(self) -> int:
         if not self._aqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to count messages')

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/cosmos.py

@@ -258,6 +258,27 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
         message.etag = str(updatedItem['_etag'])
         return True
 
+    async def release_message(self, message: CosmosMessage) -> None:
+        try:
+            await self.container.patch_item(
+                item=message.id,
+                partition_key=self.queueName,
+                patch_operations=[
+                    {'op': 'set', 'path': '/leaseId', 'value': None},
+                    {'op': 'set', 'path': '/visibleDate', 'value': time.time()},
+                ],
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except cosmos_exceptions.CosmosHttpResponseError as exception:
+            # NOTE(krishan711): 412 means another worker now owns the message and 404 that it was already completed
+            if exception.status_code not in {404, 412}:
+                raise
+
+    async def delete_expired_messages(self) -> int:
+        # NOTE(krishan711): cosmos keeps no message history, completed messages are removed straight away
+        return 0
+
     async def get_message_count(self) -> int:
         counts = self.container.query_items(
             query='SELECT COUNT(1) AS messageCount FROM c WHERE c.itemType = @itemType AND c.queueName = @queueName AND c.visibleDate <= @now',

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/message_queue.py

@@ -47,6 +47,15 @@ class MessageQueue[MessageType: Message](ABC):
     async def extend_message_lease(self, message: MessageType, expectedProcessingSeconds: int) -> bool:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def release_message(self, message: MessageType) -> None:
+        # NOTE(krishan711): hands a claimed message back to be claimed again straight away without using an attempt, e.g. when the worker shuts down
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def delete_expired_messages(self) -> int:
+        raise NotImplementedError
+
     @abc.abstractmethod
     async def get_message_count(self) -> int:
         raise NotImplementedError

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/message_queue_processor.py

@@ -1,6 +1,7 @@
 import abc
 import asyncio
 import contextlib
+import signal
 import sys
 import time
 import urllib.parse as urlparse
@@ -18,6 +19,7 @@ from core.util.value_holder import RequestIdHolder
 
 LOCKED_MAX_RETRY_COUNT = 3
 LOCKED_RETRY_DELAY_SECONDS = 30
+SHUTDOWN_SIGNALS = (signal.SIGTERM, signal.SIGINT)
 
 
 class MessageProcessor(ABC):
@@ -40,14 +42,30 @@ class MessageLeaseLostException(InternalServerErrorException):
 
 
 class MessageQueueProcessor[MessageType: Message]:
-    def __init__(self, queue: MessageQueue[MessageType], messageProcessor: MessageProcessor, notificationClients: list[NotificationClient], requestIdHolder: RequestIdHolder | None = None, pollLogIntervalSeconds: float = 600, shouldNotifyRetriedFailures: bool = True) -> None:
+    def __init__(
+        self,
+        queue: MessageQueue[MessageType],
+        messageProcessor: MessageProcessor,
+        notificationClients: list[NotificationClient],
+        requestIdHolder: RequestIdHolder | None = None,
+        pollLogIntervalSeconds: float = 600,
+        shouldNotifyRetriedFailures: bool = True,
+        deleteExpiredMessagesIntervalSeconds: float = 600,
+    ) -> None:
         self.queue = queue
         self.messageProcessor = messageProcessor
         self.notificationClients = notificationClients
         self.requestIdHolder = requestIdHolder
         self.pollLogIntervalSeconds = pollLogIntervalSeconds
         self.shouldNotifyRetriedFailures = shouldNotifyRetriedFailures
+        self.deleteExpiredMessagesIntervalSeconds = deleteExpiredMessagesIntervalSeconds
         self._lastPollLogTime: float | None = None
+        self._lastExpiredDeletionTime: float | None = None
+        self._shutdownEvent = asyncio.Event()
+
+    def request_shutdown(self) -> None:
+        # NOTE(krishan711): stops claiming new messages; the message being processed finishes and claimed messages not yet started are released
+        self._shutdownEvent.set()
 
     async def _handle_failure(self, message: MessageType, exception: Exception, requestId: str) -> int:
         logging.error('Caught exception whilst processing message:')
@@ -95,6 +113,7 @@ class MessageQueueProcessor[MessageType: Message]:
         # Cancellation (and other non-Exception errors) must still propagate, so those are never replaced.
         stopKeepAliveEvent = asyncio.Event()
         keepAliveTask = asyncio.create_task(self._keep_message_alive(message=message, expectedProcessingSeconds=expectedProcessingSeconds, stopEvent=stopKeepAliveEvent))
+        isLeaseLost = False
         try:
             try:
                 await self.messageProcessor.process_message(message=message)
@@ -105,6 +124,10 @@ class MessageQueueProcessor[MessageType: Message]:
                 if isLeaseLost and (inFlightException is None or isinstance(inFlightException, Exception)):
                     raise MessageLeaseLostException
             await self.queue.complete_message(message=message)
+        except asyncio.CancelledError:
+            if not isLeaseLost:
+                await self._release_messages(messages=[message])
+            raise
         except MessageNeedsReprocessingException as exception:
             logging.info(msg=f'Scheduling reprocessing for message:{message.command} due to: {exception.originalException!s}')
             await self.queue.retry_message(message=message, delaySeconds=((message.postCount or 0) * exception.delaySeconds))
@@ -122,16 +145,44 @@ class MessageQueueProcessor[MessageType: Message]:
         if self.requestIdHolder:
             self.requestIdHolder.set_value(value=None)
 
+    async def _release_messages(self, messages: list[MessageType]) -> None:
+        for message in messages:
+            try:
+                await self.queue.release_message(message=message)
+            except Exception as exception:  # noqa: BLE001
+                # NOTE(krishan711): a message that cannot be released is picked up again once its lease expires
+                logging.error(f'Failed to release message {message.command}:')
+                logging.exception(exception)
+
+    async def _delete_expired_messages_if_due(self) -> None:
+        now = time.monotonic()
+        if self._lastExpiredDeletionTime is not None and now - self._lastExpiredDeletionTime < self.deleteExpiredMessagesIntervalSeconds:
+            return
+        self._lastExpiredDeletionTime = now
+        try:
+            deletedCount = await self.queue.delete_expired_messages()
+        except Exception as exception:  # noqa: BLE001
+            logging.error('Failed to delete expired messages:')
+            logging.exception(exception)
+            return
+        if deletedCount > 0:
+            logging.info(f'Deleted {deletedCount} expired messages')
+
     async def execute_batch(self, batchSize: int, expectedProcessingSeconds: int = 300, longPollSeconds: int = 20, shouldProcessInParallel: bool = False) -> int:
         now = time.monotonic()
         if self._lastPollLogTime is None or now - self._lastPollLogTime >= self.pollLogIntervalSeconds:
             logging.info('Retrieving messages...')
             self._lastPollLogTime = now
         messages = await self.queue.get_messages(expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds, limit=batchSize)
-        if shouldProcessInParallel:
+        if self._shutdownEvent.is_set():
+            await self._release_messages(messages=messages)
+        elif shouldProcessInParallel:
             await asyncio.gather(*[self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds) for message in messages])
         else:
-            for message in messages:
+            for index, message in enumerate(messages):
+                if self._shutdownEvent.is_set():
+                    await self._release_messages(messages=messages[index:])
+                    break
                 await self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds)
         return len(messages)
 
@@ -139,14 +190,41 @@ class MessageQueueProcessor[MessageType: Message]:
         processedMessageCount = await self.execute_batch(batchSize=1, expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds)
         return processedMessageCount > 0
 
-    async def run_batches(self, batchSize: int, expectedProcessingSeconds: int = 300, longPollSeconds: int = 20, sleepTime: int = 30, totalMessageLimit: int | None = None) -> int:
+    def _on_shutdown_signal(self, signalNumber: int) -> None:
+        logging.info(f'Received {signal.Signals(signalNumber).name}, finishing the current message before shutting down')
+        self.request_shutdown()
+        # NOTE(krishan711): a second signal falls back to the default handler so an impatient interrupt still stops the process
+        loop = asyncio.get_running_loop()
+        for shutdownSignal in SHUTDOWN_SIGNALS:
+            loop.remove_signal_handler(shutdownSignal)
+
+    async def run_batches(
+        self,
+        batchSize: int,
+        expectedProcessingSeconds: int = 300,
+        longPollSeconds: int = 20,
+        sleepTime: int = 30,
+        totalMessageLimit: int | None = None,
+        shouldProcessInParallel: bool = False,
+    ) -> int:
+        loop = asyncio.get_running_loop()
+        for shutdownSignal in SHUTDOWN_SIGNALS:
+            loop.add_signal_handler(shutdownSignal, self._on_shutdown_signal, shutdownSignal)
         processedMessageCount = 0
-        while totalMessageLimit is None or processedMessageCount < totalMessageLimit:
-            innerProcessedMessageCount = await self.execute_batch(expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds, batchSize=batchSize)
-            if innerProcessedMessageCount == 0:
-                logging.info('No message received.. sleeping')
-                await asyncio.sleep(sleepTime)
-            processedMessageCount += innerProcessedMessageCount
+        try:
+            while not self._shutdownEvent.is_set() and (totalMessageLimit is None or processedMessageCount < totalMessageLimit):
+                await self._delete_expired_messages_if_due()
+                innerProcessedMessageCount = await self.execute_batch(expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds, batchSize=batchSize, shouldProcessInParallel=shouldProcessInParallel)
+                if innerProcessedMessageCount == 0:
+                    logging.info('No message received.. sleeping')
+                    with contextlib.suppress(TimeoutError):
+                        await asyncio.wait_for(self._shutdownEvent.wait(), timeout=sleepTime)
+                processedMessageCount += innerProcessedMessageCount
+        finally:
+            for shutdownSignal in SHUTDOWN_SIGNALS:
+                loop.remove_signal_handler(shutdownSignal)
+        if self._shutdownEvent.is_set():
+            logging.info('Shutdown requested, stopped processing messages')
         return processedMessageCount
 
     async def run(self, expectedProcessingSeconds: int = 300, longPollSeconds: int = 20, sleepTime: int = 30, totalMessageLimit: int | None = None) -> bool:

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/sql.py

@@ -211,7 +211,7 @@ class SqlMessageQueue(MessageQueue[SqlMessage]):
             return result.first() is not None
 
     async def retry_message(self, message: SqlMessage, delaySeconds: int = 0) -> None:
-        await self._reschedule_message(message=message, delaySeconds=delaySeconds, extraValues={})
+        await self._reschedule_message(message=message, delaySeconds=delaySeconds, shouldCountAttempt=True)
 
     async def fail_message(self, message: SqlMessage, errorMessage: str | None) -> bool:
         postCount = message.postCount or 0
@@ -220,21 +220,39 @@ class SqlMessageQueue(MessageQueue[SqlMessage]):
                 query = self._owned_update(message=message).values(status=MESSAGE_STATUS_FAILED, lockToken=None, completedDate=self.database.now(), lastError=errorMessage)
                 await self.database.execute(query=query, connection=connection)
             return False
-        await self._reschedule_message(message=message, delaySeconds=self.failureRetryDelaySeconds * postCount, extraValues={'lastError': errorMessage})
+        await self._reschedule_message(message=message, delaySeconds=self.failureRetryDelaySeconds * postCount, shouldCountAttempt=True, extraValues={'lastError': errorMessage})
         return True
 
-    async def _reschedule_message(self, message: SqlMessage, delaySeconds: float, extraValues: dict[str, Any]) -> None:  # type: ignore[explicit-any]
+    async def release_message(self, message: SqlMessage) -> None:
+        await self._reschedule_message(message=message, delaySeconds=0, shouldCountAttempt=False)
+
+    async def _reschedule_message(self, message: SqlMessage, delaySeconds: float, shouldCountAttempt: bool, extraValues: dict[str, Any] | None = None) -> None:  # type: ignore[explicit-any]
+        values = {'status': MESSAGE_STATUS_PENDING, 'lockToken': None, 'visibleDate': self.database.now(seconds=delaySeconds), **(extraValues or {})}
+        if shouldCountAttempt:
+            values |= {'postCount': self.table.c.postCount + 1, 'postDate': self.database.now()}
         async with self.database.create_transaction() as connection:
             try:
                 async with connection.begin_nested():
-                    query = self._owned_update(message=message).values(status=MESSAGE_STATUS_PENDING, lockToken=None, visibleDate=self.database.now(seconds=delaySeconds), postCount=self.table.c.postCount + 1, postDate=self.database.now(), **extraValues)
+                    query = self._owned_update(message=message).values(values)
                     await self.database.execute(query=query, connection=connection)
             except sqlalchemy.exc.IntegrityError:
                 # NOTE(krishan711): an identical message was enqueued while this one ran, so it will do the work instead
                 query = self._owned_update(message=message).values(status=MESSAGE_STATUS_DEDUPLICATED, lockToken=None, completedDate=self.database.now())
                 await self.database.execute(query=query, connection=connection)
 
-    async def get_message_count(self, shouldSkipRemovingNonRetained: bool = False) -> int:
+    async def delete_expired_messages(self) -> int:
+        async with self.database.create_transaction() as connection:
+            query = (
+                sqlalchemy.delete(self.table)
+                .where(self.table.c.queueName == self.queueName)
+                .where(self.table.c.status.in_([MESSAGE_STATUS_SUCCEEDED, MESSAGE_STATUS_DEDUPLICATED]))
+                .where(self.table.c.completedDate < self.database.now(seconds=-self.retentionSeconds))
+                .returning(self.table.c.id)
+            )
+            result = await self.database.execute(query=query, connection=connection)
+            return len(result.all())
+
+    async def get_message_count(self) -> int:
         countQuery = (
             sqlalchemy.select(sqlalchemy.func.count())
             .select_from(self.table)
@@ -243,14 +261,6 @@ class SqlMessageQueue(MessageQueue[SqlMessage]):
             .where(self.table.c.visibleDate <= self.database.now())
         )
         async with self.database.create_transaction() as connection:
-            if not shouldSkipRemovingNonRetained:
-                deleteQuery = (
-                    sqlalchemy.delete(self.table)
-                    .where(self.table.c.queueName == self.queueName)
-                    .where(self.table.c.status.in_([MESSAGE_STATUS_SUCCEEDED, MESSAGE_STATUS_DEDUPLICATED]))
-                    .where(self.table.c.completedDate < self.database.now(seconds=-self.retentionSeconds))
-                )
-                await self.database.execute(query=deleteQuery, connection=connection)
             result = await self.database.execute(query=countQuery, connection=connection)
         return int(result.scalar_one())
 

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/core/queues/sqs.py

@@ -131,6 +131,14 @@ class SqsMessageQueue(MessageQueue[SqsMessage]):
             raise
         return True
 
+    async def release_message(self, message: SqsMessage) -> None:
+        # NOTE(krishan711): a message whose lease was already lost is left alone; SQS still counts the next receive towards the queue's redrive policy
+        await self.extend_message_lease(message=message, expectedProcessingSeconds=0)
+
+    async def delete_expired_messages(self) -> int:
+        # NOTE(krishan711): SQS deletes messages when they are completed and expires the rest itself
+        return 0
+
     async def get_message_count(self) -> int:
         if not self._sqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to count messages')

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/kiba_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev85
+Version: 0.5.3.dev86
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kiba-core"
-version = "0.5.3.dev85"
+version = "0.5.3.dev86"
 description = "Kiba Labs' python utilities"
 readme = "README.md"
 requires-python = ">=3.12, <4"

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/tests/queues/test_aqs_queue.py

@@ -43,6 +43,13 @@ async def _claim(queue: AqsMessageQueue, expectedProcessingSeconds: int = 300) -
     return message
 
 
+async def test_released_message_can_be_received_again_straight_away(workers: tuple[AqsMessageQueue, AqsMessageQueue]):
+    workerA, workerB = workers
+    await workerA.send_message(message=_message())
+    await workerA.release_message(message=await _claim(queue=workerA))
+    await _claim(queue=workerB)
+
+
 async def test_extended_message_stays_invisible_and_can_still_be_completed(workers: tuple[AqsMessageQueue, AqsMessageQueue]):
     workerA, workerB = workers
     await workerA.send_message(message=_message())

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/tests/queues/test_cosmos_queue.py

@@ -85,6 +85,15 @@ async def test_completed_message_is_removed_and_can_be_sent_again(workers: tuple
     assert await workerA.get_message_count() == 1
 
 
+async def test_release_hands_the_message_back_without_using_an_attempt(workers: tuple[CosmosMessageQueue, CosmosMessageQueue]):
+    workerA, workerB = workers
+    await workerA.send_message(message=_message())
+    message = await _claim(queue=workerA)
+    await workerA.release_message(message=message)
+    releasedMessage = await _claim(queue=workerB)
+    assert (releasedMessage.id, releasedMessage.postCount) == (message.id, 1)
+
+
 async def test_retry_reschedules_the_same_message_with_a_delay(workers: tuple[CosmosMessageQueue, CosmosMessageQueue]):
     workerA, _ = workers
     await workerA.send_message(message=_message())

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/tests/queues/test_sql_queue.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import signal
 import typing
 from collections.abc import AsyncIterator
 from typing import Any
@@ -262,6 +263,66 @@ async def test_processor_cancellation_propagates_even_when_the_lease_was_lost(wo
         await executeTask
 
 
+async def test_processor_releases_a_cancelled_message_without_using_an_attempt(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, _ = workers
+
+    class SlowProcessor(MessageProcessor):
+        async def process_message(self, message: Message) -> None:
+            await asyncio.sleep(10)
+
+    processor = MessageQueueProcessor(queue=workerA, messageProcessor=SlowProcessor(), notificationClients=[])
+    await workerA.send_message(message=_message())
+    executeTask = asyncio.create_task(processor.execute(longPollSeconds=0))
+    await asyncio.sleep(0.3)
+    executeTask.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await executeTask
+    row = (await _rows(queue=workerA))[0]
+    assert (row['status'], row['post_count'], row['lock_token']) == ('pending', 1, None)
+    assert (await _claim(queue=workerA)).postCount == 1
+
+
+async def test_processor_finishes_the_current_message_and_releases_the_rest_on_shutdown(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, _ = workers
+    processedUserIds: list[str] = []
+
+    class ShutdownRequestingProcessor(MessageProcessor):
+        async def process_message(self, message: Message) -> None:
+            processedUserIds.append(str(message.content['userId']))
+            processor.request_shutdown()
+
+    processor = MessageQueueProcessor(queue=workerA, messageProcessor=ShutdownRequestingProcessor(), notificationClients=[])
+    await workerA.send_messages(messages=[_message(deduplicationId=None, userId=str(index)) for index in range(3)])
+    assert await processor.run_batches(batchSize=3, longPollSeconds=0) == 3  # noqa: PLR2004
+    assert processedUserIds == ['0']
+    assert [(row['status'], row['post_count']) for row in await _rows(queue=workerA)] == [('succeeded', 1), ('pending', 1), ('pending', 1)]
+
+
+async def test_processor_stops_promptly_when_signalled_while_idle(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, _ = workers
+    processor = MessageQueueProcessor(queue=workerA, messageProcessor=typing.cast(MessageProcessor, None), notificationClients=[])
+    runTask = asyncio.create_task(processor.run(longPollSeconds=0, sleepTime=30))
+    await asyncio.sleep(0.3)
+    os.kill(os.getpid(), signal.SIGTERM)
+    assert await asyncio.wait_for(runTask, timeout=5) is False
+
+
+async def test_processor_deletes_expired_messages_while_running(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, _ = workers
+    await workerA.send_message(message=_message())
+    message = await _claim(queue=workerA)
+    await workerA.complete_message(message=message)
+    async with workerA.database.create_transaction() as connection:
+        await workerA.database.execute(query=sqlalchemy.update(workerA.table).values(completedDate=date_util.datetime_from_now(days=-8)), connection=connection)
+    processor = MessageQueueProcessor(queue=workerA, messageProcessor=typing.cast(MessageProcessor, None), notificationClients=[])
+    runTask = asyncio.create_task(processor.run(longPollSeconds=0, sleepTime=30))
+    await asyncio.sleep(0.3)
+    processor.request_shutdown()
+    await asyncio.wait_for(runTask, timeout=5)
+    assert await _rows(queue=workerA) == []
+
+
+
 async def test_expired_lease_with_no_attempts_left_is_failed(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
     workerA, _ = workers
     await workerA.send_message(message=_message())
@@ -315,7 +376,27 @@ async def test_retry_is_folded_into_an_identical_pending_message(workers: tuple[
     assert [row['status'] for row in await _rows(queue=workerA)] == ['deduplicated', 'pending']
 
 
-async def test_message_count_removes_only_old_finished_messages_unless_skipped(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+async def test_release_hands_the_message_back_without_using_an_attempt(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, workerB = workers
+    await workerA.send_message(message=_message())
+    message = await _claim(queue=workerA)
+    await workerA.release_message(message=message)
+    releasedMessage = await _claim(queue=workerB)
+    assert (releasedMessage.id, releasedMessage.postCount) == (message.id, 1)
+
+
+async def test_release_after_losing_the_lease_leaves_the_new_holder_untouched(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
+    workerA, workerB = workers
+    await workerA.send_message(message=_message())
+    staleMessage = await _claim(queue=workerA)
+    await _make_visible(queue=workerA, messageId=staleMessage.id)
+    await _claim(queue=workerB)
+    await workerA.release_message(message=staleMessage)
+    row = await _row(queue=workerA, messageId=staleMessage.id)
+    assert (row['status'], row['owner']) == ('running', 'worker-b')
+
+
+async def test_delete_expired_messages_removes_only_old_finished_messages(workers: tuple[SqlMessageQueue, SqlMessageQueue]):
     workerA, _ = workers
     await workerA.send_messages(messages=[_message(deduplicationId=None, userId=str(index)) for index in range(4)])
     succeededMessage, failedMessage, recentMessage = await workerA.get_messages(limit=3)
@@ -324,9 +405,9 @@ async def test_message_count_removes_only_old_finished_messages_unless_skipped(w
     await workerA.complete_message(message=recentMessage)
     async with workerA.database.create_transaction() as connection:
         await workerA.database.execute(query=sqlalchemy.update(workerA.table).where(workerA.table.c.id.in_([succeededMessage.id, failedMessage.id])).values(completedDate=date_util.datetime_from_now(days=-8)), connection=connection)
-    assert await workerA.get_message_count(shouldSkipRemovingNonRetained=True) == 1
-    assert len(await _rows(queue=workerA)) == 4  # noqa: PLR2004
     assert await workerA.get_message_count() == 1
+    assert len(await _rows(queue=workerA)) == 4  # noqa: PLR2004
+    assert await workerA.delete_expired_messages() == 1
     assert [row['status'] for row in await _rows(queue=workerA)] == ['failed', 'succeeded', 'pending']
 
 

{kiba_core-0.5.3.dev85 → kiba_core-0.5.3.dev86}/tests/queues/test_sqs_queue.py

@@ -38,6 +38,13 @@ async def _claim(queue: SqsMessageQueue, expectedProcessingSeconds: int = 300) -
     return message
 
 
+async def test_released_message_can_be_received_again_straight_away(workers: tuple[SqsMessageQueue, SqsMessageQueue]):
+    workerA, workerB = workers
+    await workerA.send_message(message=_message())
+    await workerA.release_message(message=await _claim(queue=workerA))
+    await _claim(queue=workerB)
+
+
 async def test_extended_message_stays_invisible_and_can_still_be_completed(workers: tuple[SqsMessageQueue, SqsMessageQueue]):
     workerA, workerB = workers
     await workerA.send_message(message=_message())