✨ introduce rabbit exchange for reaching 0 credits (#4887)

Author: matusdrobuliak66

packages/models-library/src/models_library/rabbitmq_messages.py

@@ -2,7 +2,7 @@
 import logging
 from abc import abstractmethod
 from decimal import Decimal
-from enum import Enum, auto
+from enum import Enum, IntEnum, auto
 from typing import Any, Literal, TypeAlias
 
 import arrow
@@ -267,3 +267,29 @@ class WalletCreditsMessage(RabbitMessageBase):
 
     def routing_key(self) -> str | None:
         return f"{self.wallet_id}"
+
+
+class CreditsLimit(IntEnum):
+    MIN_CREDITS = 0
+
+
+class WalletCreditsLimitReachedMessage(RabbitMessageBase):
+    channel_name: Literal["io.simcore.service.wallets-credit-limit-reached"] = Field(
+        default="io.simcore.service.wallets-credit-limit-reached", const=True
+    )
+    created_at: datetime.datetime = Field(
+        default_factory=lambda: arrow.utcnow().datetime,
+        description="message creation datetime",
+    )
+    service_run_id: str = Field(
+        ..., description="uniquely identitifies the service run"
+    )
+    user_id: UserID
+    project_id: ProjectID
+    node_id: NodeID
+    wallet_id: WalletID
+    credits: Decimal
+    credits_limit: CreditsLimit
+
+    def routing_key(self) -> str | None:
+        return f"{self.wallet_id}.{self.credits_limit}"

scripts/release/monitor/monitor_release/settings.py

@@ -67,7 +67,7 @@ def get_settings(env_file, deployment):
             portainer_password=portainer_password,
             starts_with="production-simcore_production",
             swarm_stack_name="production-simcore",
-            portainer_endpoint_version=1,
+            portainer_endpoint_version=2,
         )
     if deployment == "aws-staging":
         portainer_url = os.getenv("AWS_STAGING_PORTAINER_URL")

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/modules/db/repositories/resource_tracker.py

@@ -174,10 +174,12 @@ async def update_service_run_stopped_at(
     async def list_service_runs_by_product_and_user_and_wallet(
         self,
         product_name: ProductName,
+        *,
         user_id: UserID | None,
         wallet_id: WalletID | None,
         offset: int,
         limit: int,
+        service_run_status: ServiceRunStatus | None = None,
     ) -> list[ServiceRunWithCreditsDB]:
         async with self.db_engine.begin() as conn:
             query = (
@@ -228,6 +230,11 @@ async def list_service_runs_by_product_and_user_and_wallet(
                 query = query.where(
                     resource_tracker_service_runs.c.wallet_id == wallet_id
                 )
+            if service_run_status:
+                query = query.where(
+                    resource_tracker_service_runs.c.service_run_status
+                    == service_run_status
+                )
 
             result = await conn.execute(query)
 
@@ -236,8 +243,10 @@ async def list_service_runs_by_product_and_user_and_wallet(
     async def total_service_runs_by_product_and_user_and_wallet(
         self,
         product_name: ProductName,
+        *,
         user_id: UserID | None,
         wallet_id: WalletID | None,
+        service_run_status: ServiceRunStatus | None = None,
     ) -> PositiveInt:
         async with self.db_engine.begin() as conn:
             query = (
@@ -252,6 +261,11 @@ async def total_service_runs_by_product_and_user_and_wallet(
                 query = query.where(
                     resource_tracker_service_runs.c.wallet_id == wallet_id
                 )
+            if service_run_status:
+                query = query.where(
+                    resource_tracker_service_runs.c.service_run_status
+                    == service_run_status
+                )
 
             result = await conn.execute(query)
         row = result.first()

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/resource_tracker_process_messages.py

@@ -6,6 +6,7 @@
 from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
 from models_library.rabbitmq_messages import (
+    CreditsLimit,
     RabbitResourceTrackingHeartbeatMessage,
     RabbitResourceTrackingMessages,
     RabbitResourceTrackingMessageType,
@@ -36,6 +37,7 @@
 from .modules.rabbitmq import RabbitMQClient, get_rabbitmq_client
 from .resource_tracker_utils import (
     make_negative,
+    publish_to_rabbitmq_wallet_credits_limit_reached,
     sum_credit_transactions_and_publish_to_rabbitmq,
 )
 
@@ -87,7 +89,7 @@ async def _process_start_event(
         user_id=msg.user_id,
         user_email=msg.user_email,
         project_id=msg.project_id,
-        project_name=msg.product_name,
+        project_name=msg.project_name,
         node_id=msg.node_id,
         node_name=msg.node_name,
         service_key=msg.service_key,
@@ -160,12 +162,21 @@ async def _process_heartbeat_event(
             update_credit_transaction
         )
         # Publish wallet total credits to RabbitMQ
-        await sum_credit_transactions_and_publish_to_rabbitmq(
+        wallet_total_credits = await sum_credit_transactions_and_publish_to_rabbitmq(
             resource_tracker_repo,
             rabbitmq_client,
             running_service.product_name,
             running_service.wallet_id,
         )
+        if wallet_total_credits.available_osparc_credits < CreditsLimit.MIN_CREDITS:
+            await publish_to_rabbitmq_wallet_credits_limit_reached(
+                resource_tracker_repo,
+                rabbitmq_client,
+                product_name=running_service.product_name,
+                wallet_id=running_service.wallet_id,
+                credits_=wallet_total_credits.available_osparc_credits,
+                credits_limit=CreditsLimit.MIN_CREDITS,
+            )
 
 
 async def _process_stop_event(

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/resource_tracker_utils.py

@@ -1,9 +1,23 @@
+import asyncio
 import logging
 from datetime import datetime, timezone
+from decimal import Decimal
 
+from models_library.api_schemas_resource_usage_tracker.credit_transactions import (
+    WalletTotalCredits,
+)
 from models_library.products import ProductName
-from models_library.rabbitmq_messages import WalletCreditsMessage
+from models_library.projects import ProjectID
+from models_library.projects_nodes_io import NodeID
+from models_library.rabbitmq_messages import (
+    CreditsLimit,
+    WalletCreditsLimitReachedMessage,
+    WalletCreditsMessage,
+)
+from models_library.resource_tracker import ServiceRunId, ServiceRunStatus
+from models_library.users import UserID
 from models_library.wallets import WalletID
+from pydantic import PositiveInt
 from servicelib.rabbitmq import RabbitMQClient
 
 from .modules.db.repositories.resource_tracker import ResourceTrackerRepository
@@ -20,7 +34,7 @@ async def sum_credit_transactions_and_publish_to_rabbitmq(
     rabbitmq_client: RabbitMQClient,
     product_name: ProductName,
     wallet_id: WalletID,
-):
+) -> WalletTotalCredits:
     wallet_total_credits = (
         await resource_tracker_repo.sum_credit_transactions_by_product_and_wallet(
             product_name,
@@ -33,3 +47,74 @@ async def sum_credit_transactions_and_publish_to_rabbitmq(
         credits=wallet_total_credits.available_osparc_credits,
     )
     await rabbitmq_client.publish(publish_message.channel_name, publish_message)
+    return wallet_total_credits
+
+
+_BATCH_SIZE = 20
+
+
+async def _publish_to_rabbitmq_wallet_credits_limit_reached(
+    rabbitmq_client: RabbitMQClient,
+    service_run_id: ServiceRunId,
+    user_id: UserID,
+    project_id: ProjectID,
+    node_id: NodeID,
+    wallet_id: WalletID,
+    credits_: Decimal,
+    credits_limit: CreditsLimit,
+):
+    publish_message = WalletCreditsLimitReachedMessage(
+        service_run_id=service_run_id,
+        user_id=user_id,
+        project_id=project_id,
+        node_id=node_id,
+        wallet_id=wallet_id,
+        credits=credits_,
+        credits_limit=credits_limit,
+    )
+    await rabbitmq_client.publish(publish_message.channel_name, publish_message)
+
+
+async def publish_to_rabbitmq_wallet_credits_limit_reached(
+    resource_tracker_repo: ResourceTrackerRepository,
+    rabbitmq_client: RabbitMQClient,
+    product_name: ProductName,
+    wallet_id: WalletID,
+    credits_: Decimal,
+    credits_limit: CreditsLimit,
+):
+    # Get all current running services for that wallet
+    total_count: PositiveInt = (
+        await resource_tracker_repo.total_service_runs_by_product_and_user_and_wallet(
+            product_name,
+            user_id=None,
+            wallet_id=wallet_id,
+            service_run_status=ServiceRunStatus.RUNNING,
+        )
+    )
+
+    for offset in range(0, total_count, _BATCH_SIZE):
+        batch_services = await resource_tracker_repo.list_service_runs_by_product_and_user_and_wallet(
+            product_name,
+            user_id=None,
+            wallet_id=wallet_id,
+            offset=offset,
+            limit=_BATCH_SIZE,
+            service_run_status=ServiceRunStatus.RUNNING,
+        )
+
+        await asyncio.gather(
+            *(
+                _publish_to_rabbitmq_wallet_credits_limit_reached(
+                    rabbitmq_client=rabbitmq_client,
+                    service_run_id=service.service_run_id,
+                    user_id=service.user_id,
+                    project_id=service.project_id,
+                    node_id=service.node_id,
+                    wallet_id=wallet_id,
+                    credits_=credits_,
+                    credits_limit=credits_limit,
+                )
+                for service in batch_services
+            )
+        )

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/resource_tracker_service_runs.py

@@ -30,32 +30,44 @@ async def list_service_runs(
     # Situation when we want to see all usage of a specific user
     if wallet_id is None and access_all_wallet_usage is None:
         total_service_runs: PositiveInt = await resource_tacker_repo.total_service_runs_by_product_and_user_and_wallet(
-            product_name, user_id, None
+            product_name, user_id=user_id, wallet_id=None
         )
         service_runs_db_model: list[
             ServiceRunWithCreditsDB
         ] = await resource_tacker_repo.list_service_runs_by_product_and_user_and_wallet(
-            product_name, user_id, None, page_params.offset, page_params.limit
+            product_name,
+            user_id=user_id,
+            wallet_id=None,
+            offset=page_params.offset,
+            limit=page_params.limit,
         )
     # Situation when accountant user can see all users usage of the wallet
     elif wallet_id and access_all_wallet_usage is True:
         total_service_runs: PositiveInt = await resource_tacker_repo.total_service_runs_by_product_and_user_and_wallet(  # type: ignore[no-redef]
-            product_name, None, wallet_id
+            product_name, user_id=None, wallet_id=wallet_id
         )
         service_runs_db_model: list[  # type: ignore[no-redef]
             ServiceRunWithCreditsDB
         ] = await resource_tacker_repo.list_service_runs_by_product_and_user_and_wallet(
-            product_name, None, wallet_id, page_params.offset, page_params.limit
+            product_name,
+            user_id=None,
+            wallet_id=wallet_id,
+            offset=page_params.offset,
+            limit=page_params.limit,
         )
     # Situation when regular user can see only his usage of the wallet
     elif wallet_id and access_all_wallet_usage is False:
         total_service_runs: PositiveInt = await resource_tacker_repo.total_service_runs_by_product_and_user_and_wallet(  # type: ignore[no-redef]
-            product_name, user_id, wallet_id
+            product_name, user_id=user_id, wallet_id=wallet_id
         )
         service_runs_db_model: list[  # type: ignore[no-redef]
             ServiceRunWithCreditsDB
         ] = await resource_tacker_repo.list_service_runs_by_product_and_user_and_wallet(
-            product_name, user_id, wallet_id, page_params.offset, page_params.limit
+            product_name,
+            user_id=user_id,
+            wallet_id=wallet_id,
+            offset=page_params.offset,
+            limit=page_params.limit,
         )
     else:
         msg = "wallet_id and access_all_wallet_usage parameters must be specified together"

services/resource-usage-tracker/tests/unit/with_dbs/test_process_rabbitmq_message_with_billing.py

@@ -2,14 +2,17 @@
 from collections.abc import Callable, Iterator
 from datetime import datetime, timezone
 from decimal import Decimal
+from unittest import mock
 
 import pytest
 import sqlalchemy as sa
 from models_library.rabbitmq_messages import (
     RabbitResourceTrackingHeartbeatMessage,
     RabbitResourceTrackingStoppedMessage,
     SimcorePlatformStatus,
+    WalletCreditsLimitReachedMessage,
 )
+from pytest_mock.plugin import MockerFixture
 from servicelib.rabbitmq import RabbitMQClient
 from simcore_postgres_database.models.resource_tracker_credit_transactions import (
     resource_tracker_credit_transactions,
@@ -34,6 +37,10 @@
     _process_start_event,
     _process_stop_event,
 )
+from tenacity._asyncio import AsyncRetrying
+from tenacity.retry import retry_if_exception_type
+from tenacity.stop import stop_after_delay
+from tenacity.wait import wait_fixed
 
 from .conftest import assert_credit_transactions_db_row
 
@@ -152,6 +159,11 @@ def resource_tracker_pricing_tables_db(postgres_db: sa.engine.Engine) -> Iterato
         con.execute(resource_tracker_credit_transactions.delete())
 
 
+@pytest.fixture
+async def mocked_message_parser(mocker: MockerFixture) -> mock.AsyncMock:
+    return mocker.AsyncMock(return_value=True)
+
+
 async def test_process_event_functions(
     rabbitmq_client: Callable[[str], RabbitMQClient],
     random_rabbit_message_start,
@@ -160,9 +172,16 @@ async def test_process_event_functions(
     resource_tracker_service_run_db,
     resource_tracker_pricing_tables_db,
     initialized_app,
+    mocked_message_parser,
 ):
     engine = initialized_app.state.engine
     publisher = rabbitmq_client("publisher")
+    consumer = rabbitmq_client("consumer")
+    await consumer.subscribe(
+        WalletCreditsLimitReachedMessage.get_channel_name(),
+        mocked_message_parser,
+        topics=["#"],
+    )
 
     msg = random_rabbit_message_start(
         wallet_id=1,
@@ -210,3 +229,12 @@ async def test_process_event_functions(
     )
     assert output.osparc_credits < first_credits_used
     assert output.transaction_status == "BILLED"
+
+    async for attempt in AsyncRetrying(
+        wait=wait_fixed(0.1),
+        stop=stop_after_delay(5),
+        retry=retry_if_exception_type(AssertionError),
+        reraise=True,
+    ):
+        with attempt:
+            mocked_message_parser.assert_called_once()