Skip to content

Commit 4e9fc70

Browse files
hskibahskiba
committed
Use JIT runner config to skip config.sh registration
Generate JIT config from Lambda via GitHub API instead of passing PAT to the instance. This eliminates the 15-30 second config.sh registration step on the runner. Changes: - main.go: Add generateJITConfig() to call GitHub's JIT config API - main.go: Build labels list and pass JIT config to user-data template - user-data.sh: Remove get_github_token() and config.sh steps - user-data.sh: Use ./run.sh --jitconfig instead
1 parent f6dd725 commit 4e9fc70

2 files changed

Lines changed: 101 additions & 71 deletions

File tree

‎main.go‎

Lines changed: 93 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"encoding/json"
99
"errors"
1010
"fmt"
11+
"io"
1112
"log/slog"
1213
"net/http"
1314
"os"
@@ -361,6 +362,73 @@ func getLaunchConfig() (LaunchConfig, error) {
361362
}, nil
362363
}
363364

365+
// JITConfigRequest represents the request body for generating a JIT runner config.
366+
type JITConfigRequest struct {
367+
Name string `json:"name"`
368+
RunnerGroupID int `json:"runner_group_id"`
369+
Labels []string `json:"labels"`
370+
WorkFolder string `json:"work_folder"`
371+
}
372+
373+
// JITConfigResponse represents the response from the JIT config API.
374+
type JITConfigResponse struct {
375+
Runner struct {
376+
ID int `json:"id"`
377+
Name string `json:"name"`
378+
} `json:"runner"`
379+
EncodedJITConfig string `json:"encoded_jit_config"`
380+
}
381+
382+
// generateJITConfig calls the GitHub API to generate a JIT runner configuration.
383+
// This eliminates the need for config.sh on the runner, saving 15-30 seconds.
384+
func generateJITConfig(pat, org, runnerName string, labels []string) (*JITConfigResponse, error) {
385+
reqBody := JITConfigRequest{
386+
Name: runnerName,
387+
RunnerGroupID: 1, // Default runner group
388+
Labels: labels,
389+
WorkFolder: "_work",
390+
}
391+
392+
jsonBody, err := json.Marshal(reqBody)
393+
if err != nil {
394+
return nil, fmt.Errorf("failed to marshal JIT config request: %w", err)
395+
}
396+
397+
url := fmt.Sprintf("https://api.github.com/orgs/%s/actions/runners/generate-jitconfig", org)
398+
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
399+
if err != nil {
400+
return nil, fmt.Errorf("failed to create request: %w", err)
401+
}
402+
403+
req.Header.Set("Accept", "application/vnd.github+json")
404+
req.Header.Set("Authorization", "Bearer "+pat)
405+
req.Header.Set("X-GitHub-Api-Version", "2022-11-28")
406+
req.Header.Set("Content-Type", "application/json")
407+
408+
client := &http.Client{Timeout: 30 * time.Second}
409+
resp, err := client.Do(req)
410+
if err != nil {
411+
return nil, fmt.Errorf("failed to call JIT config API: %w", err)
412+
}
413+
defer resp.Body.Close()
414+
415+
body, err := io.ReadAll(resp.Body)
416+
if err != nil {
417+
return nil, fmt.Errorf("failed to read response body: %w", err)
418+
}
419+
420+
if resp.StatusCode != http.StatusCreated {
421+
return nil, fmt.Errorf("JIT config API returned status %d: %s", resp.StatusCode, string(body))
422+
}
423+
424+
var jitResp JITConfigResponse
425+
if err := json.Unmarshal(body, &jitResp); err != nil {
426+
return nil, fmt.Errorf("failed to parse JIT config response: %w", err)
427+
}
428+
429+
return &jitResp, nil
430+
}
431+
364432
// handleMaintenance processes scheduled warm pool maintenance events.
365433
func handleMaintenance() error {
366434
slog.Info("warm pool maintenance triggered")
@@ -534,20 +602,42 @@ func handleWebhook(request events.APIGatewayProxyRequest) (events.APIGatewayProx
534602
}
535603
}
536604

537-
slog.Info("processing job", "instanceType", instanceType, "jobID", event.GetWorkflowJob().GetID())
605+
jobEventID := event.GetWorkflowJob().GetID()
606+
runnerName := fmt.Sprintf("ephemeral-i-%d", jobEventID)
607+
608+
slog.Info("processing job", "instanceType", instanceType, "jobID", jobEventID, "runnerName", runnerName)
609+
610+
// Build labels for the runner
611+
labels := []string{string(instanceType), "ephemeral", "X64"}
612+
if extraLabels != "" {
613+
// extraLabels already has leading comma, split and add non-empty labels
614+
for _, label := range strings.Split(extraLabels, ",") {
615+
if label = strings.TrimSpace(label); label != "" {
616+
labels = append(labels, label)
617+
}
618+
}
619+
}
620+
621+
// Generate JIT config from GitHub API (eliminates need for config.sh on instance)
622+
jitConfig, err := generateJITConfig(pat, "frgrisk", runnerName, labels)
623+
if err != nil {
624+
slog.Error("failed to generate JIT config", "error", err.Error())
625+
return events.APIGatewayProxyResponse{StatusCode: http.StatusInternalServerError}, err
626+
}
627+
628+
slog.Info("generated JIT config", "runnerID", jitConfig.Runner.ID, "runnerName", jitConfig.Runner.Name)
538629

539630
tpl, err := template.New("userdata").Parse(userData)
540631
if err != nil {
541632
return events.APIGatewayProxyResponse{StatusCode: http.StatusInternalServerError}, err
542633
}
543634

544635
var buf bytes.Buffer
545-
if err := tpl.Execute(&buf, map[string]string{"GitHubPAT": pat, "ExtraLabels": extraLabels}); err != nil {
636+
if err := tpl.Execute(&buf, map[string]string{"JITConfig": jitConfig.EncodedJITConfig}); err != nil {
546637
return events.APIGatewayProxyResponse{StatusCode: http.StatusInternalServerError}, err
547638
}
548639

549640
finalUserData := buf.String()
550-
jobEventID := event.GetWorkflowJob().GetID()
551641

552642
// Get warm pool target size for this instance type
553643
poolConfig := parseWarmPoolConfig()

‎user-data.sh‎

Lines changed: 8 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -86,88 +86,28 @@ INSTANCE_TYPE=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169
8686

8787
log_to_cloudwatch "INFO" "Instance: ${INSTANCE_ID}, Type: ${INSTANCE_TYPE}"
8888

89-
# Function to get GitHub registration token with retry
90-
get_github_token() {
91-
local max_attempts=5
92-
local attempt=1
93-
local delay=5
94-
95-
while [ $attempt -le $max_attempts ]; do
96-
log_to_cloudwatch "INFO" "Attempting to get GitHub registration token (attempt ${attempt}/${max_attempts})"
97-
98-
GITHUB_TOKEN=$(curl -s -L \
99-
-X POST \
100-
-H "Accept: application/vnd.github+json" \
101-
-H "Authorization: Bearer {{.GitHubPAT}}" \
102-
-H "X-GitHub-Api-Version: 2022-11-28" \
103-
https://api.github.com/orgs/frgrisk/actions/runners/registration-token | jq -r .token)
104-
105-
if [ -n "$GITHUB_TOKEN" ] && [ "$GITHUB_TOKEN" != "null" ]; then
106-
log_to_cloudwatch "INFO" "Successfully obtained GitHub registration token"
107-
return 0
108-
fi
109-
110-
log_to_cloudwatch "WARN" "Failed to get GitHub token, retrying in ${delay} seconds..."
111-
sleep $delay
112-
delay=$((delay * 2))
113-
attempt=$((attempt + 1))
114-
done
115-
116-
log_to_cloudwatch "ERROR" "Failed to get GitHub registration token after ${max_attempts} attempts"
117-
return 1
118-
}
89+
# JIT config is passed from Lambda - no need to call GitHub API or run config.sh
90+
JIT_CONFIG="{{.JITConfig}}"
11991

120-
# Get GitHub registration token
121-
if ! get_github_token; then
122-
log_to_cloudwatch "ERROR" "Unable to proceed without registration token"
92+
if [ -z "$JIT_CONFIG" ] || [ "$JIT_CONFIG" = "{{.JITConfig}}" ]; then
93+
log_to_cloudwatch "ERROR" "JIT config not provided"
12394
shutdown now
12495
exit 1
12596
fi
12697

127-
# Configure runner with retry
128-
log_to_cloudwatch "INFO" "Configuring GitHub runner"
129-
max_config_attempts=3
130-
config_attempt=1
131-
132-
while [ $config_attempt -le $max_config_attempts ]; do
133-
if sudo -u ubuntu ./config.sh \
134-
--url https://github.com/frgrisk \
135-
--token "$GITHUB_TOKEN" \
136-
--disableupdate \
137-
--ephemeral \
138-
--labels "${INSTANCE_TYPE},ephemeral,X64{{.ExtraLabels}}" \
139-
--unattended \
140-
--name "ephemeral-${INSTANCE_ID}" \
141-
--work _work; then
142-
143-
log_to_cloudwatch "INFO" "Runner configured successfully"
144-
break
145-
else
146-
log_to_cloudwatch "WARN" "Runner configuration failed (attempt ${config_attempt}/${max_config_attempts})"
147-
config_attempt=$((config_attempt + 1))
148-
if [ $config_attempt -le $max_config_attempts ]; then
149-
sleep 10
150-
fi
151-
fi
152-
done
153-
154-
if [ $config_attempt -gt $max_config_attempts ]; then
155-
log_to_cloudwatch "ERROR" "Failed to configure runner after ${max_config_attempts} attempts"
156-
shutdown now
157-
exit 1
158-
fi
98+
log_to_cloudwatch "INFO" "JIT config received, skipping config.sh"
15999

160100
END_TIME=$(date +%s)
161101
EXECUTION_TIME=$((END_TIME - START_TIME))
162102
log_to_cloudwatch "INFO" "Setup completed in ${EXECUTION_TIME} seconds"
163103

164-
# Start the runner and wait for it to complete
165-
log_to_cloudwatch "INFO" "Starting GitHub runner"
104+
# Start the runner with JIT config (skips registration entirely)
105+
log_to_cloudwatch "INFO" "Starting GitHub runner with JIT config"
166106

167107
# Create a temporary file to capture runner output
168108
RUNNER_LOG=$(mktemp /tmp/runner-output.XXXXXX)
169109

170-
if sudo -u ubuntu ./run.sh 2>&1 | tee "${RUNNER_LOG}"; then
110+
if sudo -u ubuntu ./run.sh --jitconfig "$JIT_CONFIG" 2>&1 | tee "${RUNNER_LOG}"; then
171111
log_to_cloudwatch "INFO" "Runner completed successfully"
172112
else
173113
EXIT_CODE=$?

0 commit comments

Comments
 (0)