kiba-core 0.5.3.dev64__tar.gz → 0.5.3.dev66__tar.gz

kiba_core-0.5.3.dev66/.github/workflows/pull-request.yml

@@ -0,0 +1,76 @@
+name: Build
+on:
+  pull_request:
+    branches: ["*"]
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-pull-request
+  cancel-in-progress: true
+jobs:
+  core-build:
+    runs-on: ubuntu-latest
+    env:
+      DOCKER_IMAGE: core
+    steps:
+      - name: Check out the repo
+        uses: actions/checkout@v2
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Build docker image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          load: true
+          cache-from: type=gha,scope=core
+          cache-to: type=gha,mode=max,scope=core
+          tags: ${{ env.DOCKER_IMAGE }}
+      - name: Run lint-check
+        run: |
+          docker run --name lint-check $DOCKER_IMAGE make lint-check-ci
+          docker cp lint-check:/app/lint-check-results.json ./lint-check-results.json
+      - name: Upload lint-check results
+        uses: kibalabs/github-action-create-annotations@main
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          json-file-path: ./lint-check-results.json
+          check-name: core-lint-check
+          fail-on-error: false
+      - name: Run type-check
+        run: |
+          docker run --name type-check $DOCKER_IMAGE make type-check-ci
+          docker cp type-check:/app/type-check-results.json ./type-check-results.json
+      - name: Upload type-check results
+        uses: kibalabs/github-action-create-annotations@main
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          json-file-path: ./type-check-results.json
+          check-name: core-type-check
+          fail-on-error: false
+      - name: Start test services
+        run: |
+          docker network create tests
+          docker run --detach --name postgres --network tests --env POSTGRES_PASSWORD=test --env POSTGRES_DB=core postgres:16-alpine
+          docker run --detach --name cosmos --network tests mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator@sha256:2db1f9e74c506bcf6fc347aa937aea1c00fa756061296a5a9efba530ce86ec02 --protocol http
+          docker run --detach --name elasticmq --network tests --publish 9324:9324 softwaremill/elasticmq-native@sha256:e4580ab9ad1bd5cd37b4ba04911bc5ccc8cd2d9ab4de56ece65acee71c24e05c
+          docker run --detach --name azurite --network tests --publish 10001:10001 mcr.microsoft.com/azure-storage/azurite@sha256:830430c1da1a2d537e08f3e6764dd1f5ae00cf0346bcaf625b968ec3f0971fd5 azurite-queue --queueHost 0.0.0.0 --skipApiVersionCheck
+          echo "Waiting for postgres..."
+          # NOTE(krishan711): check over TCP: during initdb the image runs a temporary server on the unix socket only, which reports ready and then shuts down
+          for i in $(seq 1 60); do docker exec postgres pg_isready --host 127.0.0.1 --username postgres --quiet && break; sleep 1; done
+          docker exec postgres pg_isready --host 127.0.0.1 --username postgres
+          echo "Waiting for sqs and aqs emulators..."
+          # NOTE(krishan711): any http response (even an error) means the emulator is accepting requests
+          for i in $(seq 1 60); do curl --silent --output /dev/null http://localhost:9324/ && curl --silent --output /dev/null http://localhost:10001/ && break; sleep 1; done
+          curl --silent --output /dev/null http://localhost:9324/ && curl --silent --output /dev/null http://localhost:10001/
+          echo "Waiting for cosmos emulator..."
+          for i in $(seq 1 120); do docker exec cosmos curl --silent --fail --output /dev/null http://localhost:8080/ready && break; sleep 1; done
+          docker exec cosmos curl --silent --fail --output /dev/null http://localhost:8080/ready
+      - name: Run tests
+        run: |
+          docker run --name tests --network tests --env CORE_TEST_PSQL_CONNECTION_STRING=postgresql+asyncpg://postgres:test@postgres:5432/core --env CORE_TEST_COSMOS_ENDPOINT=http://cosmos:8081 --env CORE_TEST_SQS_ENDPOINT=http://elasticmq:9324 --env CORE_TEST_AQS_ENDPOINT=http://azurite:10001/devstoreaccount1 $DOCKER_IMAGE make test-ci
+          docker cp tests:/app/test-check-results.json ./test-check-results.json
+      - name: Upload test results
+        uses: kibalabs/github-action-create-annotations@main
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          json-file-path: ./test-check-results.json
+          check-name: core-tests
+          fail-on-error: true

{kiba_core-0.5.3.dev64 → kiba_core-0.5.3.dev66}/CHANGELOG.md

@@ -8,6 +8,20 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 
 ### Added
+- [MAJOR] Replaced `MessageQueue.delete_message` with `complete_message`, and added abstract `retry_message` (reschedule with a delay) and `fail_message` (apply the queue's failure policy); `MessageQueueProcessor` now reschedules `MessageNeedsReprocessingException` messages with `retry_message`
+- [MINOR] Updated `MessageQueueProcessor` to retry messages that raise `LockedException` up to 3 times (30s linear backoff) before failing them
+- [MINOR] Updated `MessageQueueProcessor` to log instead of raise when a failure notification cannot be sent
+- [MAJOR] Reworked `SqlMessageQueue` to keep message history in `tbl_queue_messages` with a `status` lifecycle (`pending`, `running`, `succeeded`, `failed`, `deduplicated`), `owner`, `lastError`, `startedDate` and `completedDate`; deduplication now only applies against pending messages and keeps the earliest `visibleDate` of the duplicates, every retry, failure retry and lease-expiry reclaim increments `postCount`, failures back off and move to `failed` once `postCount` reaches `maxAttempts`, and `get_message_count` removes finished messages older than `retentionSeconds` (default 7 days) unless `shouldSkipRemovingNonRetained` is set
+- [MINOR] Added `Database.now` to build database-clock timestamps; `SqlMessageQueue` and `SqlLock` now use the database clock instead of the host clock
+- [MINOR] Capped `SqsMessageQueue` delays at SQS's 900 second maximum
+- [MINOR] Updated `CosmosMessageQueue.retry_message` to reschedule the message in place and added emulator-backed tests (run when `CORE_TEST_COSMOS_ENDPOINT` is set)
+- [MINOR] Added `shouldRaiseOnUncommittedWrites` to `Database` to raise instead of log when an isolated context connection is opened after uncommitted writes (for tests and development)
+- [MAJOR] Added abstract `MessageQueue.extend_message_lease` (implemented for SQL, SQS, AQS and Cosmos queues, each returning `False` once the lease is lost; SQS accepts any receipt handle for a message so it can only detect that the message was deleted or is no longer in flight); `MessageQueueProcessor` now keeps a message's lease alive while it is being processed and, if the lease was lost, raises `MessageLeaseLostException` once the job returns (taking priority over any exception from the job, while cancellation still propagates) so the message is neither completed nor retried; `SqsMessageQueue` and `AqsMessageQueue` `retry_message` check the lease is still held before re-sending so a lost lease cannot duplicate the message
+- [MINOR] Added optional `endpointUrl` to `SqsMessageQueue` and `AqsMessageQueue` to target local emulators, and added emulator-backed tests (run when `CORE_TEST_SQS_ENDPOINT` / `CORE_TEST_AQS_ENDPOINT` are set)
+- [MINOR] Added `Lock.ensure_held` to renew a lease and raise `LockedException` if it has been lost, for use before irreversible work; `LockLease` now records its `ttlSeconds`
+- [MINOR] Added `Lock` with a lease-based `SqlLock` implementation (`tbl_locks`) and a `with_lock` context manager that keeps the lease alive; acquiring raises `LockedException` (naming the current holder) if not acquired within `maxWaitSeconds`, and `get_owner` reports who holds a lock
+- [MINOR] Added `process_util.get_process_name` to identify the current process (`<containerName>:<containerId>` in docker, `<hostname>:<pid>` otherwise)
+- [MINOR] Added `Database.create_isolated_context_connection` to run a block in its own committed transaction inside an existing context connection; logs `ISOLATED_CONNECTION_AFTER_UNCOMMITTED_WRITES` when opened while the outer context connection has uncommitted writes
 - [MAJOR] Added abstract `MessageQueue.get_message_count` (implemented for SQS, AQS, SQL and Cosmos queues) for the number of messages waiting to be processed
 - [MAJOR] Added abstract `MessageQueue.get_inflight_message_count` (implemented for SQS, SQL and Cosmos queues; AQS raises `NotImplementedError` as Azure does not expose it) for the number of messages currently being processed
 - [MAJOR] Reworked `route` to compose JSON or streaming transport, application-resolved authorization, and rate limiting

{kiba_core-0.5.3.dev64 → kiba_core-0.5.3.dev66}/PKG-INFO

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

kiba_core-0.5.3.dev66/core/locks/lock.py

@@ -0,0 +1,67 @@
+import abc
+import asyncio
+import contextlib
+from abc import ABC
+from collections.abc import AsyncIterator
+
+from core import logging
+from core.exceptions import LockedException
+from core.locks.model import LockLease
+
+
+class Lock(ABC):
+    @abc.abstractmethod
+    async def connect(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def disconnect(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> LockLease:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def extend(self, lease: LockLease, ttlSeconds: int = 60) -> bool:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def release(self, lease: LockLease) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_owner(self, name: str) -> str | None:
+        raise NotImplementedError
+
+    async def ensure_held(self, lease: LockLease) -> None:
+        if lease.isLost or not await self.extend(lease=lease, ttlSeconds=lease.ttlSeconds):
+            lease.isLost = True
+            raise LockedException(message=f'LOCK_LOST: {lease.name}')
+
+    async def _keep_alive(self, lease: LockLease) -> None:
+        while True:
+            await asyncio.sleep(lease.ttlSeconds / 3)
+            try:
+                isExtended = await self.extend(lease=lease, ttlSeconds=lease.ttlSeconds)
+            except Exception as exception:  # noqa: BLE001
+                # NOTE(krishan711): a transient failure must not kill the keep-alive; the lease simply lapses at expiryDate if it keeps failing
+                logging.error(f'Failed to extend lease {lease.name}:')
+                logging.exception(exception)
+                continue
+            if not isExtended:
+                lease.isLost = True
+                logging.error(f'Lost lease {lease.name}')
+                return
+
+    @contextlib.asynccontextmanager
+    async def with_lock(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> AsyncIterator[LockLease]:
+        lease = await self.acquire(name=name, ttlSeconds=ttlSeconds, maxWaitSeconds=maxWaitSeconds)
+        keepAliveTask = asyncio.create_task(self._keep_alive(lease=lease))
+        try:
+            yield lease
+        finally:
+            keepAliveTask.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await keepAliveTask
+            await self.release(lease=lease)

kiba_core-0.5.3.dev66/core/locks/model.py

@@ -0,0 +1,16 @@
+import datetime
+
+from pydantic import BaseModel
+
+from core.util import date_util
+
+
+class LockLease(BaseModel):
+    name: str
+    token: str
+    expiryDate: datetime.datetime
+    ttlSeconds: int
+    isLost: bool = False
+
+    def is_valid(self) -> bool:
+        return not self.isLost and date_util.datetime_from_now() < self.expiryDate

kiba_core-0.5.3.dev66/core/locks/sql.py

@@ -0,0 +1,104 @@
+import asyncio
+import time
+import uuid
+
+import sqlalchemy
+from sqlalchemy.dialects import postgresql as sqlalchemy_psql
+from sqlalchemy.dialects import sqlite as sqlalchemy_sqlite
+
+from core.exceptions import InternalServerErrorException
+from core.exceptions import LockedException
+from core.locks.lock import Lock
+from core.locks.model import LockLease
+from core.store.database import Database
+from core.util import date_util
+from core.util import process_util
+
+# NOTE(krishan711): this table is owned here (not in the consuming app's schema.py) so it can ship with
+# core-py. To include it in an app's own alembic-tracked metadata (so autogenerate creates the table),
+# call `LocksTable.to_metadata(appMetadata)` from the app's schema module.
+LocksMetadata = sqlalchemy.MetaData()
+
+LocksTable = sqlalchemy.Table(
+    'tbl_locks',
+    LocksMetadata,
+    sqlalchemy.Column(key='name', name='name', type_=sqlalchemy.Text, primary_key=True, nullable=False),
+    sqlalchemy.Column(key='lockToken', name='lock_token', type_=sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column(key='acquiredDate', name='acquired_date', type_=sqlalchemy.DateTime(timezone=True), nullable=False),
+    sqlalchemy.Column(key='expiryDate', name='expiry_date', type_=sqlalchemy.DateTime(timezone=True), nullable=False),
+    sqlalchemy.Column(key='owner', name='owner', type_=sqlalchemy.Text, nullable=False),
+)
+
+
+class SqlLock(Lock):
+    def __init__(self, database: Database, table: sqlalchemy.Table = LocksTable, pollIntervalSeconds: float = 1.0, owner: str | None = None) -> None:
+        self.database = database
+        self.table = table
+        self.pollIntervalSeconds = pollIntervalSeconds
+        self.owner = owner or process_util.get_process_name()
+
+    async def connect(self) -> None:
+        pass
+
+    async def disconnect(self) -> None:
+        pass
+
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> LockLease:
+        deadline = time.monotonic() + maxWaitSeconds
+        while True:
+            lease = await self._try_acquire(name=name, ttlSeconds=ttlSeconds)
+            if lease is not None:
+                return lease
+            if time.monotonic() >= deadline:
+                owner = await self.get_owner(name=name)
+                raise LockedException(message=f'LOCK_HELD: {name} by {owner}')
+            await asyncio.sleep(min(self.pollIntervalSeconds, max(deadline - time.monotonic(), 0)))
+
+    async def get_owner(self, name: str) -> str | None:
+        async with self.database.create_transaction() as connection:
+            query = sqlalchemy.select(self.table.c.owner).where(self.table.c.name == name).where(self.table.c.expiryDate > self.database.now())
+            result = await self.database.execute(query=query, connection=connection)
+            return result.scalar_one_or_none()
+
+    def _create_insert(self, dialectName: str) -> sqlalchemy_psql.Insert | sqlalchemy_sqlite.Insert:
+        if dialectName == 'postgresql':
+            return sqlalchemy_psql.insert(self.table)
+        if dialectName == 'sqlite':
+            return sqlalchemy_sqlite.insert(self.table)
+        raise InternalServerErrorException(message=f'SqlLock does not support dialect: {dialectName}')
+
+    async def _try_acquire(self, name: str, ttlSeconds: int) -> LockLease | None:
+        lease = LockLease(name=name, token=str(uuid.uuid4()), expiryDate=date_util.datetime_from_now(seconds=ttlSeconds), ttlSeconds=ttlSeconds)
+        values = {self.table.c.lockToken: lease.token, self.table.c.acquiredDate: self.database.now(), self.table.c.expiryDate: self.database.now(seconds=ttlSeconds), self.table.c.owner: self.owner}
+        async with self.database.create_transaction() as connection:
+            upsertQuery = (
+                self._create_insert(dialectName=connection.dialect.name)
+                .values({self.table.c.name: name, **values})
+                .on_conflict_do_update(index_elements=[self.table.c.name], set_=values, where=self.table.c.expiryDate <= self.database.now())
+                .returning(self.table.c.name)
+            )
+            result = await self.database.execute(query=upsertQuery, connection=connection)
+            isAcquired = result.first() is not None
+        return lease if isAcquired else None
+
+    async def extend(self, lease: LockLease, ttlSeconds: int = 60) -> bool:
+        newExpiryDate = date_util.datetime_from_now(seconds=ttlSeconds)
+        async with self.database.create_transaction() as connection:
+            updateQuery = (
+                sqlalchemy.update(self.table)
+                .where(self.table.c.name == lease.name)
+                .where(self.table.c.lockToken == lease.token)
+                .where(self.table.c.expiryDate > self.database.now())
+                .values(expiryDate=self.database.now(seconds=ttlSeconds))
+                .returning(self.table.c.name)
+            )
+            result = await self.database.execute(query=updateQuery, connection=connection)
+            isExtended = result.first() is not None
+        if isExtended:
+            lease.expiryDate = newExpiryDate
+        return isExtended
+
+    async def release(self, lease: LockLease) -> None:
+        async with self.database.create_transaction() as connection:
+            deleteQuery = sqlalchemy.delete(self.table).where(self.table.c.name == lease.name).where(self.table.c.lockToken == lease.token)
+            await self.database.execute(query=deleteQuery, connection=connection)

{kiba_core-0.5.3.dev64 → kiba_core-0.5.3.dev66}/core/queues/aqs.py

@@ -3,7 +3,9 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Sequence
 
+from azure.core.exceptions import HttpResponseError
 from azure.storage.queue import QueueMessage as RawAqsMessage
+from azure.storage.queue import StorageErrorCode
 from azure.storage.queue.aio import QueueClient
 
 from core.exceptions import InternalServerErrorException
@@ -11,6 +13,8 @@ from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
 from core.util import list_util
 
+AQS_RETRY_HOLD_SECONDS = 60
+
 
 class AqsMessage(Message):
     aqsId: str
@@ -32,14 +36,18 @@ class AqsMessage(Message):
 
 
 class AqsMessageQueue(MessageQueue[AqsMessage]):
-    def __init__(self, storageAccountName: str, storageAccountKey: str, queueName: str) -> None:
+    def __init__(self, storageAccountName: str, storageAccountKey: str, queueName: str, endpointUrl: str | None = None) -> None:
         self._storageAccountName = storageAccountName
         self._storageAccountKey = storageAccountKey
         self.queueName = queueName
+        self.endpointUrl = endpointUrl
         self._aqsClient: QueueClient | None = None
 
     async def connect(self) -> None:
-        self._aqsClient = QueueClient.from_connection_string(conn_str=f'DefaultEndpointsProtocol=https;AccountName={self._storageAccountName};AccountKey={self._storageAccountKey}', queue_name=self.queueName)
+        connectionString = f'DefaultEndpointsProtocol=https;AccountName={self._storageAccountName};AccountKey={self._storageAccountKey}'
+        if self.endpointUrl:
+            connectionString += f';QueueEndpoint={self.endpointUrl}'
+        self._aqsClient = QueueClient.from_connection_string(conn_str=connectionString, queue_name=self.queueName)
         if not self._aqsClient:
             raise InternalServerErrorException('Failed to connect to queue')
         await self._aqsClient.get_queue_properties()
@@ -76,11 +84,35 @@ class AqsMessageQueue(MessageQueue[AqsMessage]):
         aqsMessages = [AqsMessage.from_aqs_message(aqsMessage=message) async for message in messagesIterator]
         return aqsMessages
 
-    async def delete_message(self, message: AqsMessage) -> None:
+    async def complete_message(self, message: AqsMessage) -> None:
         if not self._aqsClient:
-            raise InternalServerErrorException('You need to call .connect() before trying to delete messages')
+            raise InternalServerErrorException('You need to call .connect() before trying to complete messages')
         await self._aqsClient.delete_message(message=message.aqsId, pop_receipt=message.popReceipt)
 
+    async def retry_message(self, message: AqsMessage, delaySeconds: int = 0) -> None:
+        # NOTE(krishan711): check the lease is still held before re-sending, otherwise another worker owns the message and re-sending would duplicate it
+        if not await self.extend_message_lease(message=message, expectedProcessingSeconds=AQS_RETRY_HOLD_SECONDS):
+            return
+        await self.send_message(message=message, delaySeconds=delaySeconds)
+        await self.complete_message(message=message)
+
+    async def fail_message(self, message: AqsMessage, errorMessage: str | None) -> None:
+        # NOTE(krishan711): the message reappears after its visibility timeout
+        pass
+
+    async def extend_message_lease(self, message: AqsMessage, expectedProcessingSeconds: int) -> bool:
+        if not self._aqsClient:
+            raise InternalServerErrorException('You need to call .connect() before trying to extend message leases')
+        try:
+            updatedMessage = await self._aqsClient.update_message(message=message.aqsId, pop_receipt=message.popReceipt, visibility_timeout=expectedProcessingSeconds)
+        except HttpResponseError as exception:
+            if getattr(exception, 'error_code', None) in {StorageErrorCode.POP_RECEIPT_MISMATCH, StorageErrorCode.MESSAGE_NOT_FOUND}:
+                return False
+            raise
+        # NOTE(krishan711): azure invalidates the old pop receipt on every update so later calls must use the new one
+        message.popReceipt = updatedMessage.pop_receipt
+        return True
+
     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.dev64 → kiba_core-0.5.3.dev66}/core/queues/cosmos.py

@@ -13,6 +13,7 @@ from azure.cosmos.aio import ContainerProxy
 
 from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
+from core.util import date_util
 from core.util.typing_util import JsonObject
 
 
@@ -169,7 +170,7 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
                 break
         return messages
 
-    async def delete_message(self, message: CosmosMessage) -> None:
+    async def complete_message(self, message: CosmosMessage) -> None:
         if message.deduplicationId is None:
             await self.container.delete_item(
                 item=message.id,
@@ -189,6 +190,46 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
             partition_key=self.queueName,
         )
 
+    async def retry_message(self, message: CosmosMessage, delaySeconds: int = 0) -> 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() + delaySeconds},
+                    {'op': 'set', 'path': '/postCount', 'value': (message.postCount or 0) + 1},
+                    {'op': 'set', 'path': '/postDate', 'value': date_util.datetime_from_now().isoformat()},
+                ],
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except cosmos_exceptions.CosmosHttpResponseError as exception:
+            # NOTE(krishan711): 412 means the lease was lost to another worker, which now owns the message
+            if exception.status_code != 412:  # noqa: PLR2004
+                raise
+
+    async def fail_message(self, message: CosmosMessage, errorMessage: str | None) -> None:
+        # NOTE(krishan711): the message reappears once its lease expires
+        pass
+
+    async def extend_message_lease(self, message: CosmosMessage, expectedProcessingSeconds: int) -> bool:
+        try:
+            updatedItem = await self.container.patch_item(
+                item=message.id,
+                partition_key=self.queueName,
+                patch_operations=[{'op': 'set', 'path': '/visibleDate', 'value': time.time() + expectedProcessingSeconds}],
+                etag=message.etag,
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except cosmos_exceptions.CosmosHttpResponseError as exception:
+            if exception.status_code == 412:  # noqa: PLR2004
+                return False
+            raise
+        # NOTE(krishan711): the patch changes the etag so later calls must use the new one to prove they still hold the lease
+        message.etag = str(updatedItem['_etag'])
+        return True
+
     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.dev64 → kiba_core-0.5.3.dev66}/core/queues/message_queue.py

@@ -31,7 +31,19 @@ class MessageQueue[MessageType: Message](ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    async def delete_message(self, message: MessageType) -> None:
+    async def complete_message(self, message: MessageType) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def retry_message(self, message: MessageType, delaySeconds: int = 0) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def fail_message(self, message: MessageType, errorMessage: str | None) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def extend_message_lease(self, message: MessageType, expectedProcessingSeconds: int) -> bool:
         raise NotImplementedError
 
     @abc.abstractmethod

{kiba_core-0.5.3.dev64 → kiba_core-0.5.3.dev66}/core/queues/message_queue_processor.py

@@ -1,5 +1,7 @@
 import abc
 import asyncio
+import contextlib
+import sys
 import time
 import urllib.parse as urlparse
 import uuid
@@ -8,11 +10,15 @@ from abc import ABC
 from core import logging
 from core.exceptions import InternalServerErrorException
 from core.exceptions import KibaException
+from core.exceptions import LockedException
 from core.notifications.notification_client import NotificationClient
 from core.queues.message_queue import MessageQueue
 from core.queues.model import Message
 from core.util.value_holder import RequestIdHolder
 
+LOCKED_MAX_RETRY_COUNT = 3
+LOCKED_RETRY_DELAY_SECONDS = 30
+
 
 class MessageProcessor(ABC):
     @abc.abstractmethod
@@ -28,6 +34,11 @@ class MessageNeedsReprocessingException(InternalServerErrorException):
         self.originalException = originalException
 
 
+class MessageLeaseLostException(InternalServerErrorException):
+    def __init__(self) -> None:
+        super().__init__(message='MESSAGE_LEASE_LOST')
+
+
 class MessageQueueProcessor[MessageType: Message]:
     def __init__(self, queue: MessageQueue[MessageType], messageProcessor: MessageProcessor, notificationClients: list[NotificationClient], requestIdHolder: RequestIdHolder | None = None) -> None:
         self.queue = queue
@@ -35,7 +46,36 @@ class MessageQueueProcessor[MessageType: Message]:
         self.notificationClients = notificationClients
         self.requestIdHolder = requestIdHolder
 
-    async def _process_message(self, message: MessageType) -> None:
+    async def _handle_failure(self, message: MessageType, exception: Exception, requestId: str) -> int:
+        logging.error('Caught exception whilst processing message:')
+        logging.exception(exception)
+        kibaException = KibaException.from_exception(exception=exception)
+        await self.queue.fail_message(message=message, errorMessage=kibaException.message)
+        for client in self.notificationClients:
+            try:
+                await client.post(messageText=f'Error processing message: {message.command}\n```\n{requestId}\n{message.content}\n{kibaException.message}```')
+            except Exception as notificationException:  # noqa: BLE001
+                logging.error('Failed to send message failure notification:')
+                logging.exception(notificationException)
+        return exception.statusCode if isinstance(exception, KibaException) else 500
+
+    async def _keep_message_alive(self, message: MessageType, expectedProcessingSeconds: int, stopEvent: asyncio.Event) -> bool:
+        while True:
+            with contextlib.suppress(TimeoutError):
+                await asyncio.wait_for(stopEvent.wait(), timeout=expectedProcessingSeconds / 3)
+            if stopEvent.is_set():
+                return False
+            try:
+                isExtended = await self.queue.extend_message_lease(message=message, expectedProcessingSeconds=expectedProcessingSeconds)
+            except Exception as exception:  # noqa: BLE001
+                logging.error(f'Failed to extend message lease {message.command}:')
+                logging.exception(exception)
+                continue
+            if not isExtended:
+                logging.error(f'Lost message lease for {message.command}')
+                return True
+
+    async def _process_message(self, message: MessageType, expectedProcessingSeconds: int) -> None:
         requestId = message.requestId or str(uuid.uuid4()).replace('-', '')
         if self.requestIdHolder:
             self.requestIdHolder.set_value(value=requestId)
@@ -43,20 +83,34 @@ class MessageQueueProcessor[MessageType: Message]:
         logging.api(action='MESSAGE', path=message.command, pathPattern=message.command, query=query)
         startTime = time.time()
         statusCode = 200
+        # NOTE(krishan711): the keep-alive is stopped and awaited (not cancelled) so an in-flight extend finishes before the message is completed, retried or failed.
+        # A lost lease is raised once the job returns rather than interrupting it, as cancelling mid-job could abandon work that was already sent.
+        # It is raised from the finally so it replaces any exception from the job: the message belongs to another worker, so it must not be retried or completed here.
+        # 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))
         try:
-            await self.messageProcessor.process_message(message=message)
-            await self.queue.delete_message(message=message)
+            try:
+                await self.messageProcessor.process_message(message=message)
+            finally:
+                stopKeepAliveEvent.set()
+                isLeaseLost = await keepAliveTask
+                inFlightException = sys.exception()
+                if isLeaseLost and (inFlightException is None or isinstance(inFlightException, Exception)):
+                    raise MessageLeaseLostException
+            await self.queue.complete_message(message=message)
         except MessageNeedsReprocessingException as exception:
             logging.info(msg=f'Scheduling reprocessing for message:{message.command} due to: {exception.originalException!s}')
-            await self.queue.send_message(message=message, delaySeconds=((message.postCount or 0) * exception.delaySeconds))
+            await self.queue.retry_message(message=message, delaySeconds=((message.postCount or 0) * exception.delaySeconds))
+        except LockedException as exception:
+            postCount = message.postCount or 0
+            if postCount <= LOCKED_MAX_RETRY_COUNT:
+                logging.info(msg=f'Scheduling reprocessing for message:{message.command} due to: {exception!s}')
+                await self.queue.retry_message(message=message, delaySeconds=(postCount * LOCKED_RETRY_DELAY_SECONDS))
+            else:
+                statusCode = await self._handle_failure(message=message, exception=exception, requestId=requestId)
         except Exception as exception:  # noqa: BLE001
-            statusCode = exception.statusCode if isinstance(exception, KibaException) else 500
-            logging.error('Caught exception whilst processing message:')
-            logging.exception(exception)
-            kibaException = KibaException.from_exception(exception=exception)
-            for client in self.notificationClients:
-                await client.post(messageText=f'Error processing message: {message.command}\n```\n{requestId}\n{message.content}\n{kibaException.message}```')
-            # TODO(krishan711): should possibly reset the visibility timeout
+            statusCode = await self._handle_failure(message=message, exception=exception, requestId=requestId)
         duration = time.time() - startTime
         logging.api(action='MESSAGE', path=message.command, pathPattern=message.command, query=query, response=statusCode, duration=duration)
         if self.requestIdHolder:
@@ -66,10 +120,10 @@ class MessageQueueProcessor[MessageType: Message]:
         logging.info('Retrieving messages...')
         messages = await self.queue.get_messages(expectedProcessingSeconds=expectedProcessingSeconds, longPollSeconds=longPollSeconds, limit=batchSize)
         if shouldProcessInParallel:
-            await asyncio.gather(*[self._process_message(message=message) for message in messages])
+            await asyncio.gather(*[self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds) for message in messages])
         else:
             for message in messages:
-                await self._process_message(message=message)
+                await self._process_message(message=message, expectedProcessingSeconds=expectedProcessingSeconds)
         return len(messages)
 
     async def execute(self, expectedProcessingSeconds: int = 300, longPollSeconds: int = 20) -> bool: