tuneplane-node 0.3.15__py3-none-any.whl

tuneplane_node/join.py

@@ -0,0 +1,199 @@
+"""Joining a Fleet, and staying in it.
+
+Two calls and a loop. What makes this small is where the authorization lives:
+the Join Token is the whole of it (ADR-0013), so there is no login, no
+credential file to provision ahead of time, and nothing for the operator to
+have an account for. What comes back is this node's identity and its long-lived
+credential, and those are what every later request carries.
+
+The credential is returned exactly once. Losing it means asking an administrator
+for another Join Token and joining again as a new node, which is the correct
+outcome: a machine that lost its identity is not the same machine to a control
+plane that has to reconcile containers against it.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+import httpx
+
+from tuneplane_node.wire import PROTOCOL
+
+log = logging.getLogger(__name__)
+
+#: Where the node keeps what registration gave it. Beside the agent's own state
+#: rather than in the storage root, because a node that cannot see the storage
+#: root still has to be able to say who it is.
+CREDENTIAL_FILE = "node.json"
+
+#: How often a node says it is still there. Well inside the console's tightest
+#: liveness threshold (serve, seconds) so that a single missed beat is not read
+#: as a machine going away.
+HEARTBEAT_INTERVAL_S = 15.0
+
+
+@dataclass(frozen=True)
+class NodeIdentity:
+    #: Where the control plane answers. Named `server` because that is the word
+    #: an operator arriving from k3s, Nomad or Consul reaches for; the thing it
+    #: points at is this platform's console.
+    server: str
+    fleet_id: str
+    node_id: str
+    credential: str
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "server": self.server,
+            "fleet_id": self.fleet_id,
+            "node_id": self.node_id,
+            "credential": self.credential,
+        }
+
+
+class JoinError(RuntimeError):
+    """Registration was refused, with the console's reason.
+
+    Carries the console's own wording rather than a summary: the refusals are
+    written to be acted on ("mount it, or join a fleet that pulls from object
+    storage"), and rewording them here would lose that.
+    """
+
+
+def identity_path(state_dir: Path) -> Path:
+    return Path(state_dir) / CREDENTIAL_FILE
+
+
+def load_identity(state_dir: Path) -> Optional[NodeIdentity]:
+    path = identity_path(state_dir)
+    if not path.exists():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return NodeIdentity(
+            server=str(data["server"]),
+            fleet_id=str(data["fleet_id"]),
+            node_id=str(data["node_id"]),
+            credential=str(data["credential"]),
+        )
+    except (ValueError, KeyError, OSError) as exc:
+        raise JoinError(f"{path} is not a readable node identity: {exc}") from exc
+
+
+def save_identity(state_dir: Path, identity: NodeIdentity) -> Path:
+    path = identity_path(state_dir)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(identity.to_dict(), indent=2) + "\n", encoding="utf-8")
+    # It is a bearer credential; nothing else on the box needs to read it.
+    path.chmod(0o600)
+    return path
+
+
+async def join(
+    server: str,
+    *,
+    fleet_id: str,
+    token: str,
+    address: str,
+    inventory: dict[str, Any],
+    node_version: str,
+    connection: str = "direct",
+    timeout: float = 30.0,
+    transport=None,
+) -> tuple[NodeIdentity, list[str]]:
+    """Redeem a Join Token for this node's identity, plus anything the console warns about.
+
+    The warnings are returned rather than saved: they describe the state of the
+    fleet at this moment, and the identity file is meant to outlive that.
+
+    transport is injectable for tests, the same way `NodeClient` takes one.
+    """
+    base = server.rstrip("/")
+    payload = {
+        "token": token,
+        "address": address,
+        "connection": connection,
+        "node_version": node_version,
+        "protocol": PROTOCOL,
+        "inventory": inventory,
+    }
+    async with httpx.AsyncClient(timeout=timeout, transport=transport) as client:
+        response = await client.post(f"{base}/api/fleets/{fleet_id}/nodes", json=payload)
+    if response.status_code != 201:
+        raise JoinError(_detail(response))
+    body = response.json()
+    identity = NodeIdentity(
+        server=base,
+        fleet_id=fleet_id,
+        node_id=str(body["id"]),
+        credential=str(body["credential"]),
+    )
+    return identity, [str(w) for w in (body.get("warnings") or [])]
+
+
+async def beat(
+    identity: NodeIdentity,
+    *,
+    inventory: dict[str, Any],
+    node_version: str,
+    address: str = "",
+    timeout: float = 15.0,
+    transport=None,
+) -> dict[str, Any]:
+    """Tell the console this node is still there, and re-report what it has."""
+    url = f"{identity.server}/api/fleets/{identity.fleet_id}/nodes/{identity.node_id}/heartbeat"
+    payload = {
+        "node_version": node_version,
+        "protocol": PROTOCOL,
+        "address": address,
+        "inventory": inventory,
+    }
+    async with httpx.AsyncClient(timeout=timeout, transport=transport) as client:
+        response = await client.post(
+            url, json=payload, headers={"Authorization": f"Bearer {identity.credential}"}
+        )
+    if response.status_code != 200:
+        raise JoinError(_detail(response))
+    return response.json()
+
+
+async def heartbeat_forever(
+    identity: NodeIdentity,
+    settings,
+    *,
+    node_version: str,
+    address: str = "",
+    interval_s: float = HEARTBEAT_INTERVAL_S,
+) -> None:
+    """Beat until cancelled. A failed beat is logged and retried, never fatal.
+
+    Losing contact with the console is not a reason to stop running work: the
+    containers on this machine are the truth about what is running, and the
+    console reconciles against them when it comes back. That is the same rule
+    ADR-002 states from the other side -- an unreachable node is not a dead job.
+    """
+    from tuneplane_node.inventory import collect
+
+    while True:
+        try:
+            inventory = (await collect(settings)).to_dict()
+            await beat(identity, inventory=inventory, node_version=node_version, address=address)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:  # noqa: BLE001
+            log.warning("heartbeat to %s failed: %s", identity.server, exc)
+        await asyncio.sleep(interval_s)
+
+
+def _detail(response: httpx.Response) -> str:
+    try:
+        body = response.json()
+    except ValueError:
+        return f"{response.status_code}: {response.text[:200]}"
+    detail = body.get("detail") if isinstance(body, dict) else None
+    return str(detail or f"{response.status_code}: {response.text[:200]}")