healthcheck runner checkup (#376)

The healthcheck function would check if runner.id is defined in order to
retrieve its state from github.

The problem with that is that unless the runner has started a job, we
would never retrieve its id, so idle runner would never be able to get
an updated status and be forever considered as offline by the runner
manager.

To fix this issue, we are using another endpoint to create the
registration token for runners, which allow us to specify the full
runner configuration as well as retrieve the id that will be given.
Thanks to this change we will always have an `runner.id` to work with.

Author: Abubakarr99

images/runner/entrypoint.sh

@@ -2,20 +2,9 @@
 
 echo "Starting runner..."
 echo "RUNNER_ORG: ${RUNNER_ORG}"
-echo "RUNNER_TOKEN: ${RUNNER_TOKEN}"
+echo "RUNNER_JIT_CONFIG: ${RUNNER_JIT_CONFIG}"
 echo "RUNNER_NAME: ${RUNNER_NAME}"
 echo "RUNNER_LABELS: ${RUNNER_LABELS}"
 echo "RUNNER_GROUP: ${RUNNER_GROUP}"
 
-
-./config.sh --url https://github.com/${RUNNER_ORG} \
-    --token "${RUNNER_TOKEN}" \
-    --name "${RUNNER_NAME}" \
-    --labels "${RUNNER_LABELS}" \
-    --runnergroup "${RUNNER_GROUP}" \
-    --work _work \
-    --replace \
-    --unattended \
-    --ephemeral
-
-./run.sh
+./run.sh --jitconfig "${RUNNER_JIT_CONFIG}"

runner_manager/backend/gcloud.py

@@ -1,4 +1,5 @@
 import logging
+import re
 import time
 from typing import Dict, List, Literal
 
@@ -178,27 +179,31 @@ def list(self) -> List[Runner]:
             raise e
         return runners
 
+    def _sanitize_label_value(self, value: str) -> str:
+        value = value[:63]
+        value = value.lower()
+        value = re.sub(r"[^a-z0-9_-]", "-", value)
+        return value
+
     def update(self, runner: Runner) -> Runner:
         try:
             instance: Instance = self.client.get(
                 project=self.config.project_id,
                 zone=self.config.zone,
-                instance=runner.instance_id,
-            )
-            if instance.status != "RUNNING":
-                raise Exception(f"Instance {instance.name} is not running.")
-            labels = {}
-            if runner.labels is not None:
-                labels = {label.name: label.name for label in runner.labels}
-            instance.labels = labels
-            ext_operation: ExtendedOperation = self.client.update(
-                instance_resource=instance
+                instance=runner.instance_id or runner.name,
             )
-            self.wait_for_operation(
+            instance.labels["status"] = self._sanitize_label_value(runner.status)
+            instance.labels["busy"] = self._sanitize_label_value(str(runner.busy))
+
+            log.info(f"Updating {runner.name} labels to {instance.labels}")
+            self.client.update(
+                instance=runner.instance_id or runner.name,
                 project=self.config.project_id,
                 zone=self.config.zone,
-                operation=ext_operation.name,
+                instance_resource=instance,
             )
+            log.info(f"Updated {runner.name} labels to {instance.labels}")
         except Exception as e:
+            super().update(runner)
             raise e
         return super().update(runner)

runner_manager/bin/startup.sh

@@ -20,6 +20,7 @@ RUNNER_NAME=${RUNNER_NAME:-$(hostname)}
 RUNNER_ORG=${RUNNER_ORG:-"org"}
 RUNNER_LABELS=${RUNNER_LABELS:-"runner"}
 RUNNER_TOKEN=${RUNNER_TOKEN:-"token"}
+RUNNER_JIT_CONFIG=${RUNNER_JIT_CONFIG:-""}
 RUNNER_GROUP=${RUNNER_GROUP:-"default"}
 RUNNER_WORKDIR=${RUNNER_WORKDIR:-"_work"}
 RUNNER_DOWNLOAD_URL=${RUNNER_DOWNLOAD_URL:-"https://github.com/actions/runner/releases/download/v2.308.0/actions-runner-linux-x64-2.308.0.tar.gz"}
@@ -131,7 +132,7 @@ Description={{Description}}
 After=network.target
 
 [Service]
-ExecStart=/bin/bash {{RunnerRoot}}/runsvc.sh
+ExecStart=/bin/bash {{RunnerRoot}}/runsvc.sh --jitconfig \"${RUNNER_JIT_CONFIG}\"
 User={{User}}
 WorkingDirectory={{RunnerRoot}}
 KillMode=process
@@ -142,20 +143,11 @@ TimeoutStopSec=5min
 WantedBy=multi-user.target" >/home/actions/actions-runner/bin/actions.runner.service.template
 
 sudo chown -Rh actions:actions /home/actions/actions-runner
-sudo -u actions /home/actions/actions-runner/config.sh \
-	--url "https://github.com/${RUNNER_ORG}" \
-	--token "${RUNNER_TOKEN}" \
-	--name "${RUNNER_NAME}" \
-	--work "${RUNNER_WORKDIR}" \
-	--labels "${RUNNER_LABELS}" \
-	--runnergroup "${RUNNER_GROUP}" \
-	--replace \
-	--unattended \
-	--ephemeral
 
 if command -v systemctl; then
-	sudo ./svc.sh install
-	sudo ./svc.sh start
+	sudo -H -u actions bash -c 'cd /home/actions/actions-runner &&
+	sudo ./svc.sh install &&
+	sudo ./svc.sh start'
 else
-	nohup /home/actions/actions-runner/run.sh 2>/home/actions/actions-runner/logs &
+	nohup /home/actions/actions-runner/run.sh --jitconfig "${RUNNER_JIT_CONFIG}" 2>/home/actions/actions-runner/logs &
 fi

runner_manager/jobs/workflow_job.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from githubkit import Response
-from githubkit.rest.models import AuthenticationToken
 from githubkit.webhooks.models import (
     WorkflowJobCompleted,
     WorkflowJobInProgress,
@@ -51,11 +49,6 @@ def queued(webhook: WorkflowJobQueued) -> str | None:
     log.info(f"Found runner group {runner_group.name}")
     log.info(f"Creating registration token for runner {runner_group.name}")
     github: GitHub = get_github()
-    org = runner_group.organization
-    token_response: Response[
-        AuthenticationToken
-    ] = github.rest.actions.create_registration_token_for_org(org=org)
-    token: AuthenticationToken = token_response.parsed_data
     log.info("Registration token created.")
-    runner: Runner = runner_group.create_runner(token)
-    return runner.pk
+    runner: Runner | None = runner_group.create_runner(github)
+    return runner.pk if runner else None

runner_manager/models/backend.py

@@ -33,7 +33,7 @@ class RunnerEnv(BaseModel):
 
     RUNNER_NAME: Optional[str] = None
     RUNNER_LABELS: Optional[str] = None
-    RUNNER_TOKEN: Optional[str] = None
+    RUNNER_JIT_CONFIG: Optional[str] = None
     RUNNER_ORG: Optional[str] = None
     RUNNER_GROUP: Optional[str] = None
 
@@ -46,7 +46,7 @@ def runner_env(self, runner: Runner) -> RunnerEnv:
         return RunnerEnv(
             RUNNER_NAME=runner.name,
             RUNNER_LABELS=",".join([label.name for label in runner.labels]),
-            RUNNER_TOKEN=runner.token,
+            RUNNER_JIT_CONFIG=runner.encoded_jit_config,
             RUNNER_ORG=runner.organization,
             RUNNER_GROUP=runner.runner_group_name,
         )

runner_manager/models/runner.py

@@ -4,10 +4,12 @@
 
 import redis
 from githubkit.rest.models import Runner as GitHubRunner
+from githubkit.rest.types import OrgsOrgActionsRunnersGenerateJitconfigPostBodyType
 from githubkit.webhooks.types import WorkflowJobEvent
 from pydantic import BaseModel as PydanticBaseModel
 from redis_om import Field, NotFoundError
 
+from runner_manager.clients.github import GitHub
 from runner_manager.logging import log
 from runner_manager.models.base import BaseModel
 
@@ -57,6 +59,7 @@ class Runner(BaseModel):
         default=None,
     )
     token: Optional[str] = None
+    encoded_jit_config: Optional[str] = None
     backend: Optional[str] = Field(index=True, description="Backend type")
     status: RunnerStatus = Field(
         default=RunnerStatus.offline, index=True, full_text_search=True
@@ -143,12 +146,34 @@ def time_to_start_expired(self, timeout: timedelta) -> bool:
     def time_to_live_expired(self, time_to_live: timedelta) -> bool:
         return self.is_online and self.time_since_started > time_to_live
 
-    def update_status(self, github_runner: GitHubRunner):
-        self.status = RunnerStatus(github_runner.status)
-        self.busy = github_runner.busy
+    def update_from_github(self, github: GitHub) -> "Runner":
+        if self.id is not None:
+            github_runner: GitHubRunner = (
+                github.rest.actions.get_self_hosted_runner_for_org(
+                    org=self.organization, runner_id=self.id
+                ).parsed_data
+            )
+            self.status = RunnerStatus(self.status)
+            self.busy = github_runner.busy
         log.info(f"Runner {self.name} status updated to {self.status}")
         return self.save()
 
+    def generate_jit_config(self, github: GitHub) -> "Runner":
+        """Generate JIT config for the runner"""
+        assert self.organization is not None, "Organization name is required"
+        assert self.runner_group_id is not None, "Runner group id is required"
+        jitconfig = github.rest.actions.generate_runner_jitconfig_for_org(
+            org=self.organization,
+            data=OrgsOrgActionsRunnersGenerateJitconfigPostBodyType(
+                name=self.name,
+                runner_group_id=self.runner_group_id,
+                labels=[label.name for label in self.labels],
+            ),
+        ).parsed_data
+        self.id = jitconfig.runner.id
+        self.encoded_jit_config = jitconfig.encoded_jit_config
+        return self.save()
+
     def save(
         self,
         pipeline: Optional[redis.client.Pipeline] = None,

runner_manager/models/runner_group.py

@@ -7,8 +7,6 @@
 import redis
 from githubkit import Response
 from githubkit.exception import RequestFailed
-from githubkit.rest.models import AuthenticationToken
-from githubkit.rest.models import Runner as GitHubRunner
 from githubkit.webhooks.models import WorkflowJobInProgress
 from githubkit.webhooks.types import WorkflowJobEvent
 from pydantic import BaseModel as PydanticBaseModel
@@ -110,27 +108,28 @@ def get_runners(self) -> List[Runner]:
             pass
         return runners
 
-    def create_runner(self, token: AuthenticationToken) -> Runner | None:
+    def create_runner(self, github: GitHub) -> Runner | None:
         """Create a runner instance.
 
         Returns:
             Runner: Runner instance.
         """
         count = len(self.get_runners())
-        if count < self.max:
+        if count < self.max and self.id:
             runner: Runner = Runner(
                 name=self.generate_runner_name(),
                 organization=self.organization,
                 status=RunnerStatus.offline,
-                token=token.token,
                 busy=False,
                 runner_group_id=self.id,
                 created_at=datetime.now(),
                 runner_group_name=self.name,
                 labels=self.runner_labels,
                 manager=self.manager,
             )
-            runner = runner.save()
+            runner.save()
+            runner.generate_jit_config(github)
+
             return self.backend.create(runner)
         return None
 
@@ -160,7 +159,10 @@ def delete_runner(self, runner: Runner) -> int:
 
     @property
     def need_new_runner(self) -> bool:
-        return len(self.get_runners()) < (self.min or 0)
+        runners = self.get_runners()
+        idle = len([runner for runner in runners if runner.busy is False])
+        count = len(runners)
+        return idle < self.min and count < self.max
 
     def create_github_group(self, github: GitHub) -> GitHubRunnerGroup:
         """Create a GitHub runner group."""
@@ -243,26 +245,13 @@ def healthcheck(
         """Healthcheck runner group."""
         runners = self.get_runners()
         for runner in runners:
-
-            if runner.id is not None:
-                github_runner: GitHubRunner = (
-                    github.rest.actions.get_self_hosted_runner_for_org(
-                        self.organization, runner.id
-                    ).parsed_data
-                )
-                runner = runner.update_status(github_runner)
+            runner.update_from_github(github)
             if runner.time_to_live_expired(time_to_live):
                 self.delete_runner(runner)
             if runner.time_to_start_expired(timeout_runner):
                 self.delete_runner(runner)
         while self.need_new_runner:
-            token_response: Response[
-                AuthenticationToken
-            ] = github.rest.actions.create_registration_token_for_org(
-                org=self.organization
-            )
-            token: AuthenticationToken = token_response.parsed_data
-            runner: Runner = self.create_runner(token)
+            runner: Runner = self.create_runner(github)
             if runner:
                 log.info(f"Runner {runner.name} created")
 

tests/unit/conftest.py

@@ -1,14 +1,12 @@
 from datetime import timedelta
 
 import httpx
-from githubkit import Response
 from githubkit.config import Config
-from githubkit.rest.models import AuthenticationToken
 from hypothesis import HealthCheck
 from hypothesis import settings as hypothesis_settings
 from pytest import fixture
 
-from runner_manager import Runner, RunnerGroup
+from runner_manager import Runner
 from runner_manager.clients.github import GitHub
 from runner_manager.models.runner import RunnerLabel
 
@@ -56,13 +54,3 @@ def runner(settings) -> Runner:
     assert runner.Meta.global_key_prefix == settings.name
     Runner.delete_many(Runner.find().all())
     return runner
-
-
-@fixture()
-def runner_token(runner_group: RunnerGroup, github: GitHub) -> AuthenticationToken:
-    token: Response[
-        AuthenticationToken
-    ] = github.rest.actions.create_registration_token_for_org(
-        org=runner_group.organization
-    )
-    return token.parsed_data

tests/unit/jobs/test_healthchecks.py

@@ -1,6 +1,5 @@
 from datetime import datetime, timedelta
 
-from githubkit.rest.models import AuthenticationToken
 from hypothesis import given, settings
 from hypothesis import strategies as st
 from redis_om import Migrator
@@ -44,16 +43,16 @@ def test_healthchecks_hypothesis(
 
 
 def test_group_healthcheck(
-    runner_group: RunnerGroup, settings: Settings, github: GitHub, runner_token
+    runner_group: RunnerGroup, settings: Settings, github: GitHub
 ):
     runner_group.save(github=github)
-    runner_tts: Runner = runner_group.create_runner(runner_token)
+    runner_tts: Runner = runner_group.create_runner(github)
     assert runner_tts is not None
     runner_tts.created_at = datetime.now() - (
         settings.timeout_runner + timedelta(minutes=1)
     )
     runner_tts.save()
-    runner_ttl: Runner = runner_group.create_runner(runner_token)
+    runner_ttl: Runner = runner_group.create_runner(github)
     assert runner_ttl is not None
     runner_ttl.status = RunnerStatus.online
     runner_ttl.started_at = datetime.now() - (
@@ -100,17 +99,17 @@ def test_time_to_live(runner: Runner, settings: Settings):
     assert runner.time_to_live_expired(settings.time_to_live) is False
 
 
-def test_need_new_runner(runner_group: RunnerGroup, runner_token: AuthenticationToken):
+def test_need_new_runner(runner_group: RunnerGroup, github: GitHub):
     runner_group.max = 2
     runner_group.min = 1
     runner_group.save()
     assert runner_group.need_new_runner is True
-    runner_group.create_runner(runner_token)
+    runner_group.create_runner(github)
     assert runner_group.need_new_runner is False
 
 
 def test_healthcheck_job(
-    runner_group: RunnerGroup, settings: Settings, queue: Queue, runner_token
+    runner_group: RunnerGroup, settings: Settings, queue: Queue, github: GitHub
 ):
     runner_group.save()
     queue.enqueue(
@@ -120,7 +119,7 @@ def test_healthcheck_job(
         settings.timeout_runner,
     )
     assert len(runner_group.get_runners()) == 0
-    runner_group.create_runner(runner_token)
+    runner_group.create_runner(github)
     queue.enqueue(
         healthcheck.group,
         runner_group.pk,