Initial commit for open source release 🚀
Co-authored-by: Francesco Renzi <[email protected]> Co-authored-by: Nikola Jokic <[email protected]>
This commit is contained in:
Francesco Renzi
co-authored by
Francesco Renzi
Nikola Jokic
Francesco Renzi
Nikola Jokic
commit
e4a017ce06
@@ -0,0 +1,55 @@
|
||||
# Docker Runner Scale Set Example
|
||||
|
||||
This example showcases a Docker implementation of GitHub Actions runner scale sets, using the `github.com/actions/scaleset` client to provision ephemeral GitHub Actions runners as Docker containers.
|
||||
|
||||
The goal of this example is to show how simple and powerful it is when you only need to focus on the core logic of scaling runners up and down, while the client handles all the API interactions.
|
||||
|
||||
> [!WARNING]
|
||||
> This is a simplified example meant for demonstration and learning purposes. It is not intended for production use.
|
||||
|
||||
> [!NOTE]
|
||||
> When exiting normally all runners and the scale set itself are cleaned up automatically.
|
||||
|
||||
## Getting started
|
||||
|
||||
You can install the example with:
|
||||
|
||||
```bash
|
||||
go install github.com/actions/scaleset/examples/dockerscaleset@latest
|
||||
```
|
||||
|
||||
If this fails you should also try running the command with
|
||||
|
||||
```bash
|
||||
GONOSUMDB=github.com/actions/scaleset GOPRIVATE=github.com/actions/scaleset go
|
||||
install github.com/actions/scaleset/examples/dockerscaleset@latest
|
||||
```
|
||||
|
||||
You'll then need:
|
||||
|
||||
- Docker installed and running on your machine.
|
||||
- A URL for the target repository, organization, or enterprise where you want to register your scale set.
|
||||
- [Credentials that have access to the above target](https://docs.github.com/en/actions/tutorials/use-actions-runner-controller/authenticate-to-the-api): you can use either a GitHub App (recommended) or a Personal Access Token (PAT).
|
||||
- A name for your scale set (this must be unique within the runner group the scale set is created in).
|
||||
|
||||
---
|
||||
|
||||
## Flags
|
||||
|
||||
| Flag | Required | Description |
|
||||
|------|----------|-------------|
|
||||
| `--url` | Yes | Registration target (org, repo, or enterprise URL, e.g. `https://github.com/org/repo`). |
|
||||
| `--name` | Yes | Runner scale set name (must be unique within the runner group). |
|
||||
| `--labels` | No | Labels for workflow targeting (comma-separated or repeated). Defaults to `--name` if not provided. |
|
||||
| `--max-runners` | No | Upper bound of concurrently provisioned runners (default 10). |
|
||||
| `--min-runners` | No | Lower bound to maintain (default 0). |
|
||||
| `--runner-group` | No | Runner group name (default `default`). |
|
||||
| `--app-client-id` | Cond.* | GitHub App Client (App) ID. |
|
||||
| `--app-installation-id` | Cond.* | GitHub App Installation ID. |
|
||||
| `--app-private-key` | Cond.* | GitHub App private key PEM contents. |
|
||||
| `--token` | Cond.* | Personal Access Token (alternative to App). |
|
||||
| `--log-level` | No | `debug`, `info`, `warn`, `error` (default `info`). |
|
||||
| `--log-format` | No | `text`, `json`, or `none` (any invalid → no logs). |
|
||||
| `--runner-image` | No | Override container image (defaults to latest official). |
|
||||
|
||||
*Provide either App credentials (all three) OR a PAT.*
|
||||
@@ -0,0 +1,141 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/actions/scaleset"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
RegistrationURL string
|
||||
MaxRunners int
|
||||
MinRunners int
|
||||
ScaleSetName string
|
||||
Labels []string
|
||||
RunnerGroup string
|
||||
GitHubApp scaleset.GitHubAppAuth
|
||||
Token string
|
||||
RunnerImage string
|
||||
LogLevel string
|
||||
LogFormat string
|
||||
}
|
||||
|
||||
func (c *Config) defaults() {
|
||||
if c.RunnerGroup == "" {
|
||||
c.RunnerGroup = scaleset.DefaultRunnerGroup
|
||||
}
|
||||
if c.RunnerImage == "" {
|
||||
c.RunnerImage = "ghcr.io/actions/actions-runner:latest"
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) Validate() error {
|
||||
c.defaults()
|
||||
|
||||
if _, err := url.ParseRequestURI(c.RegistrationURL); err != nil {
|
||||
return fmt.Errorf("invalid registration URL: %w, it should be the full URL of where you want to register your scale set, e.g. 'https://github.com/org/repo'", err)
|
||||
}
|
||||
|
||||
appError := c.GitHubApp.Validate()
|
||||
if c.Token == "" && appError != nil {
|
||||
return fmt.Errorf("no credentials provided: either GitHub App (client id, installation id and private key) (recommended) or a Personal Access Token are required")
|
||||
}
|
||||
|
||||
if c.ScaleSetName == "" {
|
||||
return fmt.Errorf("scale set name is required")
|
||||
}
|
||||
for i, label := range c.Labels {
|
||||
if strings.TrimSpace(label) == "" {
|
||||
return fmt.Errorf("label at index %d is empty", i)
|
||||
}
|
||||
}
|
||||
if c.MaxRunners < c.MinRunners {
|
||||
return fmt.Errorf("max runners cannot be less than min-runners")
|
||||
}
|
||||
if c.RunnerGroup == "" {
|
||||
return fmt.Errorf("runner group is required")
|
||||
}
|
||||
if c.RunnerImage == "" {
|
||||
return fmt.Errorf("runner image is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// systemInfo serves as a base system info
|
||||
func systemInfo(scaleSetID int) scaleset.SystemInfo {
|
||||
return scaleset.SystemInfo{
|
||||
System: "dockerscaleset",
|
||||
Subsystem: "dockerscaleset",
|
||||
CommitSHA: "NA", // You can leverage build flags to set commit SHA
|
||||
Version: "0.1.0", // You can leverage build flags to set version
|
||||
ScaleSetID: scaleSetID,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) ScalesetClient() (*scaleset.Client, error) {
|
||||
if err := c.GitHubApp.Validate(); err == nil {
|
||||
return scaleset.NewClientWithGitHubApp(
|
||||
scaleset.ClientWithGitHubAppConfig{
|
||||
GitHubConfigURL: c.RegistrationURL,
|
||||
GitHubAppAuth: c.GitHubApp,
|
||||
SystemInfo: systemInfo(0),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
return scaleset.NewClientWithPersonalAccessToken(
|
||||
scaleset.NewClientWithPersonalAccessTokenConfig{
|
||||
GitHubConfigURL: c.RegistrationURL,
|
||||
PersonalAccessToken: c.Token,
|
||||
SystemInfo: systemInfo(0),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (c *Config) Logger() *slog.Logger {
|
||||
var lvl slog.Level
|
||||
switch strings.ToLower(c.LogLevel) {
|
||||
case "debug":
|
||||
lvl = slog.LevelDebug
|
||||
case "info":
|
||||
lvl = slog.LevelInfo
|
||||
case "warn":
|
||||
lvl = slog.LevelWarn
|
||||
case "error":
|
||||
lvl = slog.LevelError
|
||||
default:
|
||||
lvl = slog.LevelInfo
|
||||
}
|
||||
|
||||
switch c.LogFormat {
|
||||
case "json":
|
||||
return slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
|
||||
AddSource: true,
|
||||
Level: lvl,
|
||||
}))
|
||||
case "text":
|
||||
return slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
|
||||
AddSource: true,
|
||||
Level: lvl,
|
||||
}))
|
||||
default:
|
||||
return slog.New(slog.DiscardHandler)
|
||||
}
|
||||
}
|
||||
|
||||
// BuildLabels returns the labels to use for the runner scale set.
|
||||
// If custom labels are provided, those are used; otherwise, the scale set name is used as the label.
|
||||
func (c *Config) BuildLabels() []scaleset.Label {
|
||||
if len(c.Labels) > 0 {
|
||||
labels := make([]scaleset.Label, len(c.Labels))
|
||||
for i, name := range c.Labels {
|
||||
labels[i] = scaleset.Label{Name: strings.TrimSpace(name)}
|
||||
}
|
||||
return labels
|
||||
}
|
||||
return []scaleset.Label{{Name: c.ScaleSetName}}
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-github/v79/github"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestE2E(t *testing.T) {
|
||||
if os.Getenv("E2E") != "true" {
|
||||
t.Skip("Skipping E2E test; set E2E=true to run")
|
||||
}
|
||||
|
||||
configURL := mustGetEnv(t, "E2E_SCALESET_URL")
|
||||
name := mustGetEnv(t, "E2E_SCALESET_NAME")
|
||||
|
||||
workflowEnv := mustE2EWorkflowEnv(t, name)
|
||||
runArgs := mustE2ECommandArgs(t, configURL, name)
|
||||
|
||||
tempDir, err := os.MkdirTemp("", "e2e-dockerscaleset-")
|
||||
require.NoError(t, err, "Failed to create temp dir")
|
||||
defer os.RemoveAll(tempDir)
|
||||
|
||||
binaryPath := filepath.Join(tempDir, "dockerscaleset")
|
||||
|
||||
// Build the dockerscaleset binary in temp dir
|
||||
{
|
||||
cmd := exec.Command("go", "build", "-o", binaryPath, ".")
|
||||
output, err := cmd.CombinedOutput()
|
||||
require.NoError(t, err, "Failed to build dockerscaleset: %s", output)
|
||||
}
|
||||
|
||||
// Fatal channel
|
||||
testErrCh := make(chan error, 2)
|
||||
|
||||
runCmd := exec.Command(binaryPath, runArgs...)
|
||||
stdout, err := runCmd.StdoutPipe()
|
||||
runCmd.Stderr = os.Stderr
|
||||
require.NoError(t, err, "Failed to get stdout pipe")
|
||||
err = runCmd.Start()
|
||||
require.NoError(t, err, "Failed to start dockerscaleset")
|
||||
|
||||
// Command exit error
|
||||
cmdCh := make(chan error, 1)
|
||||
t.Cleanup(func() {
|
||||
_ = runCmd.Process.Signal(os.Interrupt)
|
||||
<-cmdCh
|
||||
})
|
||||
|
||||
// Wait for log line
|
||||
waitCh := make(chan struct{}, 1)
|
||||
|
||||
var (
|
||||
bufMu sync.Mutex
|
||||
buf bytes.Buffer
|
||||
)
|
||||
go func() {
|
||||
scanner := bufio.NewScanner(stdout)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
bufMu.Lock()
|
||||
buf.WriteString(line + "\n")
|
||||
bufMu.Unlock()
|
||||
if strings.Contains(line, "Getting next message") {
|
||||
close(waitCh)
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
testErrCh <- fmt.Errorf("error reading dockerscaleset stdout: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
cmdCh <- runCmd.Wait()
|
||||
close(cmdCh)
|
||||
}()
|
||||
|
||||
runID, err := workflowEnv.triggerWorkflowDispatch(t, t.Context())
|
||||
require.NoError(t, err, "Failed to trigger workflow")
|
||||
|
||||
statusCh := make(chan *WorkflowRun, 1)
|
||||
go func() {
|
||||
select {
|
||||
case <-waitCh:
|
||||
case <-time.After(30 * time.Second):
|
||||
bufMu.Lock()
|
||||
logs := buf.String()
|
||||
bufMu.Unlock()
|
||||
testErrCh <- fmt.Errorf("timeout waiting for dockerscaleset to be ready; logs:\n%s", logs)
|
||||
return
|
||||
}
|
||||
status, err := workflowEnv.waitForWorkflowCompletion(t, t.Context(), runID, 10*time.Minute)
|
||||
if err != nil {
|
||||
testErrCh <- fmt.Errorf("failed to wait for workflow completion: %w", err)
|
||||
return
|
||||
}
|
||||
statusCh <- status
|
||||
}()
|
||||
|
||||
select {
|
||||
case err := <-cmdCh:
|
||||
select {
|
||||
case status := <-statusCh:
|
||||
assert.Equal(t, "completed", status.Status)
|
||||
assert.Equal(t, "success", status.Conclusion)
|
||||
case <-time.After(30 * time.Second):
|
||||
bufMu.Lock()
|
||||
logs := buf.String()
|
||||
bufMu.Unlock()
|
||||
t.Fatalf("Timeout waiting for workflow status after dockerscaleset exited\nexit: %v\nlogs:%s\n", err, logs)
|
||||
}
|
||||
case status := <-statusCh:
|
||||
assert.NotNil(t, status, "WorkflowRun status is nil")
|
||||
assert.Equal(t, "completed", status.Status)
|
||||
assert.Equal(t, "success", status.Conclusion)
|
||||
return
|
||||
case err := <-testErrCh:
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
type e2eWorkflowEnv struct {
|
||||
targetOrg string
|
||||
targetRepo string
|
||||
targetFile string
|
||||
|
||||
scalesetName string
|
||||
client *github.Client
|
||||
}
|
||||
|
||||
func mustE2EWorkflowEnv(t *testing.T, scalesetName string) *e2eWorkflowEnv {
|
||||
return &e2eWorkflowEnv{
|
||||
targetOrg: mustGetEnv(t, "E2E_WORKFLOW_TARGET_ORG"),
|
||||
targetRepo: mustGetEnv(t, "E2E_WORKFLOW_TARGET_REPO"),
|
||||
targetFile: mustGetEnv(t, "E2E_WORKFLOW_TARGET_FILE"),
|
||||
scalesetName: scalesetName,
|
||||
client: github.NewClient(nil).WithAuthToken(mustGetEnv(t, "E2E_WORKFLOW_GITHUB_TOKEN")),
|
||||
}
|
||||
}
|
||||
|
||||
func mustE2ECommandArgs(t *testing.T, configURL, name string) []string {
|
||||
args := []string{
|
||||
"--url", configURL,
|
||||
"--name", name,
|
||||
"--log-level", "debug",
|
||||
}
|
||||
|
||||
// GitHub App credentials
|
||||
var (
|
||||
clientID string
|
||||
installationID int
|
||||
privateKeyPath string
|
||||
)
|
||||
|
||||
// GitHub token
|
||||
var token string
|
||||
|
||||
clientID = os.Getenv("E2E_SCALESET_GITHUB_APP_CLIENT_ID")
|
||||
installationIDStr := os.Getenv("E2E_SCALESET_GITHUB_APP_INSTALLATION_ID")
|
||||
privateKeyPath = os.Getenv("E2E_SCALESET_GITHUB_APP_PRIVATE_KEY_PATH")
|
||||
|
||||
if clientID != "" && installationIDStr != "" && privateKeyPath != "" {
|
||||
id, err := strconv.Atoi(installationIDStr)
|
||||
require.NoError(t, err, "Invalid E2E_SCALESET_GITHUB_APP_INSTALLATION_ID")
|
||||
installationID = id
|
||||
args = append(args,
|
||||
"--app-client-id", clientID,
|
||||
"--app-installation-id", fmt.Sprintf("%d", installationID),
|
||||
"--app-private-key", privateKeyPath,
|
||||
)
|
||||
} else {
|
||||
token = os.Getenv("E2E_SCALESET_GITHUB_TOKEN")
|
||||
require.NotEmpty(t, token, "E2E_SCALESET_GITHUB_TOKEN must be set if GitHub App credentials are not provided")
|
||||
args = append(args,
|
||||
"--token", token,
|
||||
)
|
||||
}
|
||||
|
||||
runnerGroup := os.Getenv("E2E_SCALESET_RUNNER_GROUP")
|
||||
if runnerGroup != "" {
|
||||
args = append(args,
|
||||
"--runner-group", runnerGroup,
|
||||
)
|
||||
}
|
||||
|
||||
minRunners := 0
|
||||
if minRunnersStr := os.Getenv("E2E_SCALESET_MIN_RUNNERS"); minRunnersStr != "" {
|
||||
m, err := strconv.Atoi(minRunnersStr)
|
||||
require.NoError(t, err, "Invalid E2E_SCALESET_MIN_RUNNERS")
|
||||
minRunners = m
|
||||
require.GreaterOrEqual(t, minRunners, 0, "E2E_SCALESET_MIN_RUNNERS must be >= 0")
|
||||
}
|
||||
|
||||
maxRunners := 10
|
||||
if maxRunnersStr := os.Getenv("E2E_SCALESET_MAX_RUNNERS"); maxRunnersStr != "" {
|
||||
m, err := strconv.Atoi(maxRunnersStr)
|
||||
require.NoError(t, err, "Invalid E2E_SCALESET_MAX_RUNNERS")
|
||||
maxRunners = m
|
||||
require.GreaterOrEqual(t, maxRunners, 0, "E2E_SCALESET_MAX_RUNNERS must be >= 0")
|
||||
}
|
||||
|
||||
require.GreaterOrEqual(t, maxRunners, minRunners, "E2E_SCALESET_MAX_RUNNERS must be >= E2E_SCALESET_MIN_RUNNERS")
|
||||
|
||||
args = append(args,
|
||||
"--min-runners", strconv.Itoa(minRunners),
|
||||
"--max-runners", strconv.Itoa(maxRunners),
|
||||
)
|
||||
|
||||
return args
|
||||
}
|
||||
|
||||
type WorkflowRun struct {
|
||||
ID int `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Conclusion string `json:"conclusion"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
}
|
||||
|
||||
func (env *e2eWorkflowEnv) triggerWorkflowDispatch(t *testing.T, ctx context.Context) (int, error) {
|
||||
dispatchTime := time.Now().UTC()
|
||||
|
||||
resp, err := env.client.Actions.CreateWorkflowDispatchEventByFileName(
|
||||
ctx,
|
||||
env.targetOrg,
|
||||
env.targetRepo,
|
||||
env.targetFile,
|
||||
github.CreateWorkflowDispatchEventRequest{
|
||||
Ref: "main",
|
||||
Inputs: map[string]any{
|
||||
"scaleset_name": env.scalesetName,
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err, "Failed to create workflow dispatch")
|
||||
require.Equal(t, 204, resp.StatusCode, "Unexpected status code from workflow dispatch")
|
||||
|
||||
// Wait a bit for the run to be created
|
||||
time.Sleep(10 * time.Second)
|
||||
|
||||
// List runs with event=workflow_dispatch and since=dispatchTime
|
||||
opts := &github.ListWorkflowRunsOptions{
|
||||
Event: "workflow_dispatch",
|
||||
Created: ">=" + dispatchTime.Format(time.RFC3339),
|
||||
ListOptions: github.ListOptions{
|
||||
PerPage: 10,
|
||||
},
|
||||
}
|
||||
runs, _, err := env.client.Actions.ListWorkflowRunsByFileName(
|
||||
t.Context(),
|
||||
env.targetOrg,
|
||||
env.targetRepo,
|
||||
env.targetFile,
|
||||
opts,
|
||||
)
|
||||
require.NoError(t, err, "Failed to list workflow runs")
|
||||
require.Greater(t, len(runs.WorkflowRuns), 0, "No workflow runs found after dispatch")
|
||||
|
||||
// Sort by created_at desc, take the first (most recent)
|
||||
var latestRun *github.WorkflowRun
|
||||
var latestTime time.Time
|
||||
for _, run := range runs.WorkflowRuns {
|
||||
createdAt := run.CreatedAt.Time
|
||||
if createdAt.After(latestTime) {
|
||||
latestTime = createdAt
|
||||
latestRun = run
|
||||
}
|
||||
}
|
||||
|
||||
if latestRun == nil {
|
||||
return 0, fmt.Errorf("no workflow runs found after dispatch")
|
||||
}
|
||||
|
||||
return int(latestRun.GetID()), nil
|
||||
}
|
||||
|
||||
func (env *e2eWorkflowEnv) waitForWorkflowCompletion(t *testing.T, ctx context.Context, runID int, timeout time.Duration) (*WorkflowRun, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-ticker.C:
|
||||
run, _, err := env.client.Actions.GetWorkflowRunByID(ctx, env.targetOrg, env.targetRepo, int64(runID))
|
||||
require.NoError(t, err, "Failed to get workflow run by ID")
|
||||
|
||||
if run.GetStatus() == "completed" {
|
||||
return &WorkflowRun{
|
||||
ID: int(run.GetID()),
|
||||
Status: run.GetStatus(),
|
||||
Conclusion: run.GetConclusion(),
|
||||
CreatedAt: run.GetCreatedAt().Format(time.RFC3339),
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mustGetEnv(t *testing.T, key string) string {
|
||||
value := os.Getenv(key)
|
||||
if value == "" {
|
||||
t.Fatalf("Environment variable %s not set", key)
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/signal"
|
||||
|
||||
"github.com/actions/scaleset"
|
||||
"github.com/actions/scaleset/listener"
|
||||
"github.com/docker/docker/api/types/image"
|
||||
dockerclient "github.com/docker/docker/client"
|
||||
"github.com/google/uuid"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
func init() {
|
||||
flags := cmd.Flags()
|
||||
flags.StringVar(&cfg.RegistrationURL, "url", "", "REQUIRED: URL where to register your scale set (e.g. https://github.com/org/repo)")
|
||||
flags.IntVar(&cfg.MaxRunners, "max-runners", 10, "Maximum number of runners")
|
||||
flags.IntVar(&cfg.MinRunners, "min-runners", 0, "Minimum number of runners")
|
||||
flags.StringVar(&cfg.ScaleSetName, "name", "", "REQUIRED: Name of your scale set")
|
||||
flags.StringSliceVar(&cfg.Labels, "labels", nil, "Labels for workflow targeting (comma-separated or repeated). Defaults to --name if not provided.")
|
||||
flags.StringVar(&cfg.RunnerGroup, "runner-group", scaleset.DefaultRunnerGroup, "Name of the runner group your scale set should belong to")
|
||||
flags.StringVar(&cfg.GitHubApp.ClientID, "app-client-id", "", "GitHub App client id")
|
||||
flags.Int64Var(&cfg.GitHubApp.InstallationID, "app-installation-id", 0, "GitHub App installation ID")
|
||||
flags.StringVar(&cfg.GitHubApp.PrivateKey, "app-private-key", "", "GitHub App private key")
|
||||
flags.StringVar(&cfg.Token, "token", "", "Personal access token (can be used in place of a GitHub App, although not recommended)")
|
||||
flags.StringVar(&cfg.LogLevel, "log-level", "info", "Logging level (debug, info, warn, error)")
|
||||
flags.StringVar(&cfg.LogFormat, "log-format", "text", "Logging format (text, json). If invalid value is provided, defaults to no logs.")
|
||||
|
||||
if err := cmd.MarkFlagRequired("url"); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
if err := cmd.MarkFlagRequired("name"); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
func main() {
|
||||
if err := cmd.Execute(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func run(ctx context.Context, c Config) error {
|
||||
// Ensure that the config is valid
|
||||
if err := c.Validate(); err != nil {
|
||||
return fmt.Errorf("configuration validation failed: %w", err)
|
||||
}
|
||||
|
||||
logger := c.Logger()
|
||||
|
||||
// Create a new scaleset scalesetClient
|
||||
scalesetClient, err := c.ScalesetClient()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create scaleset client: %w", err)
|
||||
}
|
||||
|
||||
// Get the runner group ID of the chosen runner group
|
||||
var runnerGroupID int
|
||||
switch c.RunnerGroup {
|
||||
case scaleset.DefaultRunnerGroup:
|
||||
runnerGroupID = 1
|
||||
default:
|
||||
runnerGroup, err := scalesetClient.GetRunnerGroupByName(ctx, c.RunnerGroup)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get runner group ID: %w", err)
|
||||
}
|
||||
runnerGroupID = runnerGroup.ID
|
||||
}
|
||||
|
||||
// Create the runner scale set
|
||||
scaleSet, err := scalesetClient.CreateRunnerScaleSet(ctx, &scaleset.RunnerScaleSet{
|
||||
Name: c.ScaleSetName,
|
||||
RunnerGroupID: runnerGroupID,
|
||||
Labels: c.BuildLabels(),
|
||||
RunnerSetting: scaleset.RunnerSetting{
|
||||
DisableUpdate: true,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create runner scale set: %w", err)
|
||||
}
|
||||
|
||||
// Set the user agent for the scaleset client now that we have the scale set ID
|
||||
scalesetClient.SetSystemInfo(systemInfo(scaleSet.ID))
|
||||
|
||||
defer func() {
|
||||
logger.Info(
|
||||
"Deleting runner scale set",
|
||||
slog.Int("scaleSetID", scaleSet.ID),
|
||||
)
|
||||
if err := scalesetClient.DeleteRunnerScaleSet(context.WithoutCancel(ctx), scaleSet.ID); err != nil {
|
||||
slog.Error(
|
||||
"Failed to delete runner scale set",
|
||||
slog.Int("scaleSetID", scaleSet.ID),
|
||||
slog.String("error", err.Error()),
|
||||
)
|
||||
}
|
||||
}()
|
||||
|
||||
dockerClient, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv, dockerclient.WithAPIVersionNegotiation())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create docker client: %w", err)
|
||||
}
|
||||
|
||||
logger.Info(
|
||||
"Pulling runner image",
|
||||
slog.String("image", c.RunnerImage),
|
||||
)
|
||||
// Pull the runner image
|
||||
pull, err := dockerClient.ImagePull(ctx, c.RunnerImage, image.PullOptions{})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to pull runner image: %w", err)
|
||||
}
|
||||
|
||||
if _, err := io.ReadAll(pull); err != nil {
|
||||
return fmt.Errorf("failed to read image pull response: %w", err)
|
||||
}
|
||||
|
||||
if err := pull.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close image pull: %w", err)
|
||||
}
|
||||
|
||||
// Get the name of the client which will be used as the owner
|
||||
hostname, err := os.Hostname()
|
||||
if err != nil {
|
||||
hostname = uuid.NewString()
|
||||
logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
|
||||
}
|
||||
|
||||
sessionClient, err := scalesetClient.MessageSessionClient(ctx, scaleSet.ID, hostname)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create message session client: %w", err)
|
||||
}
|
||||
defer sessionClient.Close(context.Background())
|
||||
|
||||
logger.Info("Initializing listener")
|
||||
listener, err := listener.New(sessionClient, listener.Config{
|
||||
ScaleSetID: scaleSet.ID,
|
||||
MaxRunners: c.MaxRunners,
|
||||
Logger: logger.WithGroup("listener"),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create listener: %w", err)
|
||||
}
|
||||
|
||||
scaler := &Scaler{
|
||||
logger: logger.WithGroup("scaler"),
|
||||
runners: runnerState{
|
||||
idle: make(map[string]string),
|
||||
busy: make(map[string]string),
|
||||
},
|
||||
runnerImage: c.RunnerImage,
|
||||
minRunners: c.MinRunners,
|
||||
maxRunners: c.MaxRunners,
|
||||
dockerClient: dockerClient,
|
||||
scalesetClient: scalesetClient,
|
||||
scaleSetID: scaleSet.ID,
|
||||
}
|
||||
|
||||
defer scaler.shutdown(context.WithoutCancel(ctx))
|
||||
|
||||
logger.Info("Starting listener")
|
||||
if err := listener.Run(ctx, scaler); !errors.Is(err, context.Canceled) {
|
||||
return fmt.Errorf("listener run failed: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var cfg Config
|
||||
|
||||
var cmd = &cobra.Command{
|
||||
Use: "dockerscaleset",
|
||||
Short: "Example CLI application scaling runners using Docker",
|
||||
Long: `This is an example CLI application that demonstrates how to scale
|
||||
runners using Docker.`,
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
ctx, cancel := signal.NotifyContext(cmd.Context(), os.Interrupt)
|
||||
defer cancel()
|
||||
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return fmt.Errorf("invalid configuration: %w", err)
|
||||
}
|
||||
|
||||
return run(ctx, cfg)
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
"github.com/actions/scaleset"
|
||||
"github.com/actions/scaleset/listener"
|
||||
"github.com/docker/docker/api/types/container"
|
||||
dockerclient "github.com/docker/docker/client"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type Scaler struct {
|
||||
runners runnerState
|
||||
runnerImage string
|
||||
scaleSetID int
|
||||
dockerClient *dockerclient.Client
|
||||
scalesetClient *scaleset.Client
|
||||
minRunners int
|
||||
maxRunners int
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func (a *Scaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) {
|
||||
currentCount := a.runners.count()
|
||||
targetRunnerCount := min(a.maxRunners, a.minRunners+count)
|
||||
|
||||
switch {
|
||||
case targetRunnerCount == currentCount:
|
||||
// No scaling needed
|
||||
return currentCount, nil
|
||||
case targetRunnerCount > currentCount:
|
||||
// Scale up
|
||||
scaleUp := targetRunnerCount - currentCount
|
||||
a.logger.Info(
|
||||
"Scaling up runners",
|
||||
slog.Int("currentCount", currentCount),
|
||||
slog.Int("desiredCount", targetRunnerCount),
|
||||
slog.Int("scaleUp", scaleUp),
|
||||
)
|
||||
|
||||
for range scaleUp {
|
||||
if _, err := a.startRunner(ctx); err != nil {
|
||||
return 0, fmt.Errorf("failed to start runner: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return a.runners.count(), nil
|
||||
default:
|
||||
// No need to handle scale down events, since:
|
||||
// 1. JobCompleted events will first remove runners
|
||||
// 2. If the count is still below the current runner count, the JobCompleted event will be delivered in the next batch.
|
||||
// 3. Removal after JobCompleted events is handled synchronously.
|
||||
// 4. If the job is cancelled, the JobCompleted event will still be delivered.
|
||||
}
|
||||
return a.runners.count(), nil
|
||||
}
|
||||
|
||||
func (a *Scaler) HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error {
|
||||
a.logger.Info(
|
||||
"Job started",
|
||||
slog.Int64("runnerRequestId", jobInfo.RunnerRequestID),
|
||||
slog.String("jobId", jobInfo.JobID),
|
||||
)
|
||||
a.runners.markBusy(jobInfo.RunnerName)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Scaler) HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error {
|
||||
a.logger.Info("Job completed", slog.Int64("runnerRequestId", jobInfo.RunnerRequestID), slog.String("jobId", jobInfo.JobID))
|
||||
|
||||
containerID := a.runners.markDone(jobInfo.RunnerName)
|
||||
if err := a.dockerClient.ContainerRemove(ctx, containerID, container.RemoveOptions{Force: true}); err != nil {
|
||||
return fmt.Errorf("failed to remove runner container: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Scaler) startRunner(ctx context.Context) (string, error) {
|
||||
name := fmt.Sprintf("runner-%s", uuid.NewString()[:8])
|
||||
|
||||
jit, err := a.scalesetClient.GenerateJitRunnerConfig(
|
||||
ctx,
|
||||
&scaleset.RunnerScaleSetJitRunnerSetting{
|
||||
Name: name,
|
||||
},
|
||||
a.scaleSetID,
|
||||
)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to generate JIT config: %w", err)
|
||||
}
|
||||
|
||||
c, err := a.dockerClient.ContainerCreate(
|
||||
ctx,
|
||||
&container.Config{
|
||||
Image: a.runnerImage,
|
||||
User: "runner",
|
||||
Cmd: []string{"/home/runner/run.sh"},
|
||||
Env: []string{
|
||||
fmt.Sprintf("ACTIONS_RUNNER_INPUT_JITCONFIG=%s", jit.EncodedJITConfig),
|
||||
},
|
||||
},
|
||||
nil,
|
||||
nil, nil,
|
||||
name,
|
||||
)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create runner container: %w", err)
|
||||
}
|
||||
|
||||
if err := a.dockerClient.ContainerStart(ctx, c.ID, container.StartOptions{}); err != nil {
|
||||
return "", fmt.Errorf("failed to start runner container: %w", err)
|
||||
}
|
||||
|
||||
a.runners.addIdle(name, c.ID)
|
||||
return name, nil
|
||||
}
|
||||
|
||||
func (a *Scaler) shutdown(ctx context.Context) {
|
||||
a.logger.Info("Shutting down runners")
|
||||
a.runners.mu.Lock()
|
||||
defer a.runners.mu.Unlock()
|
||||
|
||||
for name, containerID := range a.runners.idle {
|
||||
a.logger.Info("Removing idle runner", slog.String("name", name), slog.String("containerID", containerID))
|
||||
if err := a.dockerClient.ContainerRemove(ctx, containerID, container.RemoveOptions{Force: true}); err != nil {
|
||||
a.logger.Error("Failed to remove idle runner container", slog.String("name", name), slog.String("containerID", containerID), slog.String("error", err.Error()))
|
||||
}
|
||||
}
|
||||
clear(a.runners.idle)
|
||||
|
||||
for name, containerID := range a.runners.busy {
|
||||
a.logger.Info("Removing busy runner", slog.String("name", name), slog.String("containerID", containerID))
|
||||
if err := a.dockerClient.ContainerRemove(ctx, containerID, container.RemoveOptions{Force: true}); err != nil {
|
||||
a.logger.Error("Failed to remove busy runner container", slog.String("name", name), slog.String("containerID", containerID), slog.String("error", err.Error()))
|
||||
}
|
||||
}
|
||||
clear(a.runners.busy)
|
||||
}
|
||||
|
||||
var _ listener.Scaler = (*Scaler)(nil)
|
||||
|
||||
type runnerState struct {
|
||||
mu sync.Mutex
|
||||
idle map[string]string
|
||||
busy map[string]string
|
||||
}
|
||||
|
||||
func (r *runnerState) count() int {
|
||||
r.mu.Lock()
|
||||
count := len(r.idle) + len(r.busy)
|
||||
r.mu.Unlock()
|
||||
return count
|
||||
}
|
||||
|
||||
func (r *runnerState) markBusy(name string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
state, ok := r.idle[name]
|
||||
if !ok {
|
||||
panic("marking non-existent runner busy")
|
||||
}
|
||||
delete(r.idle, name)
|
||||
r.busy[name] = state
|
||||
}
|
||||
|
||||
func (r *runnerState) markDone(name string) string {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.markDoneUnlocked(name)
|
||||
}
|
||||
|
||||
func (r *runnerState) markDoneUnlocked(name string) string {
|
||||
containerID, ok := r.busy[name]
|
||||
if ok {
|
||||
delete(r.busy, name)
|
||||
return containerID
|
||||
}
|
||||
containerID, ok = r.idle[name]
|
||||
if ok {
|
||||
delete(r.idle, name)
|
||||
return containerID
|
||||
}
|
||||
panic("marking non-existent runner done")
|
||||
}
|
||||
|
||||
func (r *runnerState) addIdle(name, containerID string) {
|
||||
r.mu.Lock()
|
||||
r.idle[name] = containerID
|
||||
r.mu.Unlock()
|
||||
}
|
||||
Reference in New Issue
Block a user