实现一次性 Pod Runner 后端

src/gitea_microvm_runner/kubernetes.py

@@ -0,0 +1,89 @@
+"""Small in-cluster Kubernetes API client used by the Pod backend."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from urllib.parse import quote
+
+from aiohttp import ClientResponseError, ClientSession, TCPConnector
+import ssl
+
+
+class KubernetesClient:
+    def __init__(
+        self,
+        *,
+        api_url: str,
+        token_file: Path,
+        ca_file: Path,
+        namespace: str,
+    ) -> None:
+        self.api_url = api_url.rstrip("/")
+        self.token_file = token_file
+        self.ca_file = ca_file
+        self.namespace = namespace
+        self.session: ClientSession | None = None
+
+    async def __aenter__(self) -> KubernetesClient:
+        context = ssl.create_default_context(cafile=self.ca_file)
+        self.session = ClientSession(
+            connector=TCPConnector(ssl=context),
+            headers={
+                "Authorization": f"Bearer {self.token_file.read_text().strip()}",
+            },
+            raise_for_status=True,
+        )
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        if self.session is not None:
+            await self.session.close()
+
+    async def create_pod(self, manifest: dict[str, object]) -> dict[str, object]:
+        response = await self._request("POST", self._pods_path(), json=manifest)
+        return await response.json()
+
+    async def get_pod(self, name: str) -> dict[str, object] | None:
+        try:
+            response = await self._request("GET", f"{self._pods_path()}/{quote(name)}")
+        except ClientResponseError as error:
+            if error.status == 404:
+                return None
+            raise
+        return await response.json()
+
+    async def bind_identity(self, name: str, identity_path: str) -> None:
+        patch = {
+            "metadata": {
+                "labels": {"ci.ddupan.top/identity-bound": "true"},
+                "annotations": {"ci.ddupan.top/spiffe-path": identity_path},
+            }
+        }
+        response = await self._request(
+            "PATCH",
+            f"{self._pods_path()}/{quote(name)}",
+            data=json.dumps(patch),
+            headers={"Content-Type": "application/merge-patch+json"},
+        )
+        response.release()
+
+    async def delete_pod(self, name: str) -> None:
+        try:
+            response = await self._request(
+                "DELETE",
+                f"{self._pods_path()}/{quote(name)}",
+                json={"gracePeriodSeconds": 30, "propagationPolicy": "Background"},
+            )
+            response.release()
+        except ClientResponseError as error:
+            if error.status != 404:
+                raise
+
+    async def _request(self, method: str, path: str, **kwargs: object):
+        if self.session is None:
+            raise RuntimeError("KubernetesClient is not open")
+        return await self.session.request(method, f"{self.api_url}{path}", **kwargs)
+
+    def _pods_path(self) -> str:
+        return f"/api/v1/namespaces/{quote(self.namespace)}/pods"

src/gitea_microvm_runner/pod_worker.py

@@ -0,0 +1,307 @@
+#!/usr/bin/env python3
+"""JetStream worker that creates one disposable Kubernetes Pod per job."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from pathlib import Path
+import ssl
+from urllib.parse import quote
+import uuid
+
+import nats
+from aiohttp import ClientResponseError
+from nats.errors import TimeoutError
+from nats.js.api import AckPolicy, ConsumerConfig
+
+from .kubernetes import KubernetesClient
+from .models import IdentityBinding, RunnerRequest
+
+
+LOG = logging.getLogger(__name__)
+STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
+REQUEST_SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.pod")
+BINDING_SUBJECT = os.environ.get("NATS_BINDING_SUBJECT", "ci.runner.pod.binding")
+NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
+NATS_USER = os.environ.get("NATS_USER", "ci-worker")
+NATS_PASSWORD_FILE = Path(os.environ.get("NATS_PASSWORD_FILE", "/run/secrets/nats/password"))
+NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
+KUBERNETES_API = os.environ.get("KUBERNETES_API", "https://kubernetes.default.svc")
+KUBERNETES_TOKEN_FILE = Path(
+    os.environ.get(
+        "KUBERNETES_TOKEN_FILE",
+        "/var/run/secrets/kubernetes.io/serviceaccount/token",
+    )
+)
+KUBERNETES_CA_FILE = Path(
+    os.environ.get(
+        "KUBERNETES_CA_FILE",
+        "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt",
+    )
+)
+NAMESPACE = os.environ.get("RUNNER_NAMESPACE", "gitea-actions")
+RUNNER_IMAGE = os.environ.get("RUNNER_IMAGE", "docker.io/gitea/runner:2")
+RUNNER_SERVICE_ACCOUNT = os.environ.get("RUNNER_SERVICE_ACCOUNT", "gitea-dynamic-runner")
+RUNNER_TOKEN_SECRET = os.environ.get("RUNNER_TOKEN_SECRET", "gitea-dynamic-runner")
+GITEA_INSTANCE = os.environ.get("GITEA_INSTANCE", "https://git.ddupan.top")
+CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "4"))
+POD_TIMEOUT = int(os.environ.get("RUNNER_POD_TIMEOUT", str(4 * 60 * 60)))
+
+
+def identity_path(repository: str, job_name: str) -> str:
+    parts = repository.split("/")
+    if len(parts) != 2 or not all(parts):
+        raise ValueError("repository must be owner/name")
+    encoded = [quote(part, safe="-._~") for part in (*parts, job_name)]
+    if not job_name.strip() or any(not part for part in encoded):
+        raise ValueError("identity components must not be empty")
+    return "/".join(encoded)
+
+
+def pod_manifest(request: RunnerRequest, pod_name: str) -> dict[str, object]:
+    if request.backend != "pod":
+        raise ValueError("Pod backend only accepts pod requests")
+    return {
+        "apiVersion": "v1",
+        "kind": "Pod",
+        "metadata": {
+            "name": pod_name,
+            "namespace": NAMESPACE,
+            "labels": {
+                "app.kubernetes.io/name": "gitea-dynamic-runner",
+                "app.kubernetes.io/component": "runner",
+                "ci.ddupan.top/backend": "pod",
+            },
+            "annotations": {
+                "ci.ddupan.top/queued-job-id": str(request.job_id),
+                "ci.ddupan.top/queued-run-id": str(request.run_id),
+            },
+        },
+        "spec": {
+            "serviceAccountName": RUNNER_SERVICE_ACCOUNT,
+            "restartPolicy": "Never",
+            "terminationGracePeriodSeconds": 30,
+            "containers": [
+                {
+                    "name": "runner",
+                    "image": RUNNER_IMAGE,
+                    "imagePullPolicy": "IfNotPresent",
+                    "securityContext": {"privileged": True},
+                    "env": [
+                        {"name": "GITEA_INSTANCE_URL", "value": GITEA_INSTANCE},
+                        {
+                            "name": "GITEA_RUNNER_NAME",
+                            "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}},
+                        },
+                        {
+                            "name": "GITEA_RUNNER_REGISTRATION_TOKEN_FILE",
+                            "value": "/run/secrets/gitea/token",
+                        },
+                        {"name": "GITEA_RUNNER_LABELS", "value": "self-hosted:host,pod:host"},
+                        {"name": "GITEA_RUNNER_EPHEMERAL", "value": "1"},
+                        {"name": "GITEA_RUNNER_ONCE", "value": "1"},
+                        {"name": "CONFIG_FILE", "value": "/etc/gitea-runner/config.yaml"},
+                    ],
+                    "volumeMounts": [
+                        {
+                            "name": "registration-token",
+                            "mountPath": "/run/secrets/gitea",
+                            "readOnly": True,
+                        },
+                        {
+                            "name": "spire-agent-socket",
+                            "mountPath": "/run/spire/agent-sockets",
+                            "readOnly": True,
+                        },
+                    ],
+                }
+            ],
+            "volumes": [
+                {
+                    "name": "registration-token",
+                    "secret": {
+                        "secretName": RUNNER_TOKEN_SECRET,
+                        "items": [{"key": "token", "path": "token"}],
+                    },
+                },
+                {
+                    "name": "spire-agent-socket",
+                    "csi": {"driver": "csi.spiffe.io", "readOnly": True},
+                },
+            ],
+        },
+    }
+
+
+async def heartbeat(message: object, stop: asyncio.Event) -> None:
+    while True:
+        try:
+            await asyncio.wait_for(stop.wait(), timeout=60)
+            return
+        except asyncio.TimeoutError:
+            await message.in_progress()
+
+
+async def wait_for_pod(client: KubernetesClient, name: str) -> bool:
+    deadline = asyncio.get_running_loop().time() + POD_TIMEOUT
+    while asyncio.get_running_loop().time() < deadline:
+        pod = await client.get_pod(name)
+        if pod is None:
+            raise RuntimeError(f"runner Pod {name} disappeared")
+        status = pod.get("status")
+        phase = status.get("phase") if isinstance(status, dict) else None
+        if phase == "Succeeded":
+            return True
+        if phase == "Failed":
+            return False
+        await asyncio.sleep(2)
+    raise asyncio.TimeoutError(f"runner Pod {name} timed out")
+
+
+async def run_request(message: object, client: KubernetesClient) -> None:
+    try:
+        request = RunnerRequest.from_json(message.data)
+    except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as error:
+        LOG.error("discarding invalid Pod request: %s", error)
+        await message.ack()
+        return
+    if request.backend != "pod":
+        LOG.error("discarding %s request received by Pod worker", request.backend)
+        await message.ack()
+        return
+
+    pod_name = f"gitea-pod-{uuid.uuid4().hex[:12]}"
+    stop = asyncio.Event()
+    pulse = asyncio.create_task(heartbeat(message, stop))
+    try:
+        await client.create_pod(pod_manifest(request, pod_name))
+        if await wait_for_pod(client, pod_name):
+            await message.ack()
+        else:
+            LOG.error("runner Pod %s failed", pod_name)
+            await message.nak(delay=30)
+    except Exception:
+        await message.nak(delay=30)
+        raise
+    finally:
+        stop.set()
+        await pulse
+        await client.delete_pod(pod_name)
+
+
+async def bind_request(message: object, client: KubernetesClient) -> None:
+    try:
+        document = json.loads(message.data)
+        binding = IdentityBinding(**document)
+        if binding.backend != "pod" or not binding.runner_name.startswith("gitea-pod-"):
+            raise ValueError("invalid Pod identity binding")
+        path = identity_path(binding.repository, binding.job_name)
+        await client.bind_identity(binding.runner_name, path)
+    except ClientResponseError as error:
+        if error.status == 404:
+            await message.nak(delay=2)
+            return
+        await message.nak(delay=30)
+        raise
+    except (json.JSONDecodeError, TypeError, ValueError) as error:
+        LOG.error("discarding invalid identity binding: %s", error)
+        await message.ack()
+        return
+    await message.ack()
+
+
+async def consume_requests(subscription: object, client: KubernetesClient) -> None:
+    active: set[asyncio.Task[None]] = set()
+    while True:
+        active = {task for task in active if not task.done()}
+        free = CAPACITY - len(active)
+        if free < 1:
+            await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED)
+            continue
+        try:
+            messages = await subscription.fetch(batch=free, timeout=5)
+        except TimeoutError:
+            continue
+        for message in messages:
+            task = asyncio.create_task(run_request(message, client))
+            task.add_done_callback(_report_task)
+            active.add(task)
+
+
+async def consume_bindings(subscription: object, client: KubernetesClient) -> None:
+    while True:
+        try:
+            messages = await subscription.fetch(batch=16, timeout=5)
+        except TimeoutError:
+            continue
+        await asyncio.gather(*(bind_request(message, client) for message in messages))
+
+
+def _report_task(task: asyncio.Task[None]) -> None:
+    if not task.cancelled() and (error := task.exception()) is not None:
+        LOG.error("Pod runner task failed", exc_info=(type(error), error, error.__traceback__))
+
+
+async def main() -> None:
+    if CAPACITY < 1:
+        raise ValueError("RUNNER_CAPACITY must be at least 1")
+    context = ssl.create_default_context(cafile=NATS_CA_FILE)
+    nc = await nats.connect(
+        NATS_URL,
+        user=NATS_USER,
+        password=NATS_PASSWORD_FILE.read_text().strip(),
+        tls=context,
+        name="pod-runner-worker",
+    )
+    js = nc.jetstream()
+    requests = await js.pull_subscribe(
+        REQUEST_SUBJECT,
+        durable="pod",
+        stream=STREAM,
+        config=ConsumerConfig(
+            durable_name="pod",
+            filter_subject=REQUEST_SUBJECT,
+            ack_policy=AckPolicy.EXPLICIT,
+            ack_wait=5 * 60,
+            max_ack_pending=max(CAPACITY, 1),
+            max_deliver=5,
+        ),
+    )
+    bindings = await js.pull_subscribe(
+        BINDING_SUBJECT,
+        durable="pod-binding",
+        stream=STREAM,
+        config=ConsumerConfig(
+            durable_name="pod-binding",
+            filter_subject=BINDING_SUBJECT,
+            ack_policy=AckPolicy.EXPLICIT,
+            ack_wait=30,
+            max_ack_pending=64,
+            max_deliver=10,
+        ),
+    )
+    async with KubernetesClient(
+        api_url=KUBERNETES_API,
+        token_file=KUBERNETES_TOKEN_FILE,
+        ca_file=KUBERNETES_CA_FILE,
+        namespace=NAMESPACE,
+    ) as client:
+        try:
+            await asyncio.gather(
+                consume_requests(requests, client),
+                consume_bindings(bindings, client),
+            )
+        finally:
+            await nc.drain()
+
+
+def cli() -> None:
+    logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    cli()