From 4ab62a222c30729ab4db848b7db4044b72c1f105 Mon Sep 17 00:00:00 2001 From: Paulo Date: Mon, 19 Oct 2026 21:41:36 +0200 Subject: [PATCH] Read HostService as its end state The template lookup lives in create_host, its one caller. The idempotency lookup is one positive branch. Comments that narrated the code, the caller or a provider are gone; the ones that stay say why. --- src/hosts/exceptions.py | 4 + src/hosts/service.py | 223 +++++++++++++-------------------- src/hosts/tests/test_api.py | 22 ++-- src/hosts/tests/test_models.py | 29 +++-- 4 files changed, 115 insertions(+), 163 deletions(-) diff --git a/src/hosts/exceptions.py b/src/hosts/exceptions.py index 034f690..9c638a3 100644 --- a/src/hosts/exceptions.py +++ b/src/hosts/exceptions.py @@ -27,3 +27,7 @@ class HostCapacityError(AppException): class ProvisioningFailedError(AppException): status_code = 502 error_code = "PROVISIONING_FAILED" + + +class HostKeyScanError(Exception): + pass diff --git a/src/hosts/service.py b/src/hosts/service.py index 13c9909..a2ab03d 100644 --- a/src/hosts/service.py +++ b/src/hosts/service.py @@ -22,6 +22,7 @@ from host_secrets.placeholder import Placeholder from hosts.exceptions import ( HostCapacityError, + HostKeyScanError, HostStateError, IdempotencyKeyConflictError, ProvisioningFailedError, @@ -117,16 +118,14 @@ async def get_or_create_host( available = ", ".join(sorted(registered)) raise UnknownProviderError(f"unknown provider {provider!r}; available: {available}") - if idempotency_key: - existing = await self._lookup_idempotency_key(idempotency_key, service_account) - if existing: - return existing + if idempotency_key and ( + existing := await self._lookup_idempotency_key(idempotency_key, service_account) + ): + return existing host: Host | None = None - # Warm hosts are provider-specific, so the claim is scoped to the - # requested provider's pool. A request is pool-eligible only when it - # does not customize the host: no image, template, env, secrets, or - # per-request sizing — pool members are warmed at the provider's defaults. + # Pool members are warmed at the provider's defaults, so only an + # uncustomized request can claim one. requested_provider = provider or self.settings.default_host_provider customized = env or secrets or image or template or instance_type or disk_gb if not customized and self.settings.get_pool_targets().get(requested_provider): @@ -156,7 +155,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") from None + raise HostStateError("idempotency race could not be resolved") return winner return host @@ -222,9 +221,8 @@ async def create_host( disk_gb: int | None = None, pool_member: bool = False, ) -> Host: - # Always provisions a brand-new VM; the pool maintainer calls this - # directly (with pool_member=True) so it never recursively claims its - # own pool members. + # The pool maintainer calls this, not get_or_create_host, so a top-up + # never claims its own members. vm = get_vm_provider(provider) if instance_type and not vm.supports_instance_type: raise UnsupportedSizingError( @@ -241,11 +239,20 @@ async def create_host( "sandboxes dial and its certificate" ) if template: - image = await self._resolve_template_image(template_id=template, provider=vm.name) + found = await self.session.get(Template, template) + if not found or found.provider != vm.name: + raise UnknownTemplateError( + f"template {template} not found for provider {vm.name!r}" + ) + if found.status != TemplateStatus.AVAILABLE.value: + detail = f"template {found.id} is {found.status}" + if found.status == TemplateStatus.FAILED.value: + detail = f"{detail}: {found.last_error}" + raise TemplateNotAvailableError(detail) + found.last_used_at = utc_now() + image = found.image uid = uuid7() - name = Host.build_name(uid) now = utc_now() - host_image = image or vm.default_image # Safety TTL covers the strand window: if the client disconnects # mid-provision, this is what makes the janitor reap the row + VM. # Replaced with the caller's value after provisioning succeeds. A @@ -262,9 +269,9 @@ async def create_host( service_account=service_account, env=env, secrets=secrets or {}, - name=name, + name=Host.build_name(uid), provider=vm.name, - image=host_image, + image=image or vm.default_image, instance_type=instance_type, disk_gb=disk_gb, status=HostStatus.PROVISIONING.value, @@ -296,11 +303,11 @@ async def create_host( await self.session.commit() await self.session.refresh(host) - await self.provision(str(host.id)) + await self.provision(host) await self.session.refresh(host) if host.status == HostStatus.ERROR.value: - raise ProvisioningFailedError(host.last_error or "provisioning failed") + raise ProvisioningFailedError(host.last_error) # Provisioning won: replace the safety TTL with the caller's intent # in a dedicated session so we don't extend ``self.session``'s @@ -321,57 +328,30 @@ async def create_host( await self.session.refresh(host) return host - async def _resolve_template_image(self, *, template_id: uuid.UUID, provider: str) -> str: - result = await self.session.execute( - select(Template).where(Template.id == template_id).where(Template.provider == provider) - ) - template = result.scalar_one_or_none() - - if not template: - raise UnknownTemplateError( - f"template {template_id} not found for provider {provider!r}" - ) - - if template.status != TemplateStatus.AVAILABLE.value: - detail = f"template {template.id} is {template.status}" - if template.status == TemplateStatus.FAILED.value: - detail = f"{detail}: {template.last_error}" - raise TemplateNotAvailableError(detail) - - template.last_used_at = utc_now() - return template.image - async def _lookup_idempotency_key(self, key: str, service_account: str | None) -> Host | None: - record = ( + if record := ( await self.session.execute(select(IdempotencyKey).where(IdempotencyKey.key == key)) - ).scalar_one_or_none() - - if not record: - return - - host = ( - await self.session.get(Host, record.host_id) if record.expires_at > utc_now() else None - ) - if host: - if host.service_account != service_account: - raise IdempotencyKeyConflictError( - f"idempotency key {key} belongs to another service account" - ) - return host - # Stale: expired, or the host vanished without the FK cascade firing. - # GC in a dedicated session so we don't autoflush the caller's pending - # state on `self.session`. - async with async_session_factory() as gc_session: - await gc_session.execute(delete(IdempotencyKey).where(IdempotencyKey.key == key)) - await gc_session.commit() - return + ).scalar_one_or_none(): + host = ( + await self.session.get(Host, record.host_id) + if record.expires_at > utc_now() + else None + ) + if host: + if host.service_account != service_account: + raise IdempotencyKeyConflictError( + f"idempotency key {key} belongs to another service account" + ) + return host + # A stale key, expired or orphaned, goes in a dedicated session so + # the caller's pending state on self.session is not flushed. + async with async_session_factory() as gc_session: + await gc_session.execute(delete(IdempotencyKey).where(IdempotencyKey.key == key)) + await gc_session.commit() async def _record_idempotency_key(self, key: str, host: Host) -> bool: - """Persist the key→host mapping; return False if a concurrent request won. - - Runs in a dedicated session so a losing-race UNIQUE violation can't roll - back (and expire) the request's main session mid-response. - """ + """False when a concurrent request won. A dedicated session, so the lost + race's UNIQUE violation cannot roll back the request's own session.""" now = utc_now() async with async_session_factory() as record_session: record_session.add( @@ -389,9 +369,8 @@ async def _record_idempotency_key(self, key: str, host: Host) -> bool: return True async def _release_idempotency_loser(self, host: Host) -> None: - # Two shapes of loser: claimed pool host → return to pool with a - # fresh max-age TTL; freshly-created host → mark expired so the - # janitor reaps it (delete_host refuses PROVISIONING). + # A fresh loser is left to the janitor; delete_host refuses a host + # that is still provisioning. async with async_session_factory() as fix_session: if fresh := await fix_session.get(Host, host.id): now = utc_now() @@ -478,28 +457,23 @@ async def delete_host( raise ResourceNotFoundError("host not found") if pool_shed and host.claimed_at: - # A caller claimed this host between the maintainer selecting it as - # excess and this locked read — leave it for its owner, don't reap it. + # Claimed since the maintainer picked it as excess. return False if expired_only and (not host.expires_at or host.expires_at > utc_now()): - # The owner renewed this host between the janitor selecting it as - # expired and this locked read — the lease is live again, spare it. + # Renewed since the janitor picked it as expired. return False if not force and host.status in DELETE_BLOCKED_STATUSES: raise HostStateError("host is still provisioning") if force or host.status in VM_BACKED_STATUSES: - # force is the janitor reaping an abandoned provision: attempt - # teardown even from an early state, since a row stranded in - # CREATING_VM may already have a VM (delete_vm no-ops if it doesn't). + # An abandoned provision stranded in CREATING_VM may already have a + # VM; delete_vm no-ops without one. if host.tailscale_device_id and self.tailscale: - # Clear and commit the device_id before deleting the VM: - # a later delete_vm transport error must not retry the - # already-completed release. Hosts provisioned on the - # tailnet but reaped after the operator removed the Tailscale - # settings fall through and let the auth-key TTL expire the device. + # The cleared device_id commits before the VM goes, so a failed + # delete_vm does not repeat the release. Without Tailscale + # settings the auth-key TTL expires the device. await self.tailscale.release_device(host.tailscale_device_id) host.tailscale_device_id = None host.updated_at = utc_now() @@ -510,10 +484,7 @@ async def delete_host( try: await vm.delete_vm(host.name) except ProviderNotFoundError: - # VM already absent at the provider — exe.dev may have evicted - # it, or a previous delete partially succeeded. Treat as done - # so we can clean up the DB row, but log so unexpected - # evictions are visible. + # Evicted by the provider, or a previous delete got this far. logger.warning( "host VM already absent at provider during teardown: " "host_id=%s name=%s provider=%s", @@ -525,12 +496,7 @@ async def delete_host( await self.session.commit() return True - async def provision(self, host_id: str) -> None: - host = await self.get_host(uuid.UUID(host_id)) - - if not host: - raise ResourceNotFoundError("host not found") - + async def provision(self, host: Host) -> None: host.status = HostStatus.CREATING_NETWORK.value host.updated_at = utc_now() await self.session.commit() @@ -551,8 +517,8 @@ async def provision(self, host_id: str) -> None: environment = dict(host.env) setup_script: str | None = None if tailscale: - # The bootstrap script hard-requires TAILSCALE_AUTHKEY; only - # deliver it (and mint a key) when the host joins the tailnet. + # The bootstrap script needs TAILSCALE_AUTHKEY, so only a tailnet + # host runs it. try: join_credentials = await tailscale.issue_join_credentials(host_name=host.name) except NetworkError as exc: @@ -577,9 +543,7 @@ async def provision(self, host_id: str) -> None: gateway = GatewaySettings() if vm.gateway_process_class and not gateway.ssh_host: - # A gateway-provider host is reachable only through the gateway; - # provisioning one without an address would hand out dead - # coordinates. + # Without a gateway address the host's SSH coordinates would be dead. await self.mark_failed( host, ProvisioningFailedError( @@ -611,10 +575,8 @@ async def provision(self, host_id: str) -> None: host.external_ssh_host = gateway.ssh_host host.external_ssh_port = gateway.ssh_port host.ssh_username = host.name - # Stamp the per-VM key onto this instance so the POST response - # carries it. There's no column behind `private_key`, so a later - # GET that loads a fresh row sees the class default (None) and - # never echoes the key back. + # private_key has no column: this instance carries it into the POST + # response, and a later GET never sees it. host.private_key = vm_result.private_key if tailscale: host.internal_ssh_host = tailscale.build_ssh_host(host.name) @@ -638,7 +600,7 @@ async def provision(self, host_id: str) -> None: try: known_hosts_data = await self.scan_known_hosts(host) - except RuntimeError as exc: + except HostKeyScanError as exc: await self.mark_failed(host, exc) return @@ -680,8 +642,7 @@ async def mark_failed(self, host: Host, exc: Exception) -> None: host.status, ) host.status = HostStatus.ERROR.value - # Client-safe summary, not the raw traceback: last_error is echoed back - # to callers, while the full traceback stays in the log above. + # last_error reaches callers; the traceback stays in the log above. host.last_error = f"{type(exc).__name__}: {exc}" now = utc_now() # An errored host is dead weight (its VM may be half-created). Expire it @@ -691,47 +652,37 @@ async def mark_failed(self, host: Host, exc: Exception) -> None: await self.session.commit() async def scan_known_hosts(self, host: Host) -> bytes: - # Scan every reachable address. tailscaled-SSH (internal) and the - # provider's edge sshd (external) present different host keys, so - # callers picking either path need both entries to verify. Each address - # carries its own port: the internal path is always 22 by Tailscale - # convention, while the external sshd may be remapped (e.g. a published - # container port), so they're scanned separately. + # The tailnet sshd and the provider's sshd present different keys, and + # the tailnet one always listens on 22. targets: list[tuple[str, int]] = [] if host.internal_ssh_host: targets.append((host.internal_ssh_host, 22)) if host.external_ssh_host: targets.append((host.external_ssh_host, host.external_ssh_port)) - timeout = get_vm_provider(host.provider).bootstrap_ssh_timeout_seconds - deadline = time.monotonic() + timeout - last_detail = "no attempt made" + deadline = time.monotonic() + get_vm_provider(host.provider).bootstrap_ssh_timeout_seconds while True: - scans = [await self._keyscan(ssh_host, ssh_port) for ssh_host, ssh_port in targets] - collected = b"".join(stdout for stdout, _ in scans) + collected = b"" + errors: list[str] = [] + for ssh_host, ssh_port in targets: + try: + process = await asyncio.create_subprocess_exec( + "ssh-keyscan", + "-p", + str(ssh_port), + ssh_host, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as error: + raise HostKeyScanError(f"could not run ssh-keyscan: {error}") from error + stdout, stderr = await process.communicate() + collected += stdout + errors.append(stderr.decode().strip()) if all(ssh_host.encode() in collected for ssh_host, _ in targets): return collected # Tailscaled-SSH lags device discovery; keyscan can connect but - # read nothing during the gap. Retry within the budget. - last_detail = "; ".join(error for _, error in scans if error) or "empty output" + # read nothing during the gap. if time.monotonic() >= deadline: - raise RuntimeError(f"ssh-keyscan never returned host keys: {last_detail}") + detail = "; ".join(filter(None, errors)) or "empty output" + raise HostKeyScanError(f"ssh-keyscan never returned host keys: {detail}") await asyncio.sleep(0.5) - - @staticmethod - async def _keyscan(ssh_host: str, ssh_port: int) -> tuple[bytes, str]: - try: - process = await asyncio.create_subprocess_exec( - "ssh-keyscan", - "-p", - str(ssh_port), - ssh_host, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - except OSError as error: - # ssh-keyscan missing from the image (or otherwise unspawnable) raises - # here; translate to the RuntimeError provision() routes through - # mark_failed, so it surfaces as a provisioning failure, not a 500. - raise RuntimeError(f"could not run ssh-keyscan: {error}") from error - stdout, stderr = await process.communicate() - return stdout, stderr.decode().strip() diff --git a/src/hosts/tests/test_api.py b/src/hosts/tests/test_api.py index 3bc8392..c7b33d4 100644 --- a/src/hosts/tests/test_api.py +++ b/src/hosts/tests/test_api.py @@ -105,7 +105,7 @@ async def test_create_host_creates_record_and_enqueues_provisioning(client, monk mocked_provision.assert_awaited_once() assert mocked_provision.await_args is not None - assert mocked_provision.await_args.args == (str(host.id),) + assert mocked_provision.await_args.args[0].id == host.id async def test_create_host_persists_env_without_returning_it(client, monkeypatch): @@ -136,7 +136,7 @@ async def test_create_host_persists_env_without_returning_it(client, monkeypatch mocked_provision.assert_awaited_once() assert mocked_provision.await_args is not None - assert mocked_provision.await_args.args == (payload["id"],) + assert mocked_provision.await_args.args[0].id == host.id async def test_create_host_stores_and_returns_custom_image(client, monkeypatch): @@ -228,11 +228,8 @@ async def test_create_host_keeps_safety_ttl_while_provisioning(monkeypatch): # the lease itself is stamped once the host is usable. inflight = {} - async def capture_inflight_expiry(self, host_id): - async with async_session_factory() as session: - row = await session.get(Host, uuid.UUID(host_id)) - assert row is not None - inflight["expires_at"] = row.expires_at + async def capture_inflight_expiry(self, host): + inflight["expires_at"] = host.expires_at monkeypatch.setattr("hosts.service.HostService.provision", capture_inflight_expiry) grace = timedelta(seconds=get_settings().provisioning_grace_seconds) @@ -254,15 +251,12 @@ async def test_create_host_preserves_renewal_made_during_provisioning(monkeypatc # lease when provisioning finishes. renewal = {} - async def renew_mid_provision(self, host_id): - async with async_session_factory() as session: - row = await session.get(Host, uuid.UUID(host_id)) - assert row is not None - row.status = HostStatus.BOOTSTRAPPING.value - await session.commit() + async def renew_mid_provision(self, host): + host.status = HostStatus.BOOTSTRAPPING.value + await self.session.commit() async with async_session_factory() as session: renewed = await HostService(session).renew_host( - uuid.UUID(host_id), expires_at=utc_now() + timedelta(days=7) + host.id, expires_at=utc_now() + timedelta(days=7) ) renewal["expires_at"] = renewed.expires_at diff --git a/src/hosts/tests/test_models.py b/src/hosts/tests/test_models.py index 7db1355..fb24f25 100644 --- a/src/hosts/tests/test_models.py +++ b/src/hosts/tests/test_models.py @@ -8,7 +8,7 @@ from uuid6 import uuid7 from core.database import async_session_factory -from hosts.exceptions import ProvisioningFailedError +from hosts.exceptions import HostKeyScanError, ProvisioningFailedError from hosts.models import Host, HostStatus from hosts.service import HostService, utc_now from networking.tailscale import ( @@ -88,7 +88,8 @@ async def test_provision_walks_host_to_active(monkeypatch): mocks = await _patch_provision_happy_path(monkeypatch) async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) async with async_session_factory() as session: refreshed = await session.get(Host, host.id) @@ -139,7 +140,8 @@ async def test_provision_forwards_sizing_from_the_row_to_create_vm(monkeypatch): mocks = await _patch_provision_happy_path(monkeypatch) async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) assert mocks["create_vm"].await_args is not None create_vm_kwargs = mocks["create_vm"].await_args.kwargs @@ -247,7 +249,8 @@ async def test_provision_generated_authkey_overrides_caller_env(monkeypatch): mocks = await _patch_provision_happy_path(monkeypatch) async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) assert mocks["create_vm"].await_args is not None assert mocks["create_vm"].await_args.kwargs["env"]["TAILSCALE_AUTHKEY"] == "tskey-secret" @@ -259,7 +262,8 @@ async def test_provision_marks_error_when_provider_rejects_credentials(monkeypat mocks["create_vm"].side_effect = ProviderAuthError("exe.dev authentication failed") async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) async with async_session_factory() as session: refreshed = await session.get(Host, host.id) @@ -276,7 +280,8 @@ async def test_provision_marks_error_when_device_discovery_times_out(monkeypatch ) async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) async with async_session_factory() as session: refreshed = await session.get(Host, host.id) @@ -290,10 +295,11 @@ async def test_provision_marks_error_when_device_discovery_times_out(monkeypatch async def test_provision_marks_error_when_keyscan_fails(monkeypatch): host = await create_host_record(name="lb-sandbox-test", status="provisioning") mocks = await _patch_provision_happy_path(monkeypatch) - mocks["scan_known_hosts"].side_effect = RuntimeError("keyscan failed") + mocks["scan_known_hosts"].side_effect = HostKeyScanError("keyscan failed") async with async_session_factory() as session: - await HostService(session).provision(str(host.id)) + session.add(host) + await HostService(session).provision(host) async with async_session_factory() as session: refreshed = await session.get(Host, host.id) @@ -343,10 +349,7 @@ async def test_scan_known_hosts_retries_until_banner_is_returned(monkeypatch): assert create_subprocess_exec.await_count == 2 -async def test_scan_known_hosts_translates_spawn_failure_to_runtime_error(monkeypatch): - # ssh-keyscan missing from the runtime image raises FileNotFoundError at spawn; - # it must surface as RuntimeError (which provision routes through mark_failed), - # not escape as a generic server error. +async def test_scan_known_hosts_fails_when_ssh_keyscan_cannot_run(monkeypatch): host = Host( name="sb-test", image=ExeSettings().default_image, # pyright: ignore[reportCallIssue] @@ -363,7 +366,7 @@ async def test_scan_known_hosts_translates_spawn_failure_to_runtime_error(monkey async with async_session_factory() as session: service = HostService(session) - with pytest.raises(RuntimeError, match="could not run ssh-keyscan"): + with pytest.raises(HostKeyScanError, match="could not run ssh-keyscan"): await service.scan_known_hosts(host)