src/gitea_dynamic_runner/controller.py

@@ -3,14 +3,27 @@
 
 from __future__ import annotations
 
+import asyncio
 import hashlib
 import hmac
 import json
 import logging
 import os
+import ssl
 from pathlib import Path
 
+import nats
 from aiohttp import web
+from nats.errors import TimeoutError as NatsTimeoutError
+from nats.js.api import (
+    AckPolicy,
+    ConsumerConfig,
+    DiscardPolicy,
+    RetentionPolicy,
+    StorageType,
+    StreamConfig,
+)
+from nats.js.errors import NotFoundError
 
 from .models import RunnerRequest
 from .opensandbox import OpenSandboxClient
@@ -22,19 +35,28 @@ WEBHOOK_SECRET_FILE = Path(
     os.environ.get("WEBHOOK_SECRET_FILE", "/run/secrets/gitea/webhook-secret")
 )
 REGISTRATION_TOKEN_FILE = Path(
-    os.environ.get(
-        "REGISTRATION_TOKEN_FILE", "/run/secrets/gitea/registration-token"
-    )
-)
-OPENSANDBOX_API = os.environ.get(
-    "OPENSANDBOX_API", "http://10.60.0.13:8080"
+    os.environ.get("REGISTRATION_TOKEN_FILE", "/run/secrets/gitea/registration-token")
 )
+OPENSANDBOX_API = os.environ.get("OPENSANDBOX_API", "http://10.60.0.13:8080")
 OPENSANDBOX_API_KEY_FILE = Path(
     os.environ.get(
         "OPENSANDBOX_API_KEY_FILE",
         "/run/secrets/opensandbox/api-key",
     )
 )
+NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
+NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
+NATS_PRODUCER_USER = os.environ.get("NATS_PRODUCER_USER", "ci-producer")
+NATS_PRODUCER_PASSWORD_FILE = Path(
+    os.environ.get("NATS_PRODUCER_PASSWORD_FILE", "/run/secrets/nats/producer-password")
+)
+NATS_WORKER_USER = os.environ.get("NATS_WORKER_USER", "ci-worker")
+NATS_WORKER_PASSWORD_FILE = Path(
+    os.environ.get("NATS_WORKER_PASSWORD_FILE", "/run/secrets/nats/worker-password")
+)
+NATS_STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
+NATS_SUBJECT_PREFIX = os.environ.get("NATS_SUBJECT_PREFIX", "ci.runner")
+POD_CONSUMER_ENABLED = os.environ.get("POD_CONSUMER_ENABLED", "true") == "true"
 
 
 def accepts(payload: object) -> tuple[bool, str | None]:
@@ -60,21 +82,42 @@ async def webhook(request: web.Request) -> web.Response:
     context = _webhook_context(payload)
     LOG.info("workflow_job webhook received %s", context)
 
+    scheduler: OpenSandboxScheduler = request.app["scheduler"]
+    if isinstance(payload, dict) and payload.get("action") == "completed":
+        job = payload.get("workflow_job")
+        runner_name = job.get("runner_name") if isinstance(job, dict) else None
+        if isinstance(runner_name, str) and runner_name.startswith(
+            ("gitea-pod-", "gitea-vm-")
+        ):
+            cleaned = await scheduler.complete(runner_name)
+            await request.app["js"].publish(
+                f"{NATS_SUBJECT_PREFIX}.lifecycle.completed",
+                body,
+                headers={
+                    "Nats-Msg-Id": (f"gitea-workflow-job-{job.get('id')}-completed")
+                },
+            )
+            LOG.info(
+                "completed runner cleanup runner=%r cleaned=%s %s",
+                runner_name,
+                cleaned,
+                context,
+            )
+        return web.Response(status=204)
+
     runner_request = RunnerRequest.from_webhook(payload)
     if runner_request is None:
         LOG.info("workflow_job webhook ignored %s", context)
         return web.Response(status=204)
 
-    scheduler: OpenSandboxScheduler = request.app["scheduler"]
-    try:
-        sandbox_id = await scheduler.create(runner_request)
-    except ValueError as error:
-        LOG.info("duplicate workflow_job webhook ignored: %s", error)
-        return web.Response(status=202, text="already scheduled\n")
-    except Exception:
-        LOG.exception("failed to create OpenSandbox runner %s", context)
-        raise web.HTTPServiceUnavailable(text="sandbox unavailable\n")
-    return web.Response(status=202, text=f"sandbox={sandbox_id}\n")
+    subject = f"{NATS_SUBJECT_PREFIX}.{runner_request.backend}"
+    await request.app["js"].publish(
+        subject,
+        runner_request.to_json(),
+        headers={"Nats-Msg-Id": f"gitea-workflow-job-{runner_request.job_id}-queued"},
+    )
+    LOG.info("runner request persisted subject=%s %s", subject, context)
+    return web.Response(status=202, text="queued\n")
 
 
 async def registration_token(request: web.Request) -> web.Response:
@@ -108,8 +151,97 @@ async def health(request: web.Request) -> web.Response:
     )
 
 
-async def opensandbox_context(app: web.Application):
+async def ensure_stream(js: object) -> None:
+    config = StreamConfig(
+        name=NATS_STREAM,
+        subjects=[f"{NATS_SUBJECT_PREFIX}.>"],
+        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(NATS_STREAM)
+    except NotFoundError:
+        await js.add_stream(config=config)
+    else:
+        await js.update_stream(config=config)
+
+
+async def run_message(message: object, scheduler: OpenSandboxScheduler) -> None:
+    try:
+        runner_request = RunnerRequest.from_json(message.data)
+        await scheduler.create(runner_request)
+        task = scheduler.active[runner_request.job_id]
+        while not task.done():
+            try:
+                await asyncio.wait_for(asyncio.shield(task), timeout=30)
+            except asyncio.TimeoutError:
+                await message.in_progress()
+        await task
+    except ValueError as error:
+        LOG.info("runner request already active: %s", error)
+        await message.nak(delay=5)
+    except Exception:
+        LOG.exception("persistent runner request failed")
+        await message.nak(delay=15)
+    else:
+        await message.ack()
+
+
+async def consume_pod_requests(js: object, scheduler: OpenSandboxScheduler) -> None:
+    subscription = await js.pull_subscribe(
+        f"{NATS_SUBJECT_PREFIX}.pod",
+        # Reuse the existing durable so no pending Pod request is stranded
+        # during the in-process consumer migration.
+        durable="pod",
+        stream=NATS_STREAM,
+        config=ConsumerConfig(
+            durable_name="pod",
+            filter_subject=f"{NATS_SUBJECT_PREFIX}.pod",
+            ack_policy=AckPolicy.EXPLICIT,
+            ack_wait=5 * 60,
+            max_ack_pending=4,
+            max_deliver=20,
+        ),
+    )
+    active: set[asyncio.Task[None]] = set()
+    while True:
+        active = {task for task in active if not task.done()}
+        try:
+            messages = await subscription.fetch(
+                batch=max(1, 4 - len(active)), timeout=5
+            )
+        except NatsTimeoutError:
+            continue
+        for message in messages:
+            task = asyncio.create_task(run_message(message, scheduler))
+            active.add(task)
+
+
+async def runtime_context(app: web.Application):
     tokens = RegistrationTokens(REGISTRATION_TOKEN_FILE.read_bytes())
+    tls = ssl.create_default_context(cafile=NATS_CA_FILE)
+    producer = await nats.connect(
+        NATS_URL,
+        user=NATS_PRODUCER_USER,
+        password=NATS_PRODUCER_PASSWORD_FILE.read_text().strip(),
+        tls=tls,
+        name="opensandbox-runner-controller-producer",
+    )
+    worker = await nats.connect(
+        NATS_URL,
+        user=NATS_WORKER_USER,
+        password=NATS_WORKER_PASSWORD_FILE.read_text().strip(),
+        tls=ssl.create_default_context(cafile=NATS_CA_FILE),
+        name="opensandbox-runner-controller-worker",
+    )
+    app["nc"] = producer
+    app["js"] = producer.jetstream()
+    await ensure_stream(app["js"])
     async with OpenSandboxClient(
         api_url=OPENSANDBOX_API,
         api_key_file=OPENSANDBOX_API_KEY_FILE,
@@ -117,13 +249,26 @@ async def opensandbox_context(app: web.Application):
         scheduler = OpenSandboxScheduler(client, tokens)
         app["opensandbox_client"] = client
         app["scheduler"] = scheduler
-        yield
-        await scheduler.close()
+        consumer = None
+        if POD_CONSUMER_ENABLED:
+            consumer = asyncio.create_task(
+                consume_pod_requests(worker.jetstream(), scheduler),
+                name="opensandbox-pod-consumer",
+            )
+        try:
+            yield
+        finally:
+            if consumer is not None:
+                consumer.cancel()
+                await asyncio.gather(consumer, return_exceptions=True)
+            await scheduler.close()
+            await worker.drain()
+            await producer.drain()
 
 
 def create_app() -> web.Application:
     app = web.Application(client_max_size=1024 * 1024)
-    app.cleanup_ctx.append(opensandbox_context)
+    app.cleanup_ctx.append(runtime_context)
     app.router.add_post("/webhook", webhook)
     app.router.add_get("/token/{nonce}", registration_token)
     app.router.add_get("/healthz", health)

src/gitea_dynamic_runner/opensandbox_worker.py

@@ -88,6 +88,7 @@ class OpenSandboxScheduler:
         self.tokens = tokens
         self.on_finished = on_finished
         self.active: dict[int, asyncio.Task[None]] = {}
+        self.runners: dict[str, tuple[int, str, str]] = {}
 
     async def create(self, request: RunnerRequest) -> str:
         if request.job_id in self.active:
@@ -106,11 +107,12 @@ class OpenSandboxScheduler:
             raise
 
         task = asyncio.create_task(
-            self._monitor(request, sandbox_id, nonce),
+            self._monitor(request, sandbox_id, nonce, runner_name),
             name=f"opensandbox-{sandbox_id}",
         )
         task.add_done_callback(self._report)
         self.active[request.job_id] = task
+        self.runners[runner_name] = (request.job_id, sandbox_id, nonce)
         LOG.info(
             "OpenSandbox runner created sandbox=%s runner=%s job_id=%d "
             "repository=%s job_name=%r pool=ci-%s",
@@ -124,7 +126,11 @@ class OpenSandboxScheduler:
         return sandbox_id
 
     async def _monitor(
-        self, request: RunnerRequest, sandbox_id: str, nonce: str
+        self,
+        request: RunnerRequest,
+        sandbox_id: str,
+        nonce: str,
+        runner_name: str,
     ) -> None:
         try:
             deadline = asyncio.get_running_loop().time() + SANDBOX_TIMEOUT
@@ -139,13 +145,38 @@ class OpenSandboxScheduler:
                 await asyncio.sleep(2)
             raise asyncio.TimeoutError(f"OpenSandbox {sandbox_id} timed out")
         finally:
-            await self.tokens.revoke(nonce)
-            try:
-                await self.client.delete(sandbox_id)
-            finally:
-                self.active.pop(request.job_id, None)
-                if self.on_finished is not None:
-                    self.on_finished(request.job_id)
+            await self._cleanup(request.job_id, sandbox_id, nonce, runner_name)
+
+    async def complete(self, runner_name: str) -> bool:
+        """Delete the sandbox which actually ran a completed Gitea job."""
+        state = self.runners.get(runner_name)
+        if state is None:
+            return False
+        job_id, sandbox_id, nonce = state
+        task = self.active.get(job_id)
+        if task is not None:
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        await self._cleanup(job_id, sandbox_id, nonce, runner_name)
+        return True
+
+    async def _cleanup(
+        self,
+        job_id: int,
+        sandbox_id: str,
+        nonce: str,
+        runner_name: str,
+    ) -> None:
+        """Revoke and delete once, including cancellation-before-start races."""
+        if self.runners.pop(runner_name, None) is None:
+            return
+        await self.tokens.revoke(nonce)
+        try:
+            await self.client.delete(sandbox_id)
+        finally:
+            self.active.pop(job_id, None)
+            if self.on_finished is not None:
+                self.on_finished(job_id)
 
     @staticmethod
     def _report(task: asyncio.Task[None]) -> None:

tests/test_opensandbox_worker.py

@@ -26,13 +26,9 @@ def test_sandbox_request_uses_pool_identity_and_one_time_token_url():
         request(), "gitea-vm-abcd", "http://scheduler/token/nonce"
     )
     assert document["pool"] == "ci-vm"
-    assert document["entrypoint"] == [
-        "/usr/local/libexec/gitea-opensandbox-runner"
-    ]
+    assert document["entrypoint"] == ["/usr/local/libexec/gitea-opensandbox-runner"]
     assert document["env"]["GITEA_RUNNER_LABELS"] == "self-hosted:host,vm:host"
-    assert document["env"]["GITEA_RUNNER_REGISTRATION_TOKEN_URL"].endswith(
-        "/nonce"
-    )
+    assert document["env"]["GITEA_RUNNER_REGISTRATION_TOKEN_URL"].endswith("/nonce")
     assert document["env"]["CI_SPIFFE_ID"] == (
         "spiffe://ddupan.top/ci/panxiao81/example/publish-image-3e72cdc4a97e"
     )
@@ -86,3 +82,20 @@ async def test_scheduler_rejects_duplicate_active_job():
     scheduler.active[42] = asyncio.Future()
     with pytest.raises(ValueError, match="already"):
         await scheduler.create(request())
+
+
+async def test_scheduler_deletes_sandbox_for_completed_runner():
+    class RunningOpenSandbox(FakeOpenSandbox):
+        async def get(self, sandbox_id):
+            return {"status": {"state": "Running"}}
+
+    client = RunningOpenSandbox()
+    scheduler = OpenSandboxScheduler(client, RegistrationTokens(b"secret"))
+    await scheduler.create(request())
+    runner_name = next(iter(scheduler.runners))
+
+    assert await scheduler.complete(runner_name) is True
+    assert client.deleted == ["sandbox-1"]
+    assert scheduler.active == {}
+    assert scheduler.runners == {}
+    assert await scheduler.complete(runner_name) is False