fix(infrastructure): replace taskiq with remote_funcs (#62)

Author: emptybutton

pyproject.toml

@@ -25,7 +25,6 @@ dependencies = [
     "openai==1.97.0",
     "structlog==25.4.0",
     "structlog-sentry==2.2.1",
-    "taskiq==0.11.18",
 ]
 
 [dependency-groups]
@@ -117,7 +116,6 @@ ignore = [
 
 [tool.ruff.lint.per-file-ignores]
 "src/ttt/infrastructure/sqlalchemy/tables/__init__.py" = ["F401"]
-"src/ttt/infrastructure/taskiq/tasks/__init__.py" = ["F401"]
 "src/ttt/infrastructure/adapters/*" = ["ARG002"]
 "src/ttt/presentation/adapters/*" = ["ARG002"]
 "src/ttt/presentation/*" = ["RUF001"]

src/ttt/infrastructure/adapters/game_log.py

@@ -88,7 +88,7 @@ async def already_completed_game_to_make_ai_move(
         /,
     ) -> None:
         await self._logger.ainfo(
-            "already_completed_game_to_make_move",
+            "already_completed_game_to_make_ai_move",
             ai_id=ai_id.hex,
             game_id=game.id.hex,
         )

src/ttt/infrastructure/adapters/game_tasks.py

@@ -2,18 +2,20 @@
 from uuid import UUID
 
 from ttt.application.game.game.ports.game_tasks import GameTasks
-from ttt.infrastructure.taskiq.tasks.make_ai_move_in_game_task import (
-    make_ai_move_in_game_task,
+from ttt.infrastructure.remote_funcs.make_ai_move_in_game import (
+    make_ai_move_in_game_remotely,
 )
 
 
 @dataclass
-class TaskiqGameTasks(GameTasks):
+class NatsRemoteFuncGameTasks(GameTasks):
     async def make_ai_move(
         self,
         user_id: int,
         game_id: UUID,
         ai_id: UUID,
         /,
     ) -> None:
-        await make_ai_move_in_game_task.kiq(user_id, game_id, ai_id)
+        await make_ai_move_in_game_remotely(
+            user_id=user_id, game_id=game_id.hex, ai_id=ai_id.hex,
+        )

src/ttt/infrastructure/adapters/stars_purchase_tasks.py

@@ -5,21 +5,21 @@
     StarsPurchaseTasks,
 )
 from ttt.entities.finance.payment.success import PaymentSuccess
-from ttt.infrastructure.taskiq.tasks.complete_stars_purchase_payment_task import (  # noqa: E501
-    complete_stars_purchase_payment_task,
+from ttt.infrastructure.remote_funcs.complete_stars_purchase_payment import (
+    complete_stars_purchase_payment_remotely,
 )
 
 
 @dataclass
-class TaskiqStarsPurchaseTasks(StarsPurchaseTasks):
+class NatsRemoteFuncStarsPurchaseTasks(StarsPurchaseTasks):
     async def complete_stars_purchase_payment(
         self,
         purchase_id: UUID,
         success: PaymentSuccess,
         /,
     ) -> None:
-        await complete_stars_purchase_payment_task.kiq(
-            purchase_id,
-            success.id,
-            success.gateway_id,
+        await complete_stars_purchase_payment_remotely(
+            purchase_id=purchase_id.hex,
+            payment_success_id=success.id,
+            payment_success_gateway_id=success.gateway_id,
         )

src/ttt/infrastructure/multi_asynccontextmanager.py

@@ -0,0 +1,23 @@
+from asyncio import gather
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+
+
+@asynccontextmanager
+async def multi_asynccontextmanager[T](
+    *managers: AbstractAsyncContextManager[T],
+) -> AsyncIterator[list[T]]:
+    result = await gather(*(manager.__aenter__() for manager in managers))  # noqa: PLC2801
+
+    try:
+        yield result
+    except BaseException as error:  # noqa: BLE001
+        await gather(*(
+            manager.__aexit__(type(error), error, error.__traceback__)
+            for manager in managers
+        ))
+    else:
+        await gather(*(
+            manager.__aexit__(None, None, None)
+            for manager in managers
+        ))

src/ttt/infrastructure/processors/processors.py

@@ -0,0 +1,40 @@
+from uuid import UUID
+
+from dishka import AsyncContainer
+
+from ttt.application.stars_purchase.complete_stars_purchase_payment import (
+    CompleteStarsPurchasePayment,
+)
+from ttt.entities.finance.payment.success import PaymentSuccess
+from ttt.infrastructure.remote_funcs.nats_remote_func import nats_remote
+from ttt.infrastructure.retrier import Retrier
+
+
+@nats_remote(
+    subject="stars_purchase.stars_purchase.complete_stars_purchase_payment",
+    pull_subscribe=lambda js, subject: js.pull_subscribe(
+        subject,
+        "ttt-stars_purchase-stars_purchase-complete_stars_purchase_payment",
+        stream="STARS_PURCHASE",
+    ),
+)
+async def complete_stars_purchase_payment_remotely(
+    container: AsyncContainer,
+    *,
+    purchase_id: str,
+    payment_success_id: str,
+    payment_success_gateway_id: str,
+) -> None:
+    payment_success = PaymentSuccess(
+        payment_success_id, payment_success_gateway_id,
+    )
+
+    retrier = await container.get(Retrier)
+    complete_stars_purchase_payment = await container.get(
+        CompleteStarsPurchasePayment,
+    )
+    await retrier(
+        complete_stars_purchase_payment,
+        UUID(hex=purchase_id),
+        payment_success,
+    )

src/ttt/infrastructure/taskiq/__init__.py …/infrastructure/remote_funcs/__init__.pysrc/ttt/infrastructure/taskiq/__init__.py renamed to src/ttt/infrastructure/remote_funcs/__init__.py src/ttt/infrastructure/taskiq/__init__.py renamed to src/ttt/infrastructure/remote_funcs/__init__.py

@@ -0,0 +1,30 @@
+from uuid import UUID
+
+from dishka import AsyncContainer
+
+from ttt.application.game.game.make_ai_move_in_game import MakeAiMoveInGame
+from ttt.infrastructure.remote_funcs.nats_remote_func import nats_remote
+from ttt.infrastructure.retrier import Retrier
+
+
+@nats_remote(
+    subject="game.game.make_ai_move_in_game",
+    pull_subscribe=lambda js, subject: js.pull_subscribe(
+        subject,
+        durable="ttt-game-game-make_ai_move_in_game",
+        stream="GAME",
+    ),
+)
+async def make_ai_move_in_game_remotely(
+    container: AsyncContainer,
+    *,
+    user_id: int,
+    game_id: str,
+    ai_id: str,
+) -> None:
+    retrier = await container.get(Retrier)
+    make_ai_move_in_game = await container.get(MakeAiMoveInGame)
+
+    await retrier(
+        make_ai_move_in_game, user_id, UUID(hex=game_id), UUID(hex=ai_id),
+    )

src/ttt/infrastructure/remote_funcs/complete_stars_purchase_payment.py

@@ -0,0 +1,161 @@
+import json
+from asyncio import (
+    AbstractEventLoop,
+    Semaphore,
+    Task,
+    gather,
+    get_event_loop,
+)
+from collections.abc import (
+    AsyncIterator,
+    Callable,
+)
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any, Protocol, Self
+
+from dishka import AsyncContainer
+from nats.aio.msg import Msg as NatsMessage
+from nats.errors import TimeoutError as NatsTimeoutError
+from nats.js import JetStreamContext
+from structlog.types import FilteringBoundLogger
+
+from ttt.entities.tools.assertion import not_none
+from ttt.infrastructure.dishka.next_container import NextContainer
+from ttt.infrastructure.structlog.logger import (
+    unexpected_error_log,
+)
+
+
+class PullSubscribe(Protocol):
+    async def __call__(
+        self, js: JetStreamContext, subject: str,
+    ) -> JetStreamContext.PullSubscription: ...
+
+
+class NatsRemoteFuncBody[**PmT](Protocol):
+    async def __call__(
+        self,
+        container: AsyncContainer,
+        *args: PmT.args,
+        **kwargs: PmT.kwargs,
+    ) -> Any: ...  # noqa: ANN401
+
+
+@dataclass
+class NatsRemoteFunc[**PmT = ...]:
+    body: NatsRemoteFuncBody[PmT]
+    _pull_subscribe: PullSubscribe
+    _subject: str
+
+    _max_workers: int | None = None
+    _pull_consume_batch: int | None = None
+    _pull_consume_timeout: float | None = None
+
+    _js: JetStreamContext = field(init=False)
+    _pull_subscription: JetStreamContext.PullSubscription = field(init=False)
+    _workers: set[Task[None]] = field(init=False, default_factory=set)
+    _semaphore: Semaphore = field(init=False, default_factory=Semaphore)
+    _container: AsyncContainer = field(init=False)
+    _loop: AbstractEventLoop = field(init=False)
+
+    @asynccontextmanager
+    async def startup(
+        self,
+        js: JetStreamContext,
+        max_workers: int | None = None,
+        pull_consume_batch: int | None = None,
+        pull_consume_timeout: float | None = None,
+    ) -> AsyncIterator[Self]:
+        self._loop = get_event_loop()
+        self._max_workers = max_workers or self._max_workers or 1000
+        self._pull_consume_batch = (
+            pull_consume_batch or self._pull_consume_batch or 1
+        )
+        self._pull_consume_timeout = (
+            pull_consume_timeout or self._pull_consume_timeout or 5
+        )
+
+        self._js = js
+        self._pull_subscription = await self._pull_subscribe(
+            self._js,
+            self._subject,
+        )
+        self._semaphore = Semaphore(self._max_workers)
+
+        try:
+            yield self
+        finally:
+            await gather(*self._workers)
+
+    async def __call__(self, *args: PmT.args, **kwargs: PmT.kwargs) -> None:
+        payload = {"args": args, "kwargs": kwargs}
+
+        await not_none(self._js).publish(
+            self._subject,
+            payload=json.dumps(payload).encode(),
+        )
+
+    async def processor(self, next_container: NextContainer) -> None:
+        while True:
+            try:
+                messages = await self._pull_subscription.fetch(
+                    batch=not_none(self._pull_consume_batch),
+                    timeout=self._pull_consume_timeout,
+                )
+            except NatsTimeoutError:
+                continue
+
+            for message in messages:
+                await self._create_worker(message, next_container)
+
+    async def _create_worker(
+        self, message: NatsMessage, next_container: NextContainer,
+    ) -> None:
+        await self._semaphore.acquire()
+        task = self._loop.create_task(self._worker(message, next_container))
+        self._workers.add(task)
+        task.add_done_callback(self._workers.discard)
+
+    async def _worker(
+        self, message: NatsMessage, next_container: NextContainer,
+    ) -> None:
+        async with next_container() as container:
+            try:
+                json_str = message.data.decode()
+                json_ = json.loads(json_str)
+                await self.body(container, *json_["args"], **json_["kwargs"])
+            except Exception as error:  # noqa: BLE001
+                await message.nak()
+
+                logger = await container.get(FilteringBoundLogger)
+                await unexpected_error_log(logger, error)
+
+                self._semaphore.release()
+            except BaseException as error:
+                await message.nak()
+                self._semaphore.release()
+                raise error from error
+            else:
+                await message.ack()
+                self._semaphore.release()
+
+
+def nats_remote[**PmT](
+    subject: str,
+    pull_subscribe: PullSubscribe,
+    max_workers: int | None = None,
+    pull_consume_batch: int | None = None,
+    pull_consume_timeout: float | None = None,
+) -> Callable[[NatsRemoteFuncBody[PmT]], NatsRemoteFunc[PmT]]:
+    def decorator(body: NatsRemoteFuncBody[PmT]) -> NatsRemoteFunc[PmT]:
+        return NatsRemoteFunc(
+            body,
+            pull_subscribe,
+            subject,
+            max_workers,
+            pull_consume_batch,
+            pull_consume_timeout,
+        )
+
+    return decorator