实现首版 Gitea microVM runner

2026-09-16 13:06:05 +00:00
parent 0ee1d50703
commit 23695fcc98
14 changed files with 580 additions and 2 deletions
@@ -0,0 +1 @@
+"""Gitea microVM runner controller and worker."""
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""Gitea workflow_job webhook to NATS JetStream producer."""
+
+import hashlib
+import hmac
+import json
+import os
+import ssl
+from pathlib import Path
+
+import nats
+from aiohttp import web
+from nats.js.api import DiscardPolicy, RetentionPolicy, StorageType, StreamConfig
+from nats.js.errors import NotFoundError
+
+LABEL = os.environ.get("RUNNER_LABEL", "kind-microvm")
+SUBJECT = os.environ.get("NATS_SUBJECT", f"ci.runner.{LABEL}")
+STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
+NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
+NATS_USER = os.environ.get("NATS_USER", "ci-producer")
+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")
+WEBHOOK_SECRET_FILE = Path(os.environ.get("WEBHOOK_SECRET_FILE", "/run/secrets/gitea/webhook-secret"))
+
+
+def accepts(payload: object) -> tuple[bool, str | None]:
+    if not isinstance(payload, dict) or payload.get("action") != "queued":
+        return False, None
+    job = payload.get("workflow_job")
+    if not isinstance(job, dict) or LABEL not in job.get("labels", []):
+        return False, None
+    job_id = job.get("id")
+    if not isinstance(job_id, int) or isinstance(job_id, bool):
+        return False, None
+    return True, str(job_id)
+
+
+def valid_signature(body: bytes, signature: str) -> bool:
+    expected = hmac.new(WEBHOOK_SECRET_FILE.read_bytes().strip(), body, hashlib.sha256).hexdigest()
+    return hmac.compare_digest(signature.removeprefix("sha256="), expected)
+
+
+async def ensure_stream(js: object) -> None:
+    config = StreamConfig(
+        name=STREAM,
+        subjects=["ci.runner.*"],
+        retention=RetentionPolicy.WORK_QUEUE,
+        storage=StorageType.FILE,
+        discard=DiscardPolicy.OLD,
+        max_age=24 * 60 * 60,
+        max_msgs=10_000,
+        max_bytes=256 * 1024 * 1024,
+        duplicate_window=24 * 60 * 60,
+    )
+    try:
+        await js.stream_info(STREAM)
+    except NotFoundError:
+        await js.add_stream(config=config)
+    else:
+        await js.update_stream(config=config)
+
+
+async def webhook(request: web.Request) -> web.Response:
+    body = await request.read()
+    if not valid_signature(body, request.headers.get("X-Gitea-Signature", "")):
+        raise web.HTTPUnauthorized()
+    try:
+        payload = json.loads(body)
+    except json.JSONDecodeError as error:
+        raise web.HTTPBadRequest(text="invalid JSON\n") from error
+    accepted, job_id = accepts(payload)
+    if not accepted:
+        return web.Response(status=204)
+    await request.app["js"].publish(
+        SUBJECT,
+        body,
+        headers={"Nats-Msg-Id": f"gitea-workflow-job-{job_id}"},
+    )
+    return web.Response(status=202, text="queued\n")
+
+
+async def health(request: web.Request) -> web.Response:
+    connected = request.app["nc"].is_connected
+    return web.Response(text="ok\n" if connected else "disconnected\n", status=200 if connected else 503)
+
+
+async def nats_context(app: web.Application):
+    tls = 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=tls,
+        name="microvm-runner-controller",
+    )
+    app["nc"] = nc
+    app["js"] = nc.jetstream()
+    await ensure_stream(app["js"])
+    yield
+    await nc.drain()
+
+
+def create_app() -> web.Application:
+    app = web.Application(client_max_size=1024 * 1024)
+    app.cleanup_ctx.append(nats_context)
+    app.router.add_post("/webhook", webhook)
+    app.router.add_get("/healthz", health)
+    return app
+
+
+def main() -> None:
+    web.run_app(create_app(), host=os.environ.get("LISTEN", "0.0.0.0"), port=int(os.environ.get("PORT", "8787")))
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""Capacity-bounded JetStream consumer that launches one ephemeral VM per job."""
+
+import asyncio
+import logging
+import os
+import secrets
+import ssl
+import uuid
+from pathlib import Path
+
+import nats
+from aiohttp import web
+from nats.errors import TimeoutError
+from nats.js.api import AckPolicy, ConsumerConfig
+
+LOG = logging.getLogger(__name__)
+CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "1"))
+SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.kind-microvm")
+STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
+DURABLE = os.environ.get("NATS_DURABLE", "kind-microvm")
+MAX_INFLIGHT = int(os.environ.get("RUNNER_MAX_INFLIGHT", "64"))
+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", "/etc/microvm-runner/nats-password"))
+NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
+REGISTRATION_TOKEN_FILE = Path(os.environ.get("REGISTRATION_TOKEN_FILE", "/etc/microvm-runner/registration-token"))
+LAUNCHER = os.environ.get("LAUNCHER", "/usr/local/libexec/microvm-runner-launch")
+TOKEN_LISTEN = os.environ.get("TOKEN_LISTEN", "172.30.0.1")
+TOKEN_PORT = int(os.environ.get("TOKEN_PORT", "8787"))
+
+tokens: dict[str, bytes] = {}
+token_lock = asyncio.Lock()
+
+
+async def token(request: web.Request) -> web.Response:
+    nonce = request.match_info["nonce"]
+    async with token_lock:
+        value = tokens.pop(nonce, None)
+    if value is None:
+        raise web.HTTPNotFound()
+    return web.Response(body=value, headers={"Cache-Control": "no-store"})
+
+
+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 run_one(message: object) -> None:
+    instance_id = str(uuid.uuid4())
+    nonce = secrets.token_urlsafe(32)
+    async with token_lock:
+        tokens[nonce] = REGISTRATION_TOKEN_FILE.read_bytes().strip()
+    stop = asyncio.Event()
+    pulse = asyncio.create_task(heartbeat(message, stop))
+    try:
+        process = await asyncio.create_subprocess_exec(LAUNCHER, instance_id, nonce)
+        return_code = await process.wait()
+        if return_code == 0:
+            await message.ack()
+        else:
+            LOG.error("launcher for %s exited with %d", instance_id, return_code)
+            await message.nak(delay=30)
+    except Exception:
+        await message.nak(delay=30)
+        raise
+    finally:
+        stop.set()
+        await pulse
+        async with token_lock:
+            tokens.pop(nonce, None)
+
+
+def _report_task(task: asyncio.Task[None]) -> None:
+    if task.cancelled():
+        return
+    error = task.exception()
+    if error is not None:
+        LOG.error(
+            "runner task failed",
+            exc_info=(type(error), error, error.__traceback__),
+        )
+
+
+async def consume() -> None:
+    if CAPACITY < 1:
+        raise ValueError("RUNNER_CAPACITY must be at least 1")
+    tls = 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=tls,
+        name=DURABLE,
+    )
+    js = nc.jetstream()
+    subscription = await js.pull_subscribe(
+        SUBJECT,
+        durable=DURABLE,
+        stream=STREAM,
+        config=ConsumerConfig(
+            durable_name=DURABLE,
+            filter_subject=SUBJECT,
+            ack_policy=AckPolicy.EXPLICIT,
+            ack_wait=5 * 60,
+            max_ack_pending=MAX_INFLIGHT,
+            max_deliver=5,
+        ),
+    )
+    active: set[asyncio.Task[None]] = set()
+    try:
+        while True:
+            active = {task for task in active if not task.done()}
+            free = CAPACITY - len(active)
+            if free == 0:
+                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_one(message))
+                task.add_done_callback(_report_task)
+                active.add(task)
+    finally:
+        if active:
+            await asyncio.gather(*active, return_exceptions=True)
+        await nc.drain()
+
+
+async def main() -> None:
+    app = web.Application()
+    app.router.add_get("/token/{nonce}", token)
+    runner = web.AppRunner(app)
+    await runner.setup()
+    await web.TCPSite(runner, TOKEN_LISTEN, TOKEN_PORT).start()
+    try:
+        await consume()
+    finally:
+        await runner.cleanup()
+
+
+def cli() -> None:
+    logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    cli()
				`@@ -0,0 +1 @@`
				`"""Gitea microVM runner controller and worker."""`