From 59db01259461c36798cc618b5bcba01121979fd7 Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:48:42 +0000 Subject: [PATCH] feat(cloud): enforce workspace resource quotas --- src/langbot/pkg/api/http/controller/group.py | 3 + src/langbot/pkg/api/http/service/bot.py | 36 +++- src/langbot/pkg/cloud/quotas.py | 82 ++++++++ src/langbot/pkg/plugin/connector.py | 22 ++ .../test_workspace_quota_postgres.py | 137 +++++++++++++ .../api/http/test_internal_error_responses.py | 25 +++ .../api/service/test_bot_service.py | 7 +- .../test_cloud_bot_quota_concurrency.py | 129 ++++++++++++ .../unit_tests/cloud/test_workspace_quotas.py | 102 ++++++++++ .../test_cloud_plugin_quota_concurrency.py | 191 ++++++++++++++++++ 10 files changed, 722 insertions(+), 12 deletions(-) create mode 100644 src/langbot/pkg/cloud/quotas.py create mode 100644 tests/integration/persistence/test_workspace_quota_postgres.py create mode 100644 tests/unit_tests/api/service/test_cloud_bot_quota_concurrency.py create mode 100644 tests/unit_tests/cloud/test_workspace_quotas.py create mode 100644 tests/unit_tests/plugin/test_cloud_plugin_quota_concurrency.py diff --git a/src/langbot/pkg/api/http/controller/group.py b/src/langbot/pkg/api/http/controller/group.py index 1fc09a0b2..1d229772a 100644 --- a/src/langbot/pkg/api/http/controller/group.py +++ b/src/langbot/pkg/api/http/controller/group.py @@ -14,6 +14,7 @@ from ....utils import bounded_executor from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError from ....workspace.errors import WorkspaceNotFoundError from ....cloud.entitlements import EntitlementUnavailableError +from ....cloud.quotas import WorkspaceQuotaExceededError from ....core.errors import TaskCapacityError from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext @@ -219,6 +220,8 @@ class RouterGroup(abc.ABC): return self.http_status(403, e.code, str(e)) if isinstance(e, WorkspaceCollaborationError): return self.http_status(400, e.code, str(e)) + if isinstance(e, WorkspaceQuotaExceededError): + return self.http_status(409, e.error_code, str(e)) if isinstance(e, TaskCapacityError): return self.http_status(429, 'task_capacity_exceeded', str(e)) if isinstance( diff --git a/src/langbot/pkg/api/http/service/bot.py b/src/langbot/pkg/api/http/service/bot.py index 9d9211be3..7e8080060 100644 --- a/src/langbot/pkg/api/http/service/bot.py +++ b/src/langbot/pkg/api/http/service/bot.py @@ -4,6 +4,7 @@ import uuid import sqlalchemy from ....core import app +from ....cloud.quotas import require_resource_capacity, resolve_workspace_quota from ....entity.persistence import bot as persistence_bot from ....entity.persistence import pipeline as persistence_pipeline from ....workspace.errors import WorkspaceNotFoundError @@ -101,20 +102,21 @@ class BotService: async def create_bot(self, context: TenantContext, bot_data: dict) -> str: """Create bot""" workspace_uuid = require_workspace_uuid(context) - # Check limitation limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) - max_bots = limitation.get('max_bots', -1) - if max_bots >= 0: - existing_bots = await self.get_bots(context) - if len(existing_bots) >= max_bots: - raise ValueError(f'Maximum number of bots ({max_bots}) reached') + quota = await resolve_workspace_quota( + self.ap, + workspace_uuid, + 'bots.max', + fallback=limitation.get('max_bots', -1), + ) # TODO: 检查配置信息格式 bot_data = bot_data.copy() bot_data['uuid'] = str(uuid.uuid4()) bot_data['workspace_uuid'] = workspace_uuid - # bind the most recently updated pipeline if any exist + # Preserve the legacy flat-row result shape for this optional lookup; + # quota admission and insertion below still share one transaction. result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_pipeline.LegacyPipeline), @@ -129,7 +131,25 @@ class BotService: bot_data['use_pipeline_uuid'] = pipeline.uuid bot_data['use_pipeline_name'] = pipeline.name - await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_bot.Bot).values(bot_data)) + async def persist(execute) -> None: + await require_resource_capacity( + execute, + workspace_uuid=workspace_uuid, + model=persistence_bot.Bot, + quota=quota, + resource_name='bots', + ) + + await execute(sqlalchemy.insert(persistence_bot.Bot).values(bot_data)) + + tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None) + if quota.requires_transaction_lock: + if not callable(tenant_uow): + raise RuntimeError('Cloud bot quota enforcement requires transactional persistence') + async with tenant_uow(workspace_uuid) as uow: + await persist(uow.execute) + else: + await persist(self.ap.persistence_mgr.execute_async) bot = await self.get_bot(context, bot_data['uuid'], include_secret=True) diff --git a/src/langbot/pkg/cloud/quotas.py b/src/langbot/pkg/cloud/quotas.py new file mode 100644 index 000000000..d5229d0b2 --- /dev/null +++ b/src/langbot/pkg/cloud/quotas.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Awaitable, Callable + +import sqlalchemy + +from ..entity.persistence import workspace as persistence_workspace +from .entitlements import EntitlementResolver + + +Execute = Callable[[Any], Awaitable[Any]] + + +class WorkspaceQuotaExceededError(ValueError): + """A stable business error raised when a workspace has no free slots.""" + + error_code = 'workspace_quota_exceeded' + + def __init__(self, resource_name: str, limit: int) -> None: + self.resource_name = resource_name + self.limit = limit + super().__init__(f'Maximum number of {resource_name} ({limit}) reached') + + +@dataclass(frozen=True, slots=True) +class WorkspaceQuota: + limit: int + requires_transaction_lock: bool + + +async def resolve_workspace_quota( + ap: Any, + workspace_uuid: str, + limit_name: str, + *, + fallback: int = -1, +) -> WorkspaceQuota: + """Resolve a plan-agnostic Cloud limit while preserving OSS configuration.""" + + resolver = getattr(ap, 'entitlement_resolver', None) + if isinstance(resolver, EntitlementResolver): + snapshot = await resolver.resolve(workspace_uuid) + return WorkspaceQuota( + limit=snapshot.limit(limit_name), + requires_transaction_lock=True, + ) + return WorkspaceQuota(limit=fallback, requires_transaction_lock=False) + + +async def lock_workspace_for_quota(execute: Execute, workspace_uuid: str) -> None: + """Serialize quota checks on the durable Workspace row within one transaction.""" + + result = await execute( + sqlalchemy.select(persistence_workspace.Workspace.uuid) + .where(persistence_workspace.Workspace.uuid == workspace_uuid) + .with_for_update() + ) + if result.first() is None: + raise ValueError('Workspace does not exist') + + +async def require_resource_capacity( + execute: Execute, + *, + workspace_uuid: str, + model: type, + quota: WorkspaceQuota, + resource_name: str, + workspace_locked: bool = False, +) -> None: + if quota.limit < 0: + return + if quota.requires_transaction_lock and not workspace_locked: + await lock_workspace_for_quota(execute, workspace_uuid) + result = await execute( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(model) + .where(model.workspace_uuid == workspace_uuid) + ) + if int(result.scalar_one()) >= quota.limit: + raise WorkspaceQuotaExceededError(resource_name, quota.limit) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index b82b79427..cbfb4f6a2 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -19,6 +19,11 @@ from urllib.parse import urljoin, urlparse from langbot_plugin.api.entities.builtin.pipeline.query import provider_session from ..core import app +from ..cloud.quotas import ( + lock_workspace_for_quota, + require_resource_capacity, + resolve_workspace_quota, +) from . import handler from .archive import inspect_plugin_archive_metadata from .github import ( @@ -1295,6 +1300,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): install_info: dict[str, Any], artifact_digest: str, ) -> tuple[InstallationBinding, str | None, bool]: + quota = await resolve_workspace_quota( + self.ap, + execution_context.workspace_uuid, + 'plugins.max', + ) safe_install_info = { key: value for key, value in install_info.items() @@ -1316,9 +1326,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): ) async def persist(execute): + if quota.requires_transaction_lock: + await lock_workspace_for_quota(execute, execution_context.workspace_uuid) result = await execute(statement) setting = result.first() if setting is None: + await require_resource_capacity( + execute, + workspace_uuid=execution_context.workspace_uuid, + model=persistence_plugin.PluginSetting, + quota=quota, + resource_name='plugins', + workspace_locked=quota.requires_transaction_lock, + ) installation_uuid = str(uuid.uuid4()) runtime_revision = 1 previous_digest = None @@ -1373,6 +1393,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): ) tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None) + if quota.requires_transaction_lock and not callable(tenant_uow): + raise RuntimeError('Cloud plugin quota enforcement requires transactional persistence') if callable(tenant_uow): async with tenant_uow(execution_context.workspace_uuid) as uow: return await persist(uow.execute) diff --git a/tests/integration/persistence/test_workspace_quota_postgres.py b/tests/integration/persistence/test_workspace_quota_postgres.py new file mode 100644 index 000000000..b83d089fa --- /dev/null +++ b/tests/integration/persistence/test_workspace_quota_postgres.py @@ -0,0 +1,137 @@ +"""PostgreSQL integration coverage for durable workspace quota locking. + +Run with TEST_POSTGRES_URL=postgresql+asyncpg://... pytest ... +""" + +from __future__ import annotations + +import asyncio +import os +import uuid + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + +from langbot.pkg.cloud import quotas as quota_module +from langbot.pkg.cloud.quotas import WorkspaceQuota, WorkspaceQuotaExceededError, require_resource_capacity + + +pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio] + + +class _Base(DeclarativeBase): + pass + + +class _Workspace(_Base): + __tablename__ = 'quota_integration_workspaces' + + uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True) + + +class _Resource(_Base): + __tablename__ = 'quota_integration_resources' + + uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True) + workspace_uuid: Mapped[str] = mapped_column( + sa.String(36), + sa.ForeignKey('quota_integration_workspaces.uuid', ondelete='CASCADE'), + nullable=False, + index=True, + ) + + +@pytest.fixture +async def quota_postgres(monkeypatch): + url = os.environ.get('TEST_POSTGRES_URL') + if not url: + pytest.skip('TEST_POSTGRES_URL not set') + if url.startswith('postgresql://'): + url = url.replace('postgresql://', 'postgresql+asyncpg://', 1) + + engine = create_async_engine(url, pool_size=5, max_overflow=0) + monkeypatch.setattr(quota_module.persistence_workspace, 'Workspace', _Workspace) + async with engine.begin() as connection: + await connection.run_sync(_Base.metadata.drop_all) + await connection.run_sync(_Base.metadata.create_all) + try: + yield url, engine + finally: + async with engine.begin() as connection: + await connection.run_sync(_Base.metadata.drop_all) + await engine.dispose() + + +async def test_workspace_row_lock_is_atomic_isolated_and_survives_pool_restart(quota_postgres) -> None: + url, engine = quota_postgres + workspace_a = str(uuid.uuid4()) + workspace_b = str(uuid.uuid4()) + quota = WorkspaceQuota(limit=1, requires_transaction_lock=True) + sessions = async_sessionmaker(engine, expire_on_commit=False) + + async with engine.begin() as connection: + await connection.execute(sa.insert(_Workspace), [{'uuid': workspace_a}, {'uuid': workspace_b}]) + + lock_acquired = asyncio.Event() + release_first = asyncio.Event() + + async def admit(workspace_uuid: str, *, hold: bool = False) -> None: + async with sessions() as session: + async with session.begin(): + await require_resource_capacity( + session.execute, + workspace_uuid=workspace_uuid, + model=_Resource, + quota=quota, + resource_name='resources', + ) + if hold: + lock_acquired.set() + await release_first.wait() + await session.execute( + sa.insert(_Resource).values(uuid=str(uuid.uuid4()), workspace_uuid=workspace_uuid) + ) + + first = asyncio.create_task(admit(workspace_a, hold=True)) + await asyncio.wait_for(lock_acquired.wait(), timeout=2) + same_workspace = asyncio.create_task(admit(workspace_a)) + other_workspace = asyncio.create_task(admit(workspace_b)) + + await asyncio.wait_for(other_workspace, timeout=2) + assert not same_workspace.done(), 'same-workspace transaction bypassed SELECT FOR UPDATE' + + release_first.set() + await first + with pytest.raises(WorkspaceQuotaExceededError, match=r'Maximum number of resources \(1\) reached'): + await same_workspace + + async with sessions() as session: + counts = dict( + ( + await session.execute( + sa.select(_Resource.workspace_uuid, sa.func.count()) + .group_by(_Resource.workspace_uuid) + .order_by(_Resource.workspace_uuid) + ) + ).all() + ) + assert counts == {workspace_a: 1, workspace_b: 1} + + await engine.dispose() + restarted_engine = create_async_engine(url, pool_size=2, max_overflow=0) + restarted_sessions = async_sessionmaker(restarted_engine, expire_on_commit=False) + try: + async with restarted_sessions() as session: + async with session.begin(): + with pytest.raises(WorkspaceQuotaExceededError): + await require_resource_capacity( + session.execute, + workspace_uuid=workspace_a, + model=_Resource, + quota=quota, + resource_name='resources', + ) + finally: + await restarted_engine.dispose() diff --git a/tests/unit_tests/api/http/test_internal_error_responses.py b/tests/unit_tests/api/http/test_internal_error_responses.py index 25c0e83ef..96fc024ec 100644 --- a/tests/unit_tests/api/http/test_internal_error_responses.py +++ b/tests/unit_tests/api/http/test_internal_error_responses.py @@ -9,6 +9,7 @@ import quart from langbot.pkg.api.http.controller import group from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup +from langbot.pkg.cloud.quotas import WorkspaceQuotaExceededError from langbot.pkg.utils.bounded_executor import ( BlockingWorkCapacityError, current_blocking_work_scope, @@ -48,6 +49,16 @@ class _BlockingCapacityRouterGroup(group.RouterGroup): raise BlockingWorkCapacityError('Workspace blocking executor capacity reached') +class _QuotaRouterGroup(group.RouterGroup): + name = 'quota-test' + path = '/quota-test' + + async def initialize(self) -> None: + @self.route('', methods=['POST'], auth_type=group.AuthType.NONE) + async def _(): + raise WorkspaceQuotaExceededError('bots', 2) + + class _InvalidAccountRouterGroup(group.RouterGroup): name = 'invalid-account-test' path = '/invalid-account-test' @@ -130,6 +141,20 @@ async def test_blocking_work_capacity_maps_to_retryable_http_response(): } +async def test_workspace_quota_maps_to_stable_conflict_response(): + application = SimpleNamespace(logger=Mock()) + quart_app = quart.Quart(__name__) + await _QuotaRouterGroup(application, quart_app).initialize() + + response = await quart_app.test_client().post('/quota-test') + + assert response.status_code == 409 + assert await response.get_json() == { + 'code': 'workspace_quota_exceeded', + 'msg': 'Maximum number of bots (2) reached', + } + + async def test_public_webhook_carries_scope_without_holding_database_session(): class ScopeOnlyPersistenceManager: mode = SimpleNamespace(value='cloud_runtime') diff --git a/tests/unit_tests/api/service/test_bot_service.py b/tests/unit_tests/api/service/test_bot_service.py index dea5763f2..43621b480 100644 --- a/tests/unit_tests/api/service/test_bot_service.py +++ b/tests/unit_tests/api/service/test_bot_service.py @@ -311,10 +311,9 @@ class TestBotServiceCreateBot: ap.platform_mgr = SimpleNamespace() ap.platform_mgr.load_bot = AsyncMock() - # Mock get_bots to return 2 bots already - bot1 = _create_mock_bot(bot_uuid='uuid-1') - bot2 = _create_mock_bot(bot_uuid='uuid-2') - mock_result = _create_mock_result([bot1, bot2]) + # Mock the atomic count query to report 2 existing bots. + mock_result = _create_mock_result() + mock_result.scalar_one = Mock(return_value=2) ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result) ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'}) diff --git a/tests/unit_tests/api/service/test_cloud_bot_quota_concurrency.py b/tests/unit_tests/api/service/test_cloud_bot_quota_concurrency.py new file mode 100644 index 000000000..78eab3ac2 --- /dev/null +++ b/tests/unit_tests/api/service/test_cloud_bot_quota_concurrency.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import sqlalchemy + +from langbot.pkg.api.http.service.bot import BotService +from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot + + +INSTANCE_UUID = 'cloud-instance' +WORKSPACE_A = '11111111-1111-1111-1111-111111111111' +WORKSPACE_B = '22222222-2222-2222-2222-222222222222' + + +class _Provider: + async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot: + return EntitlementSnapshot( + instance_uuid=INSTANCE_UUID, + workspace_uuid=workspace_uuid, + entitlement_revision=1, + status='active', + not_before=0, + expires_at=4_102_444_800, + features={}, + limits={'bots.max': 2}, + ) + + +class _Result: + def __init__(self, *, first=None, scalar=None) -> None: + self._first = first + self._scalar = scalar + + def first(self): + return self._first + + def scalar_one(self): + return self._scalar + + +class _TenantUow: + def __init__(self, manager: '_Persistence', workspace_uuid: str) -> None: + self.manager = manager + self.workspace_uuid = workspace_uuid + self.lock = manager.locks[workspace_uuid] + + async def __aenter__(self): + await self.lock.acquire() + return self + + async def __aexit__(self, exc_type, exc, tb): + self.lock.release() + + async def execute(self, statement): + sql = str(statement) + if isinstance(statement, sqlalchemy.sql.dml.Insert): + assert statement.table.name == 'bots' + self.manager.bots[self.workspace_uuid].append(statement.compile().params) + return _Result() + if 'FROM workspaces' in sql: + assert statement._for_update_arg is not None + self.manager.workspace_locks_seen += 1 + return _Result(first=(self.workspace_uuid,)) + if 'count(' in sql.lower() and 'FROM bots' in sql: + return _Result(scalar=len(self.manager.bots[self.workspace_uuid])) + if 'FROM legacy_pipelines' in sql: + return _Result(first=None) + raise AssertionError(f'unexpected statement: {sql}') + + +class _Persistence: + def __init__(self) -> None: + self.locks = defaultdict(asyncio.Lock) + self.bots = defaultdict(list) + self.workspace_locks_seen = 0 + + def tenant_uow(self, workspace_uuid: str) -> _TenantUow: + return _TenantUow(self, workspace_uuid) + + async def execute_async(self, statement): + assert 'FROM legacy_pipelines' in str(statement) + return _Result(first=None) + + +async def _service(manager: _Persistence) -> BotService: + resolver = EntitlementResolver(INSTANCE_UUID, _Provider()) + await resolver.reconcile_active_workspaces({WORKSPACE_A, WORKSPACE_B}) + ap = SimpleNamespace( + entitlement_resolver=resolver, + persistence_mgr=manager, + instance_config=SimpleNamespace(data={'system': {'limitation': {'max_bots': 99}}}), + platform_mgr=SimpleNamespace(load_bot=AsyncMock()), + ) + service = BotService(ap) + service.get_bot = AsyncMock(return_value={'uuid': 'created'}) + return service + + +@pytest.mark.asyncio +async def test_cloud_bot_quota_is_atomic_isolated_and_persists_across_service_restart() -> None: + manager = _Persistence() + service = await _service(manager) + + async def create(workspace_uuid: str, index: int): + return await service.create_bot(workspace_uuid, {'name': f'bot-{index}'}) + + results = await asyncio.gather( + *(create(WORKSPACE_A, index) for index in range(8)), + *(create(WORKSPACE_B, index) for index in range(8)), + return_exceptions=True, + ) + + successes = [result for result in results if isinstance(result, str)] + failures = [result for result in results if isinstance(result, ValueError)] + assert len(successes) == 4 + assert len(failures) == 12 + assert len(manager.bots[WORKSPACE_A]) == 2 + assert len(manager.bots[WORKSPACE_B]) == 2 + assert manager.workspace_locks_seen == 16 + + restarted_service = await _service(manager) + with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'): + await restarted_service.create_bot(WORKSPACE_A, {'name': 'after-restart'}) + assert len(manager.bots[WORKSPACE_A]) == 2 diff --git a/tests/unit_tests/cloud/test_workspace_quotas.py b/tests/unit_tests/cloud/test_workspace_quotas.py new file mode 100644 index 000000000..58734aeeb --- /dev/null +++ b/tests/unit_tests/cloud/test_workspace_quotas.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import sqlalchemy + +from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot +from langbot.pkg.cloud.quotas import ( + WorkspaceQuota, + require_resource_capacity, + resolve_workspace_quota, +) +from langbot.pkg.entity.persistence.bot import Bot + + +WORKSPACE_UUID = '11111111-1111-1111-1111-111111111111' +INSTANCE_UUID = 'cloud-instance' + + +class _Provider: + def __init__(self, limits: dict[str, int]) -> None: + self.limits = limits + + async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot: + return EntitlementSnapshot( + instance_uuid=INSTANCE_UUID, + workspace_uuid=workspace_uuid, + entitlement_revision=1, + status='active', + not_before=0, + expires_at=4_102_444_800, + features={}, + limits=self.limits, + plan_name='test', + ) + + +async def _resolver(limits: dict[str, int]) -> EntitlementResolver: + resolver = EntitlementResolver(INSTANCE_UUID, _Provider(limits)) + await resolver.reconcile_active_workspaces({WORKSPACE_UUID}) + return resolver + + +@pytest.mark.asyncio +async def test_resolve_workspace_quota_uses_signed_cloud_limit() -> None: + ap = SimpleNamespace(entitlement_resolver=await _resolver({'bots.max': 2})) + + quota = await resolve_workspace_quota(ap, WORKSPACE_UUID, 'bots.max', fallback=99) + + assert quota == WorkspaceQuota(limit=2, requires_transaction_lock=True) + + +@pytest.mark.asyncio +async def test_resolve_workspace_quota_preserves_oss_fallback() -> None: + quota = await resolve_workspace_quota(SimpleNamespace(), WORKSPACE_UUID, 'bots.max', fallback=7) + + assert quota == WorkspaceQuota(limit=7, requires_transaction_lock=False) + + +@pytest.mark.asyncio +async def test_require_resource_capacity_locks_workspace_before_counting() -> None: + statements: list[object] = [] + lock_result = Mock() + lock_result.first.return_value = (WORKSPACE_UUID,) + count_result = Mock() + count_result.scalar_one.return_value = 1 + execute = AsyncMock(side_effect=[lock_result, count_result]) + + await require_resource_capacity( + execute, + workspace_uuid=WORKSPACE_UUID, + model=Bot, + quota=WorkspaceQuota(limit=2, requires_transaction_lock=True), + resource_name='bots', + ) + + statements.extend(call.args[0] for call in execute.await_args_list) + assert len(statements) == 2 + assert isinstance(statements[0], sqlalchemy.sql.Select) + assert statements[0]._for_update_arg is not None + assert 'workspaces' in str(statements[0]) + assert 'count' in str(statements[1]).lower() + + +@pytest.mark.asyncio +async def test_require_resource_capacity_rejects_at_boundary() -> None: + lock_result = Mock() + lock_result.first.return_value = (WORKSPACE_UUID,) + count_result = Mock() + count_result.scalar_one.return_value = 2 + execute = AsyncMock(side_effect=[lock_result, count_result]) + + with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'): + await require_resource_capacity( + execute, + workspace_uuid=WORKSPACE_UUID, + model=Bot, + quota=WorkspaceQuota(limit=2, requires_transaction_lock=True), + resource_name='bots', + ) diff --git a/tests/unit_tests/plugin/test_cloud_plugin_quota_concurrency.py b/tests/unit_tests/plugin/test_cloud_plugin_quota_concurrency.py new file mode 100644 index 000000000..18dfd52a6 --- /dev/null +++ b/tests/unit_tests/plugin/test_cloud_plugin_quota_concurrency.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import sqlalchemy + +from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot +from langbot.pkg.cloud.quotas import WorkspaceQuotaExceededError +from langbot.pkg.plugin.connector import PluginRuntimeConnector +from langbot_plugin.runtime.plugin.mgr import PluginInstallSource + + +INSTANCE_UUID = 'cloud-instance' +WORKSPACE_A = '11111111-1111-1111-1111-111111111111' +WORKSPACE_B = '22222222-2222-2222-2222-222222222222' + + +class _Provider: + async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot: + return EntitlementSnapshot( + instance_uuid=INSTANCE_UUID, + workspace_uuid=workspace_uuid, + entitlement_revision=1, + status='active', + not_before=0, + expires_at=4_102_444_800, + features={}, + limits={'plugins.max': 3}, + ) + + +class _Result: + def __init__(self, *, first=None, scalar=None) -> None: + self._first = first + self._scalar = scalar + + def first(self): + return self._first + + def scalar_one(self): + return self._scalar + + +class _TenantUow: + def __init__(self, manager: '_Persistence', workspace_uuid: str) -> None: + self.manager = manager + self.workspace_uuid = workspace_uuid + self.lock = manager.locks[workspace_uuid] + + async def __aenter__(self): + await self.lock.acquire() + return self + + async def __aexit__(self, exc_type, exc, tb): + self.lock.release() + + async def execute(self, statement): + sql = str(statement) + params = statement.compile().params + if isinstance(statement, sqlalchemy.sql.dml.Insert): + assert statement.table.name == 'plugin_settings' + key = (params['plugin_author'], params['plugin_name']) + self.manager.plugins[self.workspace_uuid][key] = dict(params) + return _Result() + if isinstance(statement, sqlalchemy.sql.dml.Update): + return _Result() + if 'FROM workspaces' in sql: + assert statement._for_update_arg is not None + self.manager.workspace_locks_seen += 1 + return _Result(first=(self.workspace_uuid,)) + if 'count(' in sql.lower() and 'FROM plugin_settings' in sql: + return _Result(scalar=len(self.manager.plugins[self.workspace_uuid])) + if 'FROM plugin_settings' in sql: + author = next(value for name, value in params.items() if 'plugin_author' in name) + name = next(value for param, value in params.items() if 'plugin_name' in param) + row = self.manager.plugins[self.workspace_uuid].get((author, name)) + if row is None: + return _Result(first=None) + return _Result( + first=SimpleNamespace( + installation_uuid=row['installation_uuid'], + runtime_revision=row['runtime_revision'], + artifact_digest=row['artifact_digest'], + install_info=row['install_info'], + ) + ) + raise AssertionError(f'unexpected statement: {sql}') + + +class _Persistence: + def __init__(self) -> None: + self.locks = defaultdict(asyncio.Lock) + self.plugins = defaultdict(dict) + self.workspace_locks_seen = 0 + + def tenant_uow(self, workspace_uuid: str) -> _TenantUow: + return _TenantUow(self, workspace_uuid) + + +async def _connector(manager: _Persistence) -> PluginRuntimeConnector: + resolver = EntitlementResolver(INSTANCE_UUID, _Provider()) + await resolver.reconcile_active_workspaces({WORKSPACE_A, WORKSPACE_B}) + connector = object.__new__(PluginRuntimeConnector) + connector.ap = SimpleNamespace(entitlement_resolver=resolver, persistence_mgr=manager) + return connector + + +def _context(workspace_uuid: str) -> ExecutionContext: + return ExecutionContext( + instance_uuid=INSTANCE_UUID, + workspace_uuid=workspace_uuid, + placement_generation=1, + entitlement_revision=1, + ) + + +@pytest.mark.asyncio +async def test_cloud_plugin_quota_is_atomic_isolated_and_persists_across_connector_restart() -> None: + manager = _Persistence() + connector = await _connector(manager) + + async def install(workspace_uuid: str, index: int): + return await connector._persist_installation_package( + _context(workspace_uuid), + plugin_author='test-author', + plugin_name=f'plugin-{index}', + install_source=PluginInstallSource.MARKETPLACE, + install_info={'author': 'test-author', 'name': f'plugin-{index}'}, + artifact_digest=f'{index:064x}', + ) + + results = await asyncio.gather( + *(install(WORKSPACE_A, index) for index in range(10)), + *(install(WORKSPACE_B, index) for index in range(10)), + return_exceptions=True, + ) + + successes = [result for result in results if isinstance(result, tuple)] + failures = [result for result in results if isinstance(result, WorkspaceQuotaExceededError)] + assert len(successes) == 6 + assert len(failures) == 14 + assert len(manager.plugins[WORKSPACE_A]) == 3 + assert len(manager.plugins[WORKSPACE_B]) == 3 + assert manager.workspace_locks_seen == 20 + + restarted_connector = await _connector(manager) + with pytest.raises(WorkspaceQuotaExceededError, match=r'Maximum number of plugins \(3\) reached'): + await restarted_connector._persist_installation_package( + _context(WORKSPACE_A), + plugin_author='test-author', + plugin_name='after-restart', + install_source=PluginInstallSource.MARKETPLACE, + install_info={}, + artifact_digest='f' * 64, + ) + assert len(manager.plugins[WORKSPACE_A]) == 3 + + installed_name = next(iter(manager.plugins[WORKSPACE_A]))[1] + + async def reinstall(): + return await restarted_connector._persist_installation_package( + _context(WORKSPACE_A), + plugin_author='test-author', + plugin_name=installed_name, + install_source=PluginInstallSource.MARKETPLACE, + install_info={'author': 'test-author', 'name': installed_name, 'revision': 2}, + artifact_digest='e' * 64, + ) + + reinstall_results = await asyncio.gather(reinstall(), reinstall()) + assert all(result[2] is True for result in reinstall_results) + + mixed_results = await asyncio.gather( + reinstall(), + restarted_connector._persist_installation_package( + _context(WORKSPACE_A), + plugin_author='test-author', + plugin_name='new-at-capacity', + install_source=PluginInstallSource.MARKETPLACE, + install_info={}, + artifact_digest='d' * 64, + ), + return_exceptions=True, + ) + assert isinstance(mixed_results[0], tuple) + assert isinstance(mixed_results[1], WorkspaceQuotaExceededError) + assert len(manager.plugins[WORKSPACE_A]) == 3