DRU-1094 -- One cloud's burst cannot take every sandbox host on a cell

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

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. |

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,

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)

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: