实现首版 Gitea microVM runner

src/gitea_microvm_runner/worker.py

@@ -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()