diff --git a/robosystems/operations/graph/capacity.py b/robosystems/operations/graph/capacity.py new file mode 100644 index 000000000..75e09020d --- /dev/null +++ b/robosystems/operations/graph/capacity.py @@ -0,0 +1,35 @@ +"""Tier capacity as the sale paths see it. + +Both places the platform commits money against a writer slot — checkout and +``change-tier`` — ask the same question first: is there a healthy writer on +the target tier with a free slot *right now*? Nothing on either path raises +desired capacity on purpose (the high tiers are provisioned on request), so +``scalable`` is not good enough, and any failure to determine capacity reads +as none: refuse the sale rather than collect against a slot that may not +exist. +""" + +from __future__ import annotations + +from robosystems.config import env +from robosystems.logger import logger + + +async def tier_capacity_status(tier: str) -> str: + """``ready`` when a writer for ``tier`` has a free slot; otherwise + ``at_capacity``. ``scalable`` (no slot, ASG below max) counts as + ``at_capacity`` because nothing on the sale paths raises desired capacity. + Any failure to determine capacity reads as ``at_capacity``. + """ + try: + from robosystems.middleware.graph.allocation_manager import ( + LadybugAllocationManager, + ) + from robosystems.middleware.graph.types import GraphTier + + manager = LadybugAllocationManager(environment=env.ENVIRONMENT) + status_value = await manager.check_tier_capacity(GraphTier(tier)) + except Exception as e: + logger.warning(f"Could not determine capacity for tier {tier}: {e}") + return "at_capacity" + return "ready" if status_value == "ready" else "at_capacity" diff --git a/robosystems/operations/graph/commands/tier.py b/robosystems/operations/graph/commands/tier.py index 15718d3cb..1b5519cd7 100644 --- a/robosystems/operations/graph/commands/tier.py +++ b/robosystems/operations/graph/commands/tier.py @@ -37,6 +37,7 @@ async def change_graph_tier_cmd( from robosystems.models.core import OrgRole, OrgUser from robosystems.models.core.graph import Graph from robosystems.models.core.graph.graph_credits import GraphCredits + from robosystems.operations.graph.capacity import tier_capacity_status from robosystems.operations.graph.tier_validation import ( validate_storage_capacity, validate_subgraph_count, @@ -93,6 +94,20 @@ async def change_graph_tier_cmd( if new_tier == old_tier: raise HTTPException(status_code=400, detail="Already on this tier") + # Refuse-the-sale rule, the same one checkout applies: the migration needs + # a healthy writer on the target tier with a free slot *now*. The worker + # can raise the ASG's desired capacity itself, but a cold boot does not fit + # its reattach window, and the high tiers are provisioned on request by + # policy — so refuse here, before Stripe has moved the customer's price. + if await tier_capacity_status(new_tier) != "ready": + raise HTTPException( + status_code=409, + detail=( + f"No capacity is currently available on the '{new_tier}' tier. " + "Request access from the tier picker or contact support." + ), + ) + new_price_cents = plan_config["base_price_cents"] old_price_cents = subscription.base_price_cents is_upgrade = new_price_cents > old_price_cents diff --git a/robosystems/operations/graph/tasks/graph_tier_upgrade.py b/robosystems/operations/graph/tasks/graph_tier_upgrade.py index 9163351a7..4414bde76 100644 --- a/robosystems/operations/graph/tasks/graph_tier_upgrade.py +++ b/robosystems/operations/graph/tasks/graph_tier_upgrade.py @@ -5,7 +5,8 @@ 1. mark the graph ``migrating`` and retag the volume with the new tier 2. snapshot the volume via the Volume Manager Lambda (the rollback point) -3. drain connections on the old instance +3. drain connections on the old instance — or refuse, if it cannot be + confirmed drained; the volume is never detached under possible writes 4. detach the volume — it then sits ``available`` carrying the new tier tag, which is what makes the target tier's ASG claim it 5. ensure the target ASG has capacity @@ -36,11 +37,19 @@ DRAIN_TIMEOUT_SECONDS = 120 DRAIN_POLL_INTERVAL_SECONDS = 5 +# Consecutive failed attempts before an unreachable graph API is taken to mean +# the container is stopped rather than the network blinked. +DRAIN_UNREACHABLE_CONFIRMATIONS = 3 REATTACH_TIMEOUT_SECONDS = 300 REATTACH_POLL_INTERVAL_SECONDS = 10 GRAPH_API_PORT = 8001 +class DrainRefusedError(RuntimeError): + """The old instance could not be confirmed drained, so the volume must + not be detached. Nothing has moved yet when this is raised.""" + + def _get_dynamodb_resource(): """DynamoDB resource, pointed at LocalStack in development.""" region = env.AWS_REGION @@ -312,24 +321,58 @@ async def execute(self) -> dict[str, Any]: raise async def _drain_instance(self, private_ip: str) -> None: - """Ask the old instance to close its connections before the detach. - - Bounded by ``DRAIN_TIMEOUT_SECONDS``; the migration proceeds either way, so - a slow drain delays rather than blocks the move. + """Confirm the old instance has stopped serving before the detach. + + Two outcomes count as drained: the instance reports zero active + connections, or its graph API is not listening at all, confirmed across + several attempts so a network blip does not pass for a stopped container + — the maintenance-window procedure stops the container before a tier + change precisely so nothing can be writing when the volume detaches, and + every write reaches the volume through that API. Everything else refuses: + no private IP to ask, a graph API that is up but has no drain endpoint, an + error response, or a drain that times out with connections still open. + Detaching under live writes is the one thing this step exists to prevent. """ if not private_ip: - logger.warning("No private IP for drain, skipping") - return + raise DrainRefusedError( + "No private IP recorded for the source instance; cannot confirm it " + "is drained, so the volume stays attached" + ) base_url = f"http://{private_ip}:{GRAPH_API_PORT}" async with httpx.AsyncClient(timeout=10) as client: - try: - await client.post(f"{base_url}/admin/drain") - except httpx.HTTPError as e: - logger.warning(f"Drain request failed (may be expected): {e}") + response = None + for attempt in range(1, DRAIN_UNREACHABLE_CONFIRMATIONS + 1): + try: + response = await client.post(f"{base_url}/admin/drain") + break + except httpx.HTTPError as e: + logger.warning( + f"Graph API on {private_ip} not reachable " + f"(attempt {attempt}/{DRAIN_UNREACHABLE_CONFIRMATIONS}): {e}" + ) + if attempt < DRAIN_UNREACHABLE_CONFIRMATIONS: + await asyncio.sleep(DRAIN_POLL_INTERVAL_SECONDS) + if response is None: + logger.warning( + f"Graph API on {private_ip} stayed unreachable across " + f"{DRAIN_UNREACHABLE_CONFIRMATIONS} attempts; treating the instance " + "as drained — nothing can be writing through it" + ) return + if response.status_code == 404: + raise DrainRefusedError( + f"Graph API on {private_ip} has no drain endpoint; stop the graph " + "container (maintenance window) and retry the tier change" + ) + if response.status_code >= 400: + raise DrainRefusedError( + f"Drain request to {private_ip} failed with HTTP " + f"{response.status_code}; not detaching" + ) + # Poll for connections to reach 0 elapsed = 0 while elapsed < DRAIN_TIMEOUT_SECONDS: @@ -346,7 +389,10 @@ async def _drain_instance(self, private_ip: str) -> None: await asyncio.sleep(DRAIN_POLL_INTERVAL_SECONDS) elapsed += DRAIN_POLL_INTERVAL_SECONDS - logger.warning("Drain timeout reached, proceeding with detach") + raise DrainRefusedError( + f"Drain of {private_ip} timed out after {DRAIN_TIMEOUT_SECONDS}s with " + "connections still open; not detaching" + ) async def _ensure_asg_capacity(self, asg_client: Any, tier: str) -> None: """Ensure the target tier ASG has capacity for a new instance.""" diff --git a/robosystems/routers/billing/checkout.py b/robosystems/routers/billing/checkout.py index 4adbfe8db..cfdbba8d8 100644 --- a/robosystems/routers/billing/checkout.py +++ b/robosystems/routers/billing/checkout.py @@ -18,29 +18,13 @@ from ...models.api.common import AUTHENTICATED_ERROR_RESPONSES, RESOURCE_ERROR_RESPONSES from ...models.core import User from ...models.core.billing import BillingCustomer, BillingSubscription +from ...operations.graph.capacity import ( + tier_capacity_status as _tier_capacity_status, +) from ...operations.providers.payment_provider import get_payment_provider logger = get_logger(__name__) - -async def _tier_capacity_status(plan_name: str) -> str: - """``ready`` when a writer for the tier has a free slot; otherwise - ``at_capacity``. ``scalable`` (no slot, ASG below max) counts as - ``at_capacity`` because nothing on the create path raises desired capacity. - Any failure to determine capacity reads as ``at_capacity``. - """ - try: - from ...middleware.graph.allocation_manager import LadybugAllocationManager - from ...middleware.graph.types import GraphTier - - manager = LadybugAllocationManager(environment=env.ENVIRONMENT) - status_value = await manager.check_tier_capacity(GraphTier(plan_name)) - except Exception as e: - logger.warning(f"Could not determine capacity for tier {plan_name}: {e}") - return "at_capacity" - return "ready" if status_value == "ready" else "at_capacity" - - router = APIRouter(prefix="/billing", tags=["Billing"]) diff --git a/tests/operations/graph/tasks/test_graph_tier_upgrade.py b/tests/operations/graph/tasks/test_graph_tier_upgrade.py index e7e2c1d0f..4ab740d16 100644 --- a/tests/operations/graph/tasks/test_graph_tier_upgrade.py +++ b/tests/operations/graph/tasks/test_graph_tier_upgrade.py @@ -9,6 +9,7 @@ import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest TASK_MODULE = "robosystems.operations.graph.tasks.graph_tier_upgrade" @@ -588,9 +589,14 @@ async def test_progress_reporting(self, mock_dynamodb): class TestDrainInstance: """Tests for the _drain_instance helper method.""" - async def test_drain_with_no_private_ip_skips(self): + async def test_drain_with_no_private_ip_refuses(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + task = _make_task() - await task._drain_instance("") + with pytest.raises(DrainRefusedError, match="No private IP"): + await task._drain_instance("") async def test_drain_succeeds_when_connections_reach_zero(self): task = _make_task() @@ -601,7 +607,7 @@ async def test_drain_succeeds_when_connections_reach_zero(self): with patch("httpx.AsyncClient") as MockClient: client = AsyncMock() - client.post = AsyncMock() + client.post = AsyncMock(return_value=MagicMock(status_code=202)) client.get = AsyncMock(return_value=mock_response) MockClient.return_value.__aenter__ = AsyncMock(return_value=client) MockClient.return_value.__aexit__ = AsyncMock(return_value=False) @@ -611,6 +617,92 @@ async def test_drain_succeeds_when_connections_reach_zero(self): client.post.assert_called_once() client.get.assert_called() + async def test_drain_treats_unreachable_api_as_drained(self): + # The maintenance-window procedure stops the graph container first; + # nothing can write through an API that is not listening. + task = _make_task() + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(side_effect=httpx.ConnectError("refused")) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + await task._drain_instance("10.0.1.5") + + # Unreachable is confirmed across several attempts, not taken on faith. + assert client.post.await_count == 3 + client.get.assert_not_called() + + async def test_drain_does_not_mistake_a_blip_for_a_stopped_container(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 404 + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + # First attempt fails, second reaches an API with no drain endpoint. + client.post = AsyncMock(side_effect=[httpx.ConnectError("blip"), post_response]) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="no drain endpoint"): + await task._drain_instance("10.0.1.5") + + assert client.post.await_count == 2 + client.get.assert_not_called() + + async def test_drain_refuses_when_api_has_no_drain_endpoint(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 404 + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(return_value=post_response) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="no drain endpoint"): + await task._drain_instance("10.0.1.5") + + client.get.assert_not_called() + + async def test_drain_refuses_on_timeout_with_connections_open(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 202 + poll_response = MagicMock() + poll_response.status_code = 200 + poll_response.json.return_value = {"active_connections": 3} + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(return_value=post_response) + client.get = AsyncMock(return_value=poll_response) + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="still open"): + await task._drain_instance("10.0.1.5") + + client.get.assert_called() + @pytest.mark.unit @pytest.mark.asyncio diff --git a/tests/operations/graph/test_capacity.py b/tests/operations/graph/test_capacity.py new file mode 100644 index 000000000..ebfdcbf16 --- /dev/null +++ b/tests/operations/graph/test_capacity.py @@ -0,0 +1,36 @@ +"""Tests for the shared refuse-the-sale capacity rule.""" + +from unittest.mock import AsyncMock, patch + +import pytest + +_MANAGER = "robosystems.middleware.graph.allocation_manager.LadybugAllocationManager" + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestTierCapacityStatus: + async def _status(self, tier="ladybug-large", *, value=None, raises=None): + from robosystems.operations.graph.capacity import tier_capacity_status + + with patch(_MANAGER) as manager_cls: + manager_cls.return_value.check_tier_capacity = AsyncMock( + side_effect=raises, return_value=value + ) + return await tier_capacity_status(tier) + + async def test_ready_passes_through(self): + assert await self._status(value="ready") == "ready" + + async def test_scalable_is_not_ready(self): + # Nothing on a sale path raises desired capacity; headroom is not a slot. + assert await self._status(value="scalable") == "at_capacity" + + async def test_at_capacity_passes_through(self): + assert await self._status(value="at_capacity") == "at_capacity" + + async def test_lookup_failure_reads_as_no_capacity(self): + assert await self._status(raises=RuntimeError("dynamodb down")) == "at_capacity" + + async def test_unknown_tier_reads_as_no_capacity(self): + assert await self._status(tier="not-a-tier", value="ready") == "at_capacity" diff --git a/tests/operations/graph/test_tier_command.py b/tests/operations/graph/test_tier_command.py index d5ffd4737..4525e8755 100644 --- a/tests/operations/graph/test_tier_command.py +++ b/tests/operations/graph/test_tier_command.py @@ -71,6 +71,14 @@ def _make_graph_credits(allocation=8000): GRAPH_ID = "kg1a2b3c4d5e6f78" +_CAPACITY = "robosystems.operations.graph.capacity.tier_capacity_status" + + +@pytest.fixture(autouse=True) +def target_tier_has_capacity(): + """The refuse-the-sale gate passes unless a test says otherwise.""" + with patch(_CAPACITY, new_callable=AsyncMock, return_value="ready") as gate: + yield gate @pytest.mark.asyncio @@ -230,6 +238,46 @@ async def test_same_tier_raises_400(self): assert exc.value.status_code == 400 assert "Already on this tier" in exc.value.detail + async def test_target_tier_without_capacity_raises_409( + self, target_tier_has_capacity + ): + """The refuse-the-sale rule fires before billing or Stripe move.""" + target_tier_has_capacity.return_value = "at_capacity" + membership = _make_membership() + subscription = _make_subscription() + graph = _make_graph() + db = MagicMock() + with ( + patch(f"{_CORE}.OrgUser") as mock_org_user, + patch(f"{_CORE}.OrgRole") as mock_org_role, + patch(_GRAPH) as mock_graph_cls, + patch(f"{_TIER}.BillingSubscription") as mock_sub_cls, + patch(f"{_TIER}.BillingConfig") as mock_billing_cfg, + patch(_PAYMENT) as mock_provider_factory, + patch(_WORKER, new_callable=AsyncMock) as mock_enqueue, + ): + mock_org_user.get_user_orgs.return_value = [membership] + mock_org_role.OWNER = "OWNER" + membership.role = mock_org_role.OWNER + mock_graph_cls.get_by_id.return_value = graph + mock_sub_cls.get_by_resource.return_value = subscription + mock_billing_cfg.get_subscription_plan.return_value = { + "name": "ladybug-large", + "base_price_cents": 24900, + } + with pytest.raises(HTTPException) as exc: + await self._call(new_tier="ladybug-large", db=db) + + assert exc.value.status_code == 409 + assert "ladybug-large" in exc.value.detail + target_tier_has_capacity.assert_awaited_once_with("ladybug-large") + # Nothing moved: no commit, no Stripe, no worker task. + assert subscription.status == "active" + assert graph.graph_tier == "ladybug-standard" + db.commit.assert_not_called() + mock_provider_factory.assert_not_called() + mock_enqueue.assert_not_called() + async def test_successful_upgrade_returns_operation_id(self): """Happy-path upgrade returns the operation_id string.""" membership = _make_membership()