diff --git a/docs/api.md b/docs/api.md index d282604..d49ae62 100644 --- a/docs/api.md +++ b/docs/api.md @@ -40,6 +40,10 @@ or claimed the host, `admin` for an admin key, or `null` for an unclaimed warm host. Callers cannot set it. An `Idempotency-Key` belongs to the service account that first used it. Another one reusing it gets `409`. +`POST /hosts` returns `503` with `CAPACITY` and `Retry-After` when the +provider has `HOST_LIMITS` live hosts. A live host is not in `error` and +has not expired. + ## Refresh a host secret `POST /hosts/{host_id}/secrets/{service}/refresh` makes the exchange drop diff --git a/docs/deploy.md b/docs/deploy.md index e537ca7..1b49f99 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -586,6 +586,7 @@ Core, optional: | `LEASE_DEFAULT_TTL` | `86400` | Lease TTL in seconds for hosts created without an explicit `expires_at`, and the extension applied by an empty `POST /hosts/{id}/renew`. An explicit `expires_at: null` at create time opts out of expiry entirely. | | `IDEMPOTENCY_KEY_TTL_HOURS` | `24` | Retention period for successful `Idempotency-Key` mappings. | | `POOL_SIZES` | `{}` | Warm hosts to keep ready per provider, as JSON (e.g. `{"exe": 2, "hetzner": 1}`). Overrides `POOL_SIZE` for the providers it names. | +| `HOST_LIMITS` | `{}` | Most live hosts per provider, as JSON (e.g. `{"exe": 28, "docker": 6}`). A create at the limit returns `503` with `CAPACITY`. Warm pool hosts count. A provider it does not name has no limit. | | `POOL_SIZE` | `0` | Warm hosts to keep ready for the default provider. `0` disables its pool. | | `POOL_HOST_MAX_AGE_HOURS` | `4` | Max age before the janitor reaps an unclaimed pool host. | | `POOL_MAX_CREATES_PER_TICK` | `2` | Upper bound on pool provisions per tick, across all providers; caps over-provision blast radius when ticks overlap. | diff --git a/src/core/settings.py b/src/core/settings.py index ae36726..dd278e2 100644 --- a/src/core/settings.py +++ b/src/core/settings.py @@ -165,6 +165,14 @@ class Settings(BaseSettings): "Overrides POOL_SIZE for the providers it names." ), ) + host_limits: dict[str, Annotated[int, Field(ge=0)]] = Field( + default_factory=dict, + validation_alias="HOST_LIMITS", + description=( + 'Most live hosts per provider, as JSON (e.g. {"exe": 28, "docker": 6}). ' + "A provider it does not name has no limit." + ), + ) pool_size: int = Field( default=0, ge=0, diff --git a/src/hosts/exceptions.py b/src/hosts/exceptions.py index c2c558c..034f690 100644 --- a/src/hosts/exceptions.py +++ b/src/hosts/exceptions.py @@ -1,3 +1,5 @@ +from typing import ClassVar + from core.exceptions import AppException @@ -16,6 +18,12 @@ class HostTeardownError(AppException): error_code = "HOST_TEARDOWN" +class HostCapacityError(AppException): + status_code = 503 + error_code = "CAPACITY" + headers: ClassVar[dict[str, str]] = {"Retry-After": "30"} + + class ProvisioningFailedError(AppException): status_code = 502 error_code = "PROVISIONING_FAILED" diff --git a/src/hosts/service.py b/src/hosts/service.py index 0ef6a16..0c09449 100644 --- a/src/hosts/service.py +++ b/src/hosts/service.py @@ -8,7 +8,7 @@ from typing import Any import httpx -from sqlalchemy import delete, or_, select, update +from sqlalchemy import delete, func, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from uuid6 import uuid7 @@ -20,7 +20,12 @@ from host_secrets import catalog from host_secrets.exceptions import SecretsProxyNotConfiguredError, SecretStaticError from host_secrets.placeholder import Placeholder -from hosts.exceptions import HostStateError, IdempotencyKeyConflictError, ProvisioningFailedError +from hosts.exceptions import ( + HostCapacityError, + HostStateError, + IdempotencyKeyConflictError, + ProvisioningFailedError, +) from hosts.models import Host, HostStatus, IdempotencyKey from networking.tailscale import ( DeviceDiscoveryTimeoutError, @@ -286,6 +291,25 @@ async def create_host( expires_at=initial_expires_at, pool_member=pool_member, ) + if vm.name in self.settings.host_limits: + # The lock serializes this provider's creates until the insert + # commits, so two requests cannot both pass the count. SQLite has + # no advisory locks; it serves single-process quickstarts. + if self.session.get_bind().dialect.name == "postgresql": + await self.session.execute( + select(func.pg_advisory_xact_lock(func.hashtext(vm.name))) + ) + live = ( + await self.session.execute( + select(func.count()) + .select_from(Host) + .where(Host.provider == vm.name) + .where(Host.status != HostStatus.ERROR.value) + .where(or_(Host.expires_at.is_(None), Host.expires_at > now)) + ) + ).scalar_one() + if live >= self.settings.host_limits[vm.name]: + raise HostCapacityError(f"provider {vm.name!r} is at its host limit") self.session.add(host) await self.session.commit() await self.session.refresh(host) diff --git a/src/hosts/tests/test_pool.py b/src/hosts/tests/test_pool.py index 7d7adb8..ac8d19a 100644 --- a/src/hosts/tests/test_pool.py +++ b/src/hosts/tests/test_pool.py @@ -7,7 +7,8 @@ from uuid6 import uuid7 from core.database import async_session_factory -from core.settings import get_settings +from core.settings import Settings, get_settings +from hosts.exceptions import HostCapacityError from hosts.models import Host, HostStatus from hosts.pool import _MAX_CONSECUTIVE_CREATE_FAILURES, maintain_pool from hosts.service import HostService, utc_now @@ -208,6 +209,27 @@ async def test_create_host_never_claims_another_providers_pool_member( mocked_provision.assert_awaited_once() +async def test_create_host_refuses_a_provider_at_its_host_limit(monkeypatch, stub_provider): + # The warm pool member counts toward exe's limit; stub has none. + monkeypatch.setenv("HOST_LIMITS", '{"exe": 1}') + settings = Settings() # pyright: ignore[reportCallIssue] + await _seed_pool_host(name="lb-pool-exe") + monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock()) + + async with async_session_factory() as session: + with pytest.raises(HostCapacityError): + await HostService(session, settings=settings).create_host( + env={}, image=None, provider="exe" + ) + async with async_session_factory() as session: + await HostService(session, settings=settings).create_host( + env={}, image=None, provider="stub" + ) + counted = await session.execute(select(Host.provider, func.count()).group_by(Host.provider)) + + assert {provider: count for provider, count in counted.all()} == {"exe": 1, "stub": 1} + + async def test_create_host_skips_pool_when_sized(multi_pool_settings, monkeypatch, stub_provider): # Pool members are warmed at the provider's default size, so a sized # request must provision fresh even when its provider has a warm host.