diff --git a/docs/api.md b/docs/api.md index 28170b8..19615ec 100644 --- a/docs/api.md +++ b/docs/api.md @@ -41,8 +41,9 @@ 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. +provider has `HOST_LIMITS` live hosts, or when the caller's service account +has `HOST_LIMIT_PER_SERVICE_ACCOUNT` live hosts. A live host is not in +`error` and has not expired. The admin key has no service account limit. `POST /hosts` returns `422` with `SECRET_MISSING` when a secret is missing. The detail is the issuer's reason, for example that the GitHub App is not diff --git a/docs/deploy.md b/docs/deploy.md index 6474e8c..2e4d98d 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -599,6 +599,7 @@ Core, optional: | `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. | +| `HOST_LIMIT_PER_SERVICE_ACCOUNT` | unset | Most live hosts per service account. A create or a warm pool claim at the limit returns `503` with `CAPACITY`. The admin key has no limit. Unset means 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 a5426ba..db32357 100644 --- a/src/core/settings.py +++ b/src/core/settings.py @@ -163,6 +163,15 @@ class Settings(BaseSettings): "A provider it does not name has no limit." ), ) + host_limit_per_service_account: int | None = Field( + default=None, + ge=1, + validation_alias="HOST_LIMIT_PER_SERVICE_ACCOUNT", + description=( + "Most live hosts per service account. Unset means no limit. The admin key " + "has none; HOST_LIMITS still bounds it." + ), + ) pool_size: int = Field( default=0, ge=0, diff --git a/src/hosts/service.py b/src/hosts/service.py index 18a458e..28190d3 100644 --- a/src/hosts/service.py +++ b/src/hosts/service.py @@ -7,7 +7,7 @@ from typing import Any import httpx -from sqlalchemy import delete, func, or_, select, update +from sqlalchemy import ColumnElement, delete, func, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from uuid6 import uuid7 @@ -19,30 +19,20 @@ from host_secrets import catalog from host_secrets.exceptions import SecretsProxyNotConfiguredError, SecretStaticError from host_secrets.placeholder import Placeholder -from hosts.exceptions import ( - HostCapacityError, - HostKeyScanError, - HostStateError, - IdempotencyKeyConflictError, - ProvisioningFailedError, -) +from hosts import exceptions from hosts.models import Host, HostStatus, IdempotencyKey from networking.tailscale import ( DeviceDiscoveryTimeoutError, NetworkError, Tailscale, ) +from providers import exceptions as provider_exceptions from providers.base import VMProvider from providers.environment import get_persist -from providers.exceptions import ( - ProviderError, - ProviderNotFoundError, - UnknownProviderError, - UnsupportedSizingError, -) from providers.registry import get_provider_names, get_vm_provider from secrets_exchange.client import SecretsExchange from secrets_exchange.secrets import IssuerError, MissingSecretError, Secret +from service_accounts.models import ServiceAccount from templates.exceptions import TemplateNotAvailableError, UnknownTemplateError from templates.models import Template, TemplateStatus @@ -111,13 +101,26 @@ async def get_or_create_host( registered = get_provider_names() if provider not in registered: available = ", ".join(sorted(registered)) - raise UnknownProviderError(f"unknown provider {provider!r}; available: {available}") + raise provider_exceptions.UnknownProviderError( + f"unknown provider {provider!r}; available: {available}" + ) if idempotency_key and ( existing := await self._lookup_idempotency_key(idempotency_key, service_account) ): return existing + limit = self.settings.host_limit_per_service_account + if limit and service_account not in (None, ServiceAccount.ADMIN): + # The warm claim or the insert below commits the transaction that holds the lock. + live = await self._count_live_hosts( + f"service_account:{service_account}", Host.service_account == service_account + ) + if live >= limit: + raise exceptions.HostCapacityError( + f"service account {service_account!r} is at its host limit" + ) + host: Host | None = None # Pool members are warmed at the provider's defaults, so only an # uncustomized request can claim one. @@ -154,7 +157,7 @@ async def get_or_create_host( await self._release_idempotency_loser(host) winner = await self._lookup_idempotency_key(idempotency_key, service_account) if not winner: - raise HostStateError("idempotency race could not be resolved") + raise exceptions.HostStateError("idempotency race could not be resolved") return winner return host @@ -204,6 +207,23 @@ async def _try_claim_pool_host( logger.info("pool: claimed host_id=%s name=%s", host.id, host.name) return host + async def _count_live_hosts(self, key: str, condition: ColumnElement[bool]) -> int: + """The hosts that match ``condition``, are not in error, and have not expired. + The advisory lock on ``key`` holds until the transaction commits, so two creates + cannot both pass the same 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(key)))) + return ( + await self.session.execute( + select(func.count()) + .select_from(Host) + .where(condition) + .where(Host.status != HostStatus.ERROR.value) + .where(or_(Host.expires_at.is_(None), Host.expires_at > utc_now())) + ) + ).scalar_one() + async def create_host( self, *, @@ -223,11 +243,11 @@ async def create_host( # never claims its own members. vm = get_vm_provider(provider) if instance_type and not vm.supports_instance_type: - raise UnsupportedSizingError( + raise provider_exceptions.UnsupportedSizingError( f"provider {vm.name!r} does not support a per-request instance_type" ) if disk_gb and not vm.supports_disk_gb: - raise UnsupportedSizingError( + raise provider_exceptions.UnsupportedSizingError( f"provider {vm.name!r} does not support a per-request disk_gb" ) proxy = self.settings.secrets_proxy_url and self.settings.secrets_proxy_ca_file @@ -272,24 +292,9 @@ async def create_host( 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() + live = await self._count_live_hosts(f"provider:{vm.name}", Host.provider == vm.name) if live >= self.settings.host_limits[vm.name]: - raise HostCapacityError(f"provider {vm.name!r} is at its host limit") + raise exceptions.HostCapacityError(f"provider {vm.name!r} is at its host limit") self.session.add(host) await self.session.commit() await self.session.refresh(host) @@ -298,7 +303,7 @@ async def create_host( await self.session.refresh(host) if host.status == HostStatus.ERROR.value: - raise ProvisioningFailedError(host.last_error) + raise exceptions.ProvisioningFailedError(host.last_error) # A dedicated session, so self.session's transaction stays short for # callers that hold an advisory lock. The match on the in-flight value @@ -326,7 +331,7 @@ async def _lookup_idempotency_key(self, key: str, service_account: str | None) - ) if host: if host.service_account != service_account: - raise IdempotencyKeyConflictError( + raise exceptions.IdempotencyKeyConflictError( f"idempotency key {key} belongs to another service account" ) return host @@ -402,10 +407,10 @@ async def renew_host(self, host_id: uuid.UUID, *, expires_at: datetime | None = raise ResourceNotFoundError("host not found") if host.pool_member and not host.claimed_at: - raise HostStateError("unclaimed pool host is managed by pool maintenance") + raise exceptions.HostStateError("unclaimed pool host is managed by pool maintenance") if host.status not in RENEWABLE_STATUSES: - raise HostStateError(f"cannot renew a host in status {host.status}") + raise exceptions.HostStateError(f"cannot renew a host in status {host.status}") host.expires_at = get_vm_provider(host.provider).lease_end( host.created_at, expires_at or self._default_lease_expires_at() @@ -452,7 +457,7 @@ async def delete_host( return False if not force and host.status in DELETE_BLOCKED_STATUSES: - raise HostStateError("host is still provisioning") + raise exceptions.HostStateError("host is still provisioning") if force or host.status in VM_BACKED_STATUSES: # An abandoned provision stranded in CREATING_VM may already have a @@ -470,7 +475,7 @@ async def delete_host( await vm.secrets.delete_secrets(vm=host.name) try: await vm.delete_vm(host.name) - except ProviderNotFoundError: + except provider_exceptions.ProviderNotFoundError: # Evicted by the provider, or a previous delete got this far. logger.warning( "host VM already absent at provider during teardown: " @@ -492,7 +497,7 @@ async def provision(self, host: Host) -> None: if vm.tailnet and not self.tailscale: await self.mark_failed( host, - ProvisioningFailedError( + exceptions.ProvisioningFailedError( f"provider {vm.name!r} joins the tailnet; set TAILSCALE_TAILNET, " "TAILSCALE_AUTH_TAGS, TAILSCALE_OAUTH_CLIENT_ID and " "TAILSCALE_OAUTH_CLIENT_SECRET" @@ -520,7 +525,7 @@ async def provision(self, host: Host) -> None: except MissingSecretError as exc: await self.mark_failed(host, exc) raise - except (IssuerError, ProviderError, ValueError) as exc: + except (IssuerError, provider_exceptions.ProviderError, ValueError) as exc: await self.mark_failed(host, exc) return @@ -533,7 +538,7 @@ async def provision(self, host: Host) -> None: # Without a gateway address the host's SSH coordinates would be dead. await self.mark_failed( host, - ProvisioningFailedError( + exceptions.ProvisioningFailedError( f"provider {vm.name!r} requires the SSH gateway; set GATEWAY_SSH_HOST" ), ) @@ -547,7 +552,7 @@ async def provision(self, host: Host) -> None: instance_type=host.instance_type, disk_gb=host.disk_gb, ) - except ProviderError as exc: + except provider_exceptions.ProviderError as exc: await self.mark_failed(host, exc) return @@ -587,7 +592,7 @@ async def provision(self, host: Host) -> None: try: known_hosts_data = await self.scan_known_hosts(host) - except HostKeyScanError as exc: + except exceptions.HostKeyScanError as exc: await self.mark_failed(host, exc) return @@ -661,7 +666,9 @@ async def scan_known_hosts(self, host: Host) -> bytes: stderr=asyncio.subprocess.PIPE, ) except OSError as error: - raise HostKeyScanError(f"could not run ssh-keyscan: {error}") from error + raise exceptions.HostKeyScanError( + f"could not run ssh-keyscan: {error}" + ) from error stdout, stderr = await process.communicate() collected += stdout errors.append(stderr.decode().strip()) @@ -671,5 +678,5 @@ async def scan_known_hosts(self, host: Host) -> bytes: # read nothing during the gap. if time.monotonic() >= deadline: detail = "; ".join(filter(None, errors)) or "empty output" - raise HostKeyScanError(f"ssh-keyscan never returned host keys: {detail}") + raise exceptions.HostKeyScanError(f"ssh-keyscan never returned host keys: {detail}") await asyncio.sleep(0.5) diff --git a/src/hosts/tests/test_host_service_account.py b/src/hosts/tests/test_host_service_account.py index 73dd72c..62928fa 100644 --- a/src/hosts/tests/test_host_service_account.py +++ b/src/hosts/tests/test_host_service_account.py @@ -54,7 +54,7 @@ async def test_pool_claim_records_service_account_and_release_clears_it( async with async_session_factory() as session: service = HostService(session) pool_host = await service.create_host(env={}, image=None, pool_member=True) - assert pool_host.service_account is None + assert not pool_host.service_account pool_host.status = HostStatus.ACTIVE.value await session.commit() pool_id = pool_host.id @@ -70,7 +70,24 @@ async def test_pool_claim_records_service_account_and_release_clears_it( await HostService(session)._release_idempotency_loser(pool_host) released = await client.get(f"/hosts/{pool_id}", headers=headers) - assert released.json()["service_account"] is None + assert not released.json()["service_account"] + + +async def test_a_service_account_at_its_host_limit_waits( + client: AsyncClient, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(get_settings(), "host_limit_per_service_account", 1) + account_a = await issue(client, "account-a") + account_b = await issue(client, "account-b") + assert (await client.post("/hosts", headers=account_a)).status_code == 201 + + refused = await client.post("/hosts", headers=account_a) + + assert refused.status_code == 503 + assert refused.json()["error_code"] == "CAPACITY" + assert refused.headers["Retry-After"] == "30" + assert (await client.post("/hosts", headers=account_b)).status_code == 201 + assert (await client.post("/hosts", headers=ADMIN)).status_code == 201 async def test_idempotency_key_belongs_to_the_caller_that_used_it(client: AsyncClient) -> None: