|
|
|
@@ -1,5 +1,5 @@
|
|
|
|
|
#!/usr/bin/env python3
|
|
|
|
|
"""Gitea workflow_job webhook to the OpenSandbox Lifecycle API."""
|
|
|
|
|
"""Persist Gitea workflow_job events and optionally schedule VM sandboxes."""
|
|
|
|
|
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
|
@@ -25,7 +25,7 @@ from nats.js.api import (
|
|
|
|
|
)
|
|
|
|
|
from nats.js.errors import NotFoundError
|
|
|
|
|
|
|
|
|
|
from .models import RunnerRequest
|
|
|
|
|
from .models import IdentityBinding, RunnerRequest
|
|
|
|
|
from .opensandbox import OpenSandboxClient
|
|
|
|
|
from .opensandbox_worker import OpenSandboxScheduler, RegistrationTokens
|
|
|
|
|
|
|
|
|
@@ -56,7 +56,8 @@ NATS_WORKER_PASSWORD_FILE = Path(
|
|
|
|
|
)
|
|
|
|
|
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"
|
|
|
|
|
VM_CONSUMER_ENABLED = os.environ.get("VM_CONSUMER_ENABLED", "false") == "true"
|
|
|
|
|
VM_CAPACITY = int(os.environ.get("VM_CAPACITY", "1"))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def accepts(payload: object) -> tuple[bool, str | None]:
|
|
|
|
@@ -82,21 +83,16 @@ 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-")
|
|
|
|
|
scheduler = request.app.get("scheduler")
|
|
|
|
|
if (
|
|
|
|
|
isinstance(runner_name, str)
|
|
|
|
|
and runner_name.startswith("gitea-vm-")
|
|
|
|
|
and scheduler is not None
|
|
|
|
|
):
|
|
|
|
|
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,
|
|
|
|
@@ -107,8 +103,20 @@ async def webhook(request: web.Request) -> web.Response:
|
|
|
|
|
|
|
|
|
|
runner_request = RunnerRequest.from_webhook(payload)
|
|
|
|
|
if runner_request is None:
|
|
|
|
|
binding = IdentityBinding.from_webhook(payload)
|
|
|
|
|
if binding is None:
|
|
|
|
|
LOG.info("workflow_job webhook ignored %s", context)
|
|
|
|
|
return web.Response(status=204)
|
|
|
|
|
subject = f"{NATS_SUBJECT_PREFIX}.{binding.backend}.binding"
|
|
|
|
|
await request.app["js"].publish(
|
|
|
|
|
subject,
|
|
|
|
|
binding.to_json(),
|
|
|
|
|
headers={
|
|
|
|
|
"Nats-Msg-Id": f"gitea-workflow-job-{binding.job_id}-in-progress"
|
|
|
|
|
},
|
|
|
|
|
)
|
|
|
|
|
LOG.info("runner identity binding persisted subject=%s %s", subject, context)
|
|
|
|
|
return web.Response(status=202, text="binding queued\n")
|
|
|
|
|
|
|
|
|
|
subject = f"{NATS_SUBJECT_PREFIX}.{runner_request.backend}"
|
|
|
|
|
await request.app["js"].publish(
|
|
|
|
@@ -121,7 +129,9 @@ async def webhook(request: web.Request) -> web.Response:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def registration_token(request: web.Request) -> web.Response:
|
|
|
|
|
scheduler: OpenSandboxScheduler = request.app["scheduler"]
|
|
|
|
|
scheduler = request.app.get("scheduler")
|
|
|
|
|
if scheduler is None:
|
|
|
|
|
raise web.HTTPNotFound()
|
|
|
|
|
value = await scheduler.tokens.consume(request.match_info["nonce"])
|
|
|
|
|
if value is None:
|
|
|
|
|
raise web.HTTPNotFound()
|
|
|
|
@@ -144,10 +154,10 @@ def _webhook_context(payload: object) -> str:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def health(request: web.Request) -> web.Response:
|
|
|
|
|
client: OpenSandboxClient = request.app["opensandbox_client"]
|
|
|
|
|
connected = request.app["nc"].is_connected
|
|
|
|
|
return web.Response(
|
|
|
|
|
text="ok\n" if client.session is not None else "disconnected\n",
|
|
|
|
|
status=200 if client.session is not None else 503,
|
|
|
|
|
text="ok\n" if connected else "disconnected\n",
|
|
|
|
|
status=200 if connected else 503,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@@ -200,19 +210,17 @@ async def run_message(message: object, scheduler: OpenSandboxScheduler) -> None:
|
|
|
|
|
await message.ack()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def consume_pod_requests(js: object, scheduler: OpenSandboxScheduler) -> None:
|
|
|
|
|
async def consume_vm_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",
|
|
|
|
|
f"{NATS_SUBJECT_PREFIX}.vm",
|
|
|
|
|
durable="vm",
|
|
|
|
|
stream=NATS_STREAM,
|
|
|
|
|
config=ConsumerConfig(
|
|
|
|
|
durable_name="pod",
|
|
|
|
|
filter_subject=f"{NATS_SUBJECT_PREFIX}.pod",
|
|
|
|
|
durable_name="vm",
|
|
|
|
|
filter_subject=f"{NATS_SUBJECT_PREFIX}.vm",
|
|
|
|
|
ack_policy=AckPolicy.EXPLICIT,
|
|
|
|
|
ack_wait=5 * 60,
|
|
|
|
|
max_ack_pending=4,
|
|
|
|
|
max_ack_pending=VM_CAPACITY,
|
|
|
|
|
max_deliver=20,
|
|
|
|
|
),
|
|
|
|
|
)
|
|
|
|
@@ -221,7 +229,7 @@ async def consume_pod_requests(js: object, scheduler: OpenSandboxScheduler) -> N
|
|
|
|
|
active = {task for task in active if not task.done()}
|
|
|
|
|
try:
|
|
|
|
|
messages = await subscription.fetch(
|
|
|
|
|
batch=max(1, 4 - len(active)), timeout=5
|
|
|
|
|
batch=max(1, VM_CAPACITY - len(active)), timeout=5
|
|
|
|
|
)
|
|
|
|
|
except NatsTimeoutError:
|
|
|
|
|
continue
|
|
|
|
@@ -231,7 +239,6 @@ async def consume_pod_requests(js: object, scheduler: OpenSandboxScheduler) -> N
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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,
|
|
|
|
@@ -240,16 +247,26 @@ async def runtime_context(app: web.Application):
|
|
|
|
|
tls=tls,
|
|
|
|
|
name="opensandbox-runner-controller-producer",
|
|
|
|
|
)
|
|
|
|
|
app["nc"] = producer
|
|
|
|
|
app["js"] = producer.jetstream()
|
|
|
|
|
await ensure_stream(app["js"])
|
|
|
|
|
if not VM_CONSUMER_ENABLED:
|
|
|
|
|
try:
|
|
|
|
|
yield
|
|
|
|
|
finally:
|
|
|
|
|
await producer.drain()
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
if VM_CAPACITY < 1:
|
|
|
|
|
raise ValueError("VM_CAPACITY must be at least 1")
|
|
|
|
|
tokens = RegistrationTokens(REGISTRATION_TOKEN_FILE.read_bytes())
|
|
|
|
|
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",
|
|
|
|
|
name="opensandbox-vm-runner-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,
|
|
|
|
@@ -257,16 +274,13 @@ async def runtime_context(app: web.Application):
|
|
|
|
|
scheduler = OpenSandboxScheduler(client, tokens)
|
|
|
|
|
app["opensandbox_client"] = client
|
|
|
|
|
app["scheduler"] = scheduler
|
|
|
|
|
consumer = None
|
|
|
|
|
if POD_CONSUMER_ENABLED:
|
|
|
|
|
consumer = asyncio.create_task(
|
|
|
|
|
consume_pod_requests(worker.jetstream(), scheduler),
|
|
|
|
|
name="opensandbox-pod-consumer",
|
|
|
|
|
consume_vm_requests(worker.jetstream(), scheduler),
|
|
|
|
|
name="opensandbox-vm-consumer",
|
|
|
|
|
)
|
|
|
|
|
try:
|
|
|
|
|
yield
|
|
|
|
|
finally:
|
|
|
|
|
if consumer is not None:
|
|
|
|
|
consumer.cancel()
|
|
|
|
|
await asyncio.gather(consumer, return_exceptions=True)
|
|
|
|
|
await scheduler.close()
|
|
|
|
|