gitea-dynamic-runner/src/gitea_dynamic_runner/opensandbox_worker.py

"""Direct OpenSandbox lifecycle scheduler used by the webhook controller."""

from __future__ import annotations

import asyncio
import logging
import os
import secrets
import uuid
from collections.abc import Callable

from .models import RunnerRequest
from .opensandbox import OpenSandboxClient
from .pod_worker import identity_path


LOG = logging.getLogger(__name__)
SANDBOX_TIMEOUT = int(os.environ.get("RUNNER_SANDBOX_TIMEOUT", str(4 * 60 * 60)))
SPIFFE_TRUST_DOMAIN = os.environ.get("SPIFFE_TRUST_DOMAIN", "ddupan.top")
TOKEN_BASE_URL = os.environ.get(
    "RUNNER_TOKEN_BASE_URL", "http://192.168.10.127:8787/token"
).rstrip("/")


def sandbox_request(
    request: RunnerRequest,
    runner_name: str,
    registration_token_url: str,
) -> dict[str, object]:
    path = identity_path(request.repository, request.job_name)
    return {
        "pool": f"ci-{request.backend}",
        "timeout": SANDBOX_TIMEOUT,
        "entrypoint": ["/usr/local/libexec/gitea-opensandbox-runner"],
        "env": {
            "GITEA_RUNNER_NAME": runner_name,
            "GITEA_RUNNER_LABELS": f"self-hosted:host,{request.backend}:host",
            "GITEA_RUNNER_REGISTRATION_TOKEN_URL": registration_token_url,
            "GITEA_RUNNER_EPHEMERAL": "1",
            "GITEA_RUNNER_ONCE": "1",
            "CONFIG_FILE": "/etc/gitea-runner/config.yaml",
            "CI_SPIFFE_ID": f"spiffe://{SPIFFE_TRUST_DOMAIN}/ci/{path}",
            "SPIFFE_ENDPOINT_SOCKET": (
                "unix:///run/spire/agent-sockets/spire-agent.sock"
            ),
        },
        "metadata": {
            "ci.ddupan.top/runner": "true",
            "ci.ddupan.top/job-id": str(request.job_id),
            "ci.ddupan.top/run-id": str(request.run_id),
            "ci.ddupan.top/runner-name": runner_name,
        },
    }


class RegistrationTokens:
    """Single-use registration-token URLs; values never enter Sandbox CRs."""

    def __init__(self, token: bytes) -> None:
        self._token = token.strip()
        self._values: dict[str, bytes] = {}
        self._lock = asyncio.Lock()

    async def issue(self) -> tuple[str, str]:
        nonce = secrets.token_urlsafe(32)
        async with self._lock:
            self._values[nonce] = self._token
        return nonce, f"{TOKEN_BASE_URL}/{nonce}"

    async def consume(self, nonce: str) -> bytes | None:
        async with self._lock:
            return self._values.pop(nonce, None)

    async def revoke(self, nonce: str) -> None:
        async with self._lock:
            self._values.pop(nonce, None)


class OpenSandboxScheduler:
    def __init__(
        self,
        client: OpenSandboxClient,
        tokens: RegistrationTokens,
        *,
        on_finished: Callable[[int], None] | None = None,
    ) -> None:
        self.client = client
        self.tokens = tokens
        self.on_finished = on_finished
        self.active: dict[int, asyncio.Task[None]] = {}

    async def create(self, request: RunnerRequest) -> str:
        if request.job_id in self.active:
            raise ValueError(f"job {request.job_id} already has an active sandbox")
        runner_name = f"gitea-{request.backend}-{uuid.uuid4().hex[:12]}"
        nonce, token_url = await self.tokens.issue()
        try:
            response = await self.client.create(
                **sandbox_request(request, runner_name, token_url)
            )
            sandbox_id = response.get("id")
            if not isinstance(sandbox_id, str) or not sandbox_id:
                raise RuntimeError("OpenSandbox create response has no id")
        except Exception:
            await self.tokens.revoke(nonce)
            raise

        task = asyncio.create_task(
            self._monitor(request, sandbox_id, nonce),
            name=f"opensandbox-{sandbox_id}",
        )
        task.add_done_callback(self._report)
        self.active[request.job_id] = task
        LOG.info(
            "OpenSandbox runner created sandbox=%s runner=%s job_id=%d "
            "repository=%s job_name=%r pool=ci-%s",
            sandbox_id,
            runner_name,
            request.job_id,
            request.repository,
            request.job_name,
            request.backend,
        )
        return sandbox_id

    async def _monitor(
        self, request: RunnerRequest, sandbox_id: str, nonce: str
    ) -> None:
        try:
            deadline = asyncio.get_running_loop().time() + SANDBOX_TIMEOUT
            while asyncio.get_running_loop().time() < deadline:
                sandbox = await self.client.get(sandbox_id)
                if sandbox is None:
                    return
                status = sandbox.get("status")
                state = status.get("state") if isinstance(status, dict) else None
                if state in {"Terminated", "Failed"}:
                    return
                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)

    @staticmethod
    def _report(task: asyncio.Task[None]) -> None:
        if not task.cancelled() and (error := task.exception()) is not None:
            LOG.error(
                "OpenSandbox lifecycle task failed",
                exc_info=(type(error), error, error.__traceback__),
            )

    async def close(self) -> None:
        tasks = list(self.active.values())
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)