From 5348ca9891e5eede742b2bc26ce902216e900f9f Mon Sep 17 00:00:00 2001 From: panxiao81 Date: Fri, 18 Sep 2026 19:25:09 +0000 Subject: [PATCH] =?UTF-8?q?=E5=B0=86=20OpenSandbox=20=E8=B0=83=E5=BA=A6?= =?UTF-8?q?=E6=94=B9=E4=B8=BA=E6=8C=81=E4=B9=85=E4=BA=8B=E4=BB=B6=E5=BE=AA?= =?UTF-8?q?=E7=8E=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gitea_dynamic_runner/controller.py | 185 ++++++++++++++++-- .../opensandbox_worker.py | 49 ++++- tests/test_opensandbox_worker.py | 25 ++- 3 files changed, 224 insertions(+), 35 deletions(-) diff --git a/src/gitea_dynamic_runner/controller.py b/src/gitea_dynamic_runner/controller.py index d464d04..cccc368 100644 --- a/src/gitea_dynamic_runner/controller.py +++ b/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) diff --git a/src/gitea_dynamic_runner/opensandbox_worker.py b/src/gitea_dynamic_runner/opensandbox_worker.py index 88baa55..3cb274e 100644 --- a/src/gitea_dynamic_runner/opensandbox_worker.py +++ b/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: diff --git a/tests/test_opensandbox_worker.py b/tests/test_opensandbox_worker.py index f623df2..ef5fe2a 100644 --- a/tests/test_opensandbox_worker.py +++ b/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