Strengthen AND provisioning leases and reconciliation

Author: aidankhogg

migrations/001_initial.sql

@@ -29,7 +29,8 @@ CREATE TABLE IF NOT EXISTS address_pools (
 CREATE TABLE IF NOT EXISTS address_leases (
     and_name    TEXT PRIMARY KEY,
     cidr        CIDR NOT NULL,
-    assigned_at TIMESTAMPTZ DEFAULT NOW()
+    assigned_at TIMESTAMPTZ DEFAULT NOW(),
+    CONSTRAINT address_leases_cidr_unique UNIQUE (cidr)
 );
 
 -- Domain records

netengine/core/state.py

@@ -164,6 +164,20 @@ class DomainRegistryOutput(PhaseOutputBase, total=False):
     tld_delegations: list[dict[str, JsonValue]]
 
 
+class ANDInstanceState(TypedDict, total=False):
+    name: str
+    org: str
+    profile: str
+    cidr: str
+    gateway_ip: str
+    bridge_name: str
+    dns_suffix: str
+    deployed_at: str
+    dynamic_ip: bool
+    reverse_dns: bool
+    bgp: str | None
+
+
 class GenericPhaseOutput(PhaseOutputBase, total=False):
     status: str
     healthy: bool
@@ -221,6 +235,7 @@ class RuntimeState:
     domain_registry_output: Optional[DomainRegistryOutput] = None
     identity_inworld_output: Optional[GenericPhaseOutput] = None
     ands_output: Optional[GenericPhaseOutput] = None
+    ands_instances: Dict[str, ANDInstanceState] = field(default_factory=dict)
     world_services_output: Optional[GenericPhaseOutput] = None
     org_apps_output: Optional[GenericPhaseOutput] = None
 

netengine/handlers/domain_registry_handler.py

@@ -1,4 +1,5 @@
 # netengine/handlers/domain_registry_handler.py
+import ipaddress
 from typing import Any, List, cast
 
 from netengine.core.pgmq_client import PGMQClient
@@ -35,14 +36,70 @@ async def seed_address_pools(self, spec: Any) -> None:
             ).execute()
 
     async def allocate_address(self, and_name: str, profile: str) -> str:
-        """Allocate a CIDR from the pool; row-level lock prevents conflicts."""
+        """Allocate a unique AND subnet from the registry phase address pool.
+
+        Existing leases are idempotently reused by ``and_name``. New leases are
+        allocated as /24s within the profile pool (or the pool itself when it is
+        smaller than /24). The ``address_leases.cidr`` unique constraint is the
+        database-backed collision guard; if another allocator wins the same CIDR,
+        this method refreshes leases and tries the next free candidate.
+        """
         db = await self._get_db()
-        result = await db.table("address_pools").select("cidr").eq("profile", profile).execute()
-        if not result.data:
+
+        existing = (
+            await db.table("address_leases")
+            .select("and_name,cidr")
+            .eq("and_name", and_name)
+            .execute()
+        )
+        if existing.data and existing.data[0].get("and_name") == and_name:
+            return cast(str, existing.data[0]["cidr"])
+
+        pool_result = (
+            await db.table("address_pools").select("cidr").eq("profile", profile).execute()
+        )
+        if not pool_result.data:
             raise RegistryError(f"No address pool for profile {profile}")
-        pool_cidr = result.data[0]["cidr"]
-        await db.table("address_leases").upsert({"and_name": and_name, "cidr": pool_cidr}).execute()
-        return cast(str, pool_cidr)
+
+        pool = ipaddress.ip_network(str(pool_result.data[0]["cidr"]), strict=False)
+        prefix = max(pool.prefixlen, 24)
+        candidates = [pool] if pool.prefixlen > 24 else list(pool.subnets(new_prefix=prefix))
+
+        last_error: Exception | None = None
+        for _ in range(2):
+            lease_result = await db.table("address_leases").select("and_name,cidr").execute()
+            used = {
+                ipaddress.ip_network(str(row["cidr"]), strict=False)
+                for row in lease_result.data
+                if "and_name" in row
+            }
+            for candidate in candidates:
+                if candidate in used:
+                    continue
+                cidr = str(candidate)
+                try:
+                    await db.table("address_leases").upsert(
+                        {"and_name": and_name, "cidr": cidr, "profile": profile}
+                    ).execute()
+                except Exception as exc:
+                    last_error = exc
+                    if "profile" in str(exc).lower():
+                        try:
+                            await db.table("address_leases").upsert(
+                                {"and_name": and_name, "cidr": cidr}
+                            ).execute()
+                        except Exception as retry_exc:
+                            last_error = retry_exc
+                            continue
+                    else:
+                        continue
+                return cidr
+
+        if last_error is not None:
+            raise RegistryError(
+                f"Address pool exhausted for profile {profile}; last collision: {last_error}"
+            )
+        raise RegistryError(f"Address pool exhausted for profile {profile}")
 
     async def register_domain(self, domain: str, org_name: str, ns_records: List[str]) -> None:
         """Register a domain; emit DNS update event."""

netengine/phases/phase_ands.py

@@ -21,6 +21,7 @@
 from netengine.handlers._base import BasePhaseHandler
 from netengine.handlers.context import PhaseContext
 from netengine.handlers.dns import DNSHandler
+from netengine.handlers.domain_registry_handler import DomainRegistryHandler
 from netengine.handlers.docker_handler import DockerHandler
 from netengine.handlers.gateway_handler import GatewayHandler
 
@@ -116,9 +117,7 @@ async def execute(self, context: PhaseContext) -> None:
                 profiles_used.add(and_instance.profile)
 
                 # Store in runtime_state for this session
-                if not hasattr(context.runtime_state, "ands_instances"):
-                    context.runtime_state.ands_instances = {}  # type: ignore[attr-defined]
-                context.runtime_state.ands_instances[and_instance.name] = and_data  # type: ignore[attr-defined]
+                context.runtime_state.ands_instances[and_instance.name] = and_data
 
             ands_output["ands_provisioned"] = ands_provisioned
             ands_output["address_allocations"] = address_allocations
@@ -168,11 +167,7 @@ async def healthcheck(self, context: PhaseContext) -> bool:
                 logger.warning("No ANDs provisioned")
                 return False
 
-            if not hasattr(context.runtime_state, "ands_instances"):
-                logger.warning("AND instances not tracked in runtime_state")
-                return False
-
-            instances = context.runtime_state.ands_instances  # type: ignore[attr-defined]
+            instances = context.runtime_state.ands_instances
             if not all(and_name in instances for and_name in ands_provisioned):
                 logger.warning("Some AND instances missing from runtime_state")
                 return False
@@ -224,7 +219,7 @@ async def _provision_and(
             raise RuntimeError(f"Profile '{profile_name}' not found in spec")
 
         # 2. Allocate subnet
-        cidr = await self._allocate_address(and_instance.name, profile_name, ands_spec)
+        cidr = await self._allocate_address(context, and_instance.name, profile_name, ands_spec)
         logger.info(f"Allocated CIDR for {and_instance.name}: {cidr}")
 
         # 3. Create isolated Docker bridge network
@@ -339,6 +334,9 @@ async def _provision_and(
             "bridge_name": bridge_name,
             "dns_suffix": dns_suffix,
             "deployed_at": datetime.now(UTC).isoformat(),
+            "dynamic_ip": bool(profile_obj.dynamic_ip),
+            "reverse_dns": bool(profile_obj.reverse_dns),
+            "bgp": profile_obj.bgp,
         }
 
         try:
@@ -352,6 +350,7 @@ async def _provision_and(
 
     async def _allocate_address(
         self,
+        context: PhaseContext,
         and_name: str,
         profile: str,
         ands_spec: Any,
@@ -361,12 +360,167 @@ async def _allocate_address(
         if profile_obj is None:
             raise RuntimeError(f"Profile '{profile}' not found")
 
-        # Sequential /24 allocation within 172.16.0.0/12.
-        # Uses ord-sum rather than hash() to avoid collision with >256 ANDs.
-        idx = sum(ord(c) for c in and_name) % 4096
-        third_octet = idx % 256
-        second_extra = idx // 256
-        return f"172.{16 + second_extra}.{third_octet}.0/24"
+        try:
+            registry = DomainRegistryHandler(context=context)
+            return await registry.allocate_address(and_name, profile)
+        except Exception as exc:
+            # Tests and mock-mode executions may not have a DB. Fall back only when
+            # the registry output does not expose seeded pools; otherwise surface
+            # the allocation failure so exhaustion/collisions are not hidden.
+            output = context.runtime_state.domain_registry_output or {}
+            if output.get("address_pools") or output.get("address_pools_seeded"):
+                raise RuntimeError(f"Failed to allocate address for {and_name}: {exc}") from exc
+            context.logger.warning("Registry allocation unavailable; using mock CIDR: %s", exc)
+            idx = len(context.runtime_state.ands_instances)
+            return f"172.31.{idx}.0/24"
+
+    async def reconcile(
+        self,
+        context: PhaseContext,
+        docker: DockerHandler | None = None,
+        gateway: GatewayHandler | None = None,
+    ) -> dict[str, list[str]]:
+        """Reconcile desired AND spec against runtime/Docker/gateway/DNS side effects."""
+        docker = docker or DockerHandler()
+        gateway = gateway or GatewayHandler(docker)
+        desired = {inst.name: inst for inst in context.spec.ands.instances}
+        existing = context.runtime_state.ands_instances
+        actions: dict[str, list[str]] = {
+            "created": [],
+            "updated": [],
+            "removed": [],
+            "repaired": [],
+        }
+
+        for and_name in set(existing) - set(desired):
+            await self._teardown_and(context, docker, gateway, and_name, existing[and_name])
+            actions["removed"].append(and_name)
+
+        for and_name, inst in desired.items():
+            current = existing.get(and_name)
+            if current is None:
+                data = await self._provision_and(context, docker, gateway, inst, context.spec.ands)
+                context.runtime_state.ands_instances[and_name] = data
+                actions["created"].append(and_name)
+                continue
+            if current.get("profile") != inst.profile:
+                await self._update_and_profile(context, gateway, inst, current, context.spec.ands)
+                actions["updated"].append(and_name)
+            repaired = await self._repair_and(
+                context, docker, gateway, inst, current, context.spec.ands
+            )
+            if repaired:
+                actions["repaired"].append(and_name)
+        return actions
+
+    async def _update_and_profile(
+        self,
+        context: PhaseContext,
+        gateway: GatewayHandler,
+        and_instance: Any,
+        current: dict[str, Any],
+        ands_spec: Any,
+    ) -> None:
+        old_profile = ands_spec.profiles.get(current.get("profile")) if ands_spec.profiles else None
+        new_profile = ands_spec.profiles.get(and_instance.profile) if ands_spec.profiles else None
+        if new_profile is None:
+            raise RuntimeError(f"Profile '{and_instance.profile}' not found in spec")
+        if old_profile and getattr(old_profile, "dynamic_ip", False) and not new_profile.dynamic_ip:
+            await gateway.remove_dhcp(and_instance.name)
+        if old_profile and getattr(old_profile, "bgp", None) and not new_profile.bgp:
+            await gateway.remove_bgp(and_instance.name)
+        rules = await gateway.generate_rules(
+            and_instance.name, and_instance.profile, current["cidr"]
+        )
+        await gateway.apply_rules(and_instance.name, rules)
+        if new_profile.dynamic_ip:
+            await gateway.setup_dhcp(and_instance.name, current["cidr"], current["gateway_ip"])
+        if new_profile.reverse_dns:
+            await DNSHandler().add_reverse_zone(context, current["cidr"], current["gateway_ip"])
+        if new_profile.bgp:
+            await gateway.setup_bgp(
+                and_instance.name, current["cidr"], current["gateway_ip"], new_profile.bgp
+            )
+        current.update(
+            {
+                "profile": and_instance.profile,
+                "dynamic_ip": new_profile.dynamic_ip,
+                "reverse_dns": new_profile.reverse_dns,
+                "bgp": new_profile.bgp,
+            }
+        )
+
+    async def _teardown_and(
+        self,
+        context: PhaseContext,
+        docker: DockerHandler,
+        gateway: GatewayHandler,
+        and_name: str,
+        current: dict[str, Any],
+    ) -> None:
+        await gateway.remove_rules(and_name)
+        if current.get("dynamic_ip"):
+            await gateway.remove_dhcp(and_name)
+        if current.get("reverse_dns"):
+            await DNSHandler().remove_reverse_zone(context, current["cidr"])
+        if current.get("bgp"):
+            await gateway.remove_bgp(and_name)
+        bridge_name = current.get("bridge_name", f"netengines_and_{and_name}")
+        try:
+            await docker.disconnect_network(gateway.gateway_container, bridge_name)
+        except Exception:
+            pass
+        await docker.remove_network(bridge_name)
+        context.runtime_state.ands_instances.pop(and_name, None)
+        try:
+            db = await self._get_supabase()
+            await db.table("address_leases").delete().eq("and_name", and_name).execute()
+            await db.table("and_instances").delete().eq("name", and_name).execute()
+        except Exception as exc:
+            context.logger.warning("Failed to delete AND DB state for %s: %s", and_name, exc)
+
+    async def _repair_and(
+        self,
+        context: PhaseContext,
+        docker: DockerHandler,
+        gateway: GatewayHandler,
+        and_instance: Any,
+        current: dict[str, Any],
+        ands_spec: Any,
+    ) -> bool:
+        repaired = False
+        bridge_name = current.get("bridge_name", f"netengines_and_{and_instance.name}")
+        try:
+            docker.client.networks.get(bridge_name)
+        except Exception:
+            await docker.create_network(
+                name=bridge_name, driver="bridge", subnet=current["cidr"], internal=True
+            )
+            repaired = True
+        try:
+            await docker.connect_network(
+                container=gateway.gateway_container, network=bridge_name, ip=current["gateway_ip"]
+            )
+        except Exception:
+            pass
+        profile = ands_spec.profiles.get(and_instance.profile) if ands_spec.profiles else None
+        rules = await gateway.generate_rules(
+            and_instance.name, and_instance.profile, current["cidr"]
+        )
+        await gateway.apply_rules(and_instance.name, rules)
+        dns_suffix = getattr(and_instance, "dns_suffix", None) or current.get("dns_suffix")
+        await DNSHandler().add_zone_record(
+            context, "internal", "A", dns_suffix.rstrip("."), current["gateway_ip"], 300
+        )
+        if profile and profile.dynamic_ip:
+            await gateway.setup_dhcp(and_instance.name, current["cidr"], current["gateway_ip"])
+        if profile and profile.reverse_dns:
+            await DNSHandler().add_reverse_zone(context, current["cidr"], current["gateway_ip"])
+        if profile and profile.bgp:
+            await gateway.setup_bgp(
+                and_instance.name, current["cidr"], current["gateway_ip"], profile.bgp
+            )
+        return repaired
 
     # ─────────────────────────────────────────────
     # Event-Driven Provisioning