From 51aa04a4b2305002e42cd5958f67e93f7811cca4 Mon Sep 17 00:00:00 2001 From: Nikola Jokic Date: Fri, 31 Oct 2025 12:35:41 +0100 Subject: [PATCH] Add Docker example (#5) * Renames, style changes and unesports * use signal handling * remove unused fields, unexport packages, have a working example * remove viper * fix logger * fix test * remove unused var * Added few log lines * max runners fixed * Restructure example * Improve naming for app auth * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Remove old comment --------- Co-authored-by: Francesco Renzi Co-authored-by: Francesco Renzi Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- client.go | 31 ++- client_runner_test.go | 12 +- examples/dockerscaleset/config.go | 103 ++++++++ examples/dockerscaleset/go.mod | 47 ++++ examples/dockerscaleset/go.sum | 123 +++++++++ examples/dockerscaleset/main.go | 181 +++++++++++++ examples/dockerscaleset/scaler.go | 199 +++++++++++++++ go.mod | 2 - go.sum | 4 - listener/listener.go | 405 ++++++++++++++++++++++++++++++ types.go | 17 +- 11 files changed, 1100 insertions(+), 24 deletions(-) create mode 100644 examples/dockerscaleset/config.go create mode 100644 examples/dockerscaleset/go.mod create mode 100644 examples/dockerscaleset/go.sum create mode 100644 examples/dockerscaleset/main.go create mode 100644 examples/dockerscaleset/scaler.go create mode 100644 listener/listener.go diff --git a/client.go b/client.go index 19d5cd4..1cfad2a 100644 --- a/client.go +++ b/client.go @@ -77,12 +77,25 @@ type Client struct { } type GitHubAppAuth struct { - // AppID is the ID or the Client ID of the application - AppID string - // AppInstallationID is the installation ID of the GitHub App - AppInstallationID int64 - // AppPrivateKey is the private key of the GitHub App in PEM format - AppPrivateKey string + // ClientID is the Client ID of the application (app id also works) + ClientID string + // InstallationID is the installation ID of the GitHub App + InstallationID int64 + // PrivateKey is the private key of the GitHub App in PEM format + PrivateKey string +} + +func (a *GitHubAppAuth) Validate() error { + if a.ClientID == "" { + return fmt.Errorf("client ID is required") + } + if a.InstallationID == 0 { + return fmt.Errorf("app installation ID is required") + } + if a.PrivateKey == "" { + return fmt.Errorf("app private key is required") + } + return nil } type ActionsAuth struct { @@ -913,7 +926,7 @@ func (c *Client) fetchAccessToken(ctx context.Context, creds *GitHubAppAuth) (*a return nil, fmt.Errorf("failed to create JWT for GitHub app: %w", err) } - path := fmt.Sprintf("/app/installations/%v/access_tokens", creds.AppInstallationID) + path := fmt.Sprintf("/app/installations/%v/access_tokens", creds.InstallationID) req, err := c.newGitHubAPIRequest(ctx, http.MethodPost, path, nil) if err != nil { return nil, fmt.Errorf("failed to create new GitHub API request: %w", err) @@ -1098,12 +1111,12 @@ func createJWTForGitHubApp(appAuth *GitHubAppAuth) (string, error) { claims := &jwt.RegisteredClaims{ IssuedAt: jwt.NewNumericDate(issuedAt), ExpiresAt: jwt.NewNumericDate(expiresAt), - Issuer: appAuth.AppID, + Issuer: appAuth.ClientID, } token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims) - privateKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(appAuth.AppPrivateKey)) + privateKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(appAuth.PrivateKey)) if err != nil { return "", fmt.Errorf("failed to parse RSA private key from PEM: %w", err) } diff --git a/client_runner_test.go b/client_runner_test.go index 120aca4..42dcef2 100644 --- a/client_runner_test.go +++ b/client_runner_test.go @@ -66,7 +66,7 @@ func TestGetRunnerByName(t *testing.T) { t.Run("Get Runner by Name", func(t *testing.T) { var runnerID int64 = 1 - var runnerName = "self-hosted-ubuntu" + runnerName := "self-hosted-ubuntu" want := &RunnerReference{ ID: int(runnerID), Name: runnerName, @@ -86,7 +86,7 @@ func TestGetRunnerByName(t *testing.T) { }) t.Run("Get Runner by name with not exist runner", func(t *testing.T) { - var runnerName = "self-hosted-ubuntu" + runnerName := "self-hosted-ubuntu" response := []byte(`{"count": 0, "value": []}`) server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -102,7 +102,7 @@ func TestGetRunnerByName(t *testing.T) { }) t.Run("Default retries on server error", func(t *testing.T) { - var runnerName = "self-hosted-ubuntu" + runnerName := "self-hosted-ubuntu" retryWaitMax := 1 * time.Millisecond retryMax := 1 @@ -179,8 +179,8 @@ func TestGetRunnerGroupByName(t *testing.T) { } t.Run("Get RunnerGroup by Name", func(t *testing.T) { - var runnerGroupID int64 = 1 - var runnerGroupName = "test-runner-group" + runnerGroupID := 1 + runnerGroupName := "test-runner-group" want := &RunnerGroup{ ID: runnerGroupID, Name: runnerGroupName, @@ -200,7 +200,7 @@ func TestGetRunnerGroupByName(t *testing.T) { }) t.Run("Get RunnerGroup by name with not exist runner group", func(t *testing.T) { - var runnerGroupName = "test-runner-group" + runnerGroupName := "test-runner-group" response := []byte(`{"count": 0, "value": []}`) server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/examples/dockerscaleset/config.go b/examples/dockerscaleset/config.go new file mode 100644 index 0000000..8412bc6 --- /dev/null +++ b/examples/dockerscaleset/config.go @@ -0,0 +1,103 @@ +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 + 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") + } + 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 +} + +func (c *Config) ActionsAuth() *scaleset.ActionsAuth { + if err := c.GitHubApp.Validate(); err == nil { + return &scaleset.ActionsAuth{ + App: &c.GitHubApp, + } + } + + return &scaleset.ActionsAuth{ + Token: c.Token, + } +} + +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) + } +} diff --git a/examples/dockerscaleset/go.mod b/examples/dockerscaleset/go.mod new file mode 100644 index 0000000..bdffcec --- /dev/null +++ b/examples/dockerscaleset/go.mod @@ -0,0 +1,47 @@ +module github.com/actions/scaleset/examples/docker + +go 1.25.3 + +require ( + github.com/actions/scaleset v0.0.0-20240605120000-abcdef123456 + github.com/docker/docker v28.5.1+incompatible + github.com/google/uuid v1.6.0 + github.com/moby/moby v28.5.1+incompatible + github.com/spf13/cobra v1.10.1 +) + +require ( + github.com/Microsoft/go-winio v0.6.2 // indirect + github.com/containerd/errdefs v1.0.0 // indirect + github.com/containerd/errdefs/pkg v0.3.0 // indirect + github.com/containerd/log v0.1.0 // indirect + github.com/distribution/reference v0.6.0 // indirect + github.com/docker/go-connections v0.6.0 // indirect + github.com/docker/go-units v0.5.0 // indirect + github.com/felixge/httpsnoop v1.0.4 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/golang-jwt/jwt/v4 v4.5.2 // indirect + github.com/hashicorp/go-cleanhttp v0.5.2 // indirect + github.com/hashicorp/go-retryablehttp v0.7.8 // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/moby/docker-image-spec v1.3.1 // indirect + github.com/moby/sys/atomicwriter v0.1.0 // indirect + github.com/moby/term v0.5.2 // indirect + github.com/morikuni/aec v1.0.0 // indirect + github.com/opencontainers/go-digest v1.0.0 // indirect + github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/spf13/pflag v1.0.10 // indirect + go.opentelemetry.io/auto/sdk v1.1.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect + go.opentelemetry.io/otel v1.38.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 // indirect + go.opentelemetry.io/otel/metric v1.38.0 // indirect + go.opentelemetry.io/otel/trace v1.38.0 // indirect + golang.org/x/sys v0.35.0 // indirect + golang.org/x/time v0.14.0 // indirect + gotest.tools/v3 v3.5.2 // indirect +) + +replace github.com/actions/scaleset => ../.. diff --git a/examples/dockerscaleset/go.sum b/examples/dockerscaleset/go.sum new file mode 100644 index 0000000..e377099 --- /dev/null +++ b/examples/dockerscaleset/go.sum @@ -0,0 +1,123 @@ +github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEKWjV8V+WSxDXJ4NFATAsZjh8iIbsQIg= +github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= +github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= +github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= +github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= +github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= +github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= +github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= +github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= +github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk= +github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I= +github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= +github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= +github.com/docker/docker v28.5.1+incompatible h1:Bm8DchhSD2J6PsFzxC35TZo4TLGR2PdW/E69rU45NhM= +github.com/docker/docker v28.5.1+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= +github.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94= +github.com/docker/go-connections v0.6.0/go.mod h1:AahvXYshr6JgfUJGdDCs2b5EZG/vmaMAntpSFH5BFKE= +github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= +github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= +github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM= +github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE= +github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= +github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI= +github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnVTyacbefKhmbLhIhU= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs= +github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ= +github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48= +github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB11/k= +github.com/hashicorp/go-hclog v1.6.3/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M= +github.com/hashicorp/go-retryablehttp v0.7.8 h1:ylXZWnqa7Lhqpk0L1P1LzDtGcCR0rPVUrx/c8Unxc48= +github.com/hashicorp/go-retryablehttp v0.7.8/go.mod h1:rjiScheydd+CxvumBsIrFKlx3iS0jrZ7LvzFGFmuKbw= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= +github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= +github.com/moby/moby v28.5.1+incompatible h1:JD8lBdCDBF2oiHWLqIRofPqI8qvkppRjMJ6EnwrhvX0= +github.com/moby/moby v28.5.1+incompatible/go.mod h1:fDXVQ6+S340veQPv35CzDahGBmHsiclFwfEygB/TWMc= +github.com/moby/sys/atomicwriter v0.1.0 h1:kw5D/EqkBwsBFi0ss9v1VG3wIkVhzGvLklJ+w3A14Sw= +github.com/moby/sys/atomicwriter v0.1.0/go.mod h1:Ul8oqv2ZMNHOceF643P6FKPXeCmYtlQMvpizfsSoaWs= +github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU= +github.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko= +github.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ= +github.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc= +github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= +github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= +github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= +github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= +github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= +github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/spf13/cobra v1.10.1 h1:lJeBwCfmrnXthfAupyUTzJ/J4Nc1RsHC/mSRU2dll/s= +github.com/spf13/cobra v1.10.1/go.mod h1:7SmJGaTHFVBY0jW4NXGluQoLvhqFQM+6XSKD+P4XaB0= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= +github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg= +go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8= +go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 h1:GqRJVj7UmLjCVyVJ3ZFLdPRmhDUp2zFmQe3RHIOsw24= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0/go.mod h1:ri3aaHSmCTVYu2AWv44YMauwAQc0aqI9gHKIcSbI1pU= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 h1:aTL7F04bJHUlztTsNGJ2l+6he8c+y/b//eR0jjjemT4= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0/go.mod h1:kldtb7jDTeol0l3ewcmd8SDvx3EmIE7lyvqbasU3QC4= +go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA= +go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI= +go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E= +go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg= +go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM= +go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA= +go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE= +go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs= +go.opentelemetry.io/proto/otlp v1.7.1 h1:gTOMpGDb0WTBOP8JaO72iL3auEZhVmAQg4ipjOVAtj4= +go.opentelemetry.io/proto/otlp v1.7.1/go.mod h1:b2rVh6rfI/s2pHWNlB7ILJcRALpcNDzKhACevjI+ZnE= +golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4= +golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210= +golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= +golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k= +golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM= +golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI= +golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4= +google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5 h1:BIRfGDEjiHRrk0QKZe3Xv2ieMhtgRGeLcZQ0mIVn4EY= +google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5/go.mod h1:j3QtIyytwqGr1JUDtYXwtMXWPKsEa5LtzIFN1Wn5WvE= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5 h1:eaY8u2EuxbRv7c3NiGK0/NedzVsCcV6hDuU5qPX5EGE= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5/go.mod h1:M4/wBTSeyLxupu3W3tJtOgB14jILAS/XWPSSa3TAlJc= +google.golang.org/grpc v1.75.0 h1:+TW+dqTd2Biwe6KKfhE5JpiYIBWq865PhKGSXiivqt4= +google.golang.org/grpc v1.75.0/go.mod h1:JtPAzKiq4v1xcAB2hydNlWI2RnF85XXcV0mhKXr2ecQ= +google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc= +google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= +gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA= diff --git a/examples/dockerscaleset/main.go b/examples/dockerscaleset/main.go new file mode 100644 index 0000000..30550f9 --- /dev/null +++ b/examples/dockerscaleset/main.go @@ -0,0 +1,181 @@ +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/moby/moby/client" + "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.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 := scaleset.NewClient(c.RegistrationURL, c.ActionsAuth()) + 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: []scaleset.Label{ + { + Name: c.ScaleSetName, + Type: "System", + }, + }, + RunnerSetting: scaleset.RunnerSetting{ + Ephemeral: true, + DisableUpdate: true, + }, + }) + if err != nil { + return fmt.Errorf("failed to create runner scale set: %w", err) + } + + 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) + } + + logger.Info("Initializing listener") + listener, err := listener.New(scalesetClient, listener.Config{ + ScaleSetID: scaleSet.ID, + MinRunners: c.MinRunners, + 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, + 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) + }, +} diff --git a/examples/dockerscaleset/scaler.go b/examples/dockerscaleset/scaler.go new file mode 100644 index 0000000..c062f06 --- /dev/null +++ b/examples/dockerscaleset/scaler.go @@ -0,0 +1,199 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + "sync" + + "github.com/actions/scaleset" + "github.com/actions/scaleset/listener" + "github.com/docker/docker/api/types/container" + "github.com/google/uuid" + dockerclient "github.com/moby/moby/client" +) + +type Scaler struct { + runners runnerState + runnerImage string + scaleSetID int + dockerClient *dockerclient.Client + scalesetClient *scaleset.Client + minRunners int + logger *slog.Logger +} + +func (a *Scaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) { + currentCount := a.runners.count() + targetRunnerCount := min(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), + ) + + var errs []error + for range scaleUp { + if _, err := a.startRunner(ctx); err != nil { + errs = append(errs, err) + } + } + + for _, err := range errs { + // TODO: should we return error? + a.logger.Error("Failed to start runner", slog.String("error", err.Error())) + } + 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.Handler = (*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() +} diff --git a/go.mod b/go.mod index a00d8e7..48077e0 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module github.com/actions/scaleset go 1.25.3 require ( - github.com/go-logr/logr v1.4.3 github.com/golang-jwt/jwt/v4 v4.5.2 github.com/google/uuid v1.6.0 github.com/hashicorp/go-retryablehttp v0.7.8 @@ -17,7 +16,6 @@ require ( github.com/kr/pretty v0.3.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rogpeppe/go-internal v1.13.1 // indirect - github.com/stretchr/objx v0.5.2 // indirect golang.org/x/text v0.30.0 // indirect gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index a6f910e..b9754b3 100644 --- a/go.sum +++ b/go.sum @@ -3,8 +3,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM= github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE= -github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= -github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI= github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -29,8 +27,6 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= -github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= -github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4= diff --git a/listener/listener.go b/listener/listener.go new file mode 100644 index 0000000..a47cd42 --- /dev/null +++ b/listener/listener.go @@ -0,0 +1,405 @@ +// Package listener provides a listener for GitHub Actions runner scale set messages. +package listener + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "os" + "time" + + "github.com/actions/scaleset" + "github.com/google/uuid" +) + +const ( + sessionCreationMaxRetries = 10 +) + +type Config struct { + ScaleSetID int + MinRunners int + MaxRunners int + Logger *slog.Logger +} + +func (c *Config) defaults() { + if c.Logger == nil { + c.Logger = slog.New(slog.DiscardHandler) + } +} + +func (c *Config) Validate() error { + c.defaults() + + if c.ScaleSetID == 0 { + return errors.New("scaleSetID is required") + } + if c.MinRunners < 0 { + return errors.New("minRunners must be greater than or equal to 0") + } + if c.MaxRunners < 0 { + return errors.New("maxRunners must be greater than or equal to 0") + } + if c.MaxRunners > 0 && c.MinRunners > c.MaxRunners { + return errors.New("minRunners must be less than or equal to maxRunners") + } + return nil +} + +type Listener struct { + // The main client responsible for communicating with the scaleset service + client *scaleset.Client + + // Configuration for the listener + scaleSetID int + minRunners int + maxRunners int + + // lastMessageID keeps track of the last processed message ID + lastMessageID int64 + // hostname of the current machine + hostname string + // session represents the current message session + session *scaleset.RunnerScaleSetSession + + // configuration for the listener + logger *slog.Logger +} + +func New(client *scaleset.Client, config Config) (*Listener, error) { + if client == nil { + return nil, errors.New("client is required") + } + + if err := config.Validate(); err != nil { + return nil, err + } + + hostname, err := os.Hostname() + if err != nil { + hostname = uuid.NewString() + config.Logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err) + } + + return &Listener{ + client: client, + scaleSetID: config.ScaleSetID, + minRunners: config.MinRunners, + maxRunners: config.MaxRunners, + hostname: hostname, + logger: config.Logger, + }, nil +} + +type Handler interface { + HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error + HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error + HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) +} + +func (l *Listener) Run(ctx context.Context, handler Handler) error { + l.logger.Info("Creating message session") + if err := l.createSession(ctx); err != nil { + return fmt.Errorf("failed to create session: %w", err) + } + + defer func() { + l.logger.Debug("Deleting message session") + if err := l.deleteMessageSession(); err != nil { + l.logger.Error("failed to delete message session", "error", err.Error()) + } + }() + + if l.session.Statistics == nil { + return fmt.Errorf("session statistics is nil") + } + + l.logger.Info("Message session created; listening for messages", "sessionID", l.session.SessionID) + + // Handle initial statistics + if _, err := handler.HandleDesiredRunnerCount(ctx, l.session.Statistics.TotalAssignedJobs); err != nil { + return fmt.Errorf("handling initial message failed: %w", err) + } + + for { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + msg, err := l.getMessage(ctx) + if err != nil { + return fmt.Errorf("failed to get message: %w", err) + } + + if msg == nil { + _, err := handler.HandleDesiredRunnerCount(ctx, 0) + if err != nil { + return fmt.Errorf("handling nil message failed: %w", err) + } + + continue + } + + // Remove cancellation from the context to avoid cancelling the message handling. + if err := l.handleMessage(context.WithoutCancel(ctx), handler, msg); err != nil { + return fmt.Errorf("failed to handle message: %w", err) + } + } +} + +func (l *Listener) handleMessage(ctx context.Context, handler Handler, msg *scaleset.RunnerScaleSetMessage) error { + parsedMsg, err := l.parseMessage(ctx, msg) + if err != nil { + return fmt.Errorf("failed to parse message: %w", err) + } + + l.lastMessageID = msg.MessageID + + if err := l.deleteLastMessage(ctx); err != nil { + return fmt.Errorf("failed to delete message: %w", err) + } + + for _, jobStarted := range parsedMsg.jobsStarted { + if err := handler.HandleJobStarted(ctx, jobStarted); err != nil { + return fmt.Errorf("failed to handle job started: %w", err) + } + } + for _, jobCompleted := range parsedMsg.jobsCompleted { + if err := handler.HandleJobCompleted(ctx, jobCompleted); err != nil { + return fmt.Errorf("failed to handle job completed: %w", err) + } + } + + if _, err := handler.HandleDesiredRunnerCount(ctx, parsedMsg.statistics.TotalAssignedJobs); err != nil { + return fmt.Errorf("failed to handle desired runner count: %w", err) + } + + return nil +} + +func (l *Listener) createSession(ctx context.Context) error { + var session *scaleset.RunnerScaleSetSession + var retries int + + for { + var err error + session, err = l.client.CreateMessageSession(ctx, l.scaleSetID, l.hostname) + if err == nil { + break + } + + clientErr := &scaleset.HttpClientSideError{} + if !errors.As(err, &clientErr) { + return fmt.Errorf("failed to create session: %w", err) + } + + if clientErr.Code != http.StatusConflict { + return fmt.Errorf("failed to create session: %w", err) + } + + retries++ + if retries >= sessionCreationMaxRetries { + return fmt.Errorf("failed to create session after %d retries: %w", retries, err) + } + + l.logger.Info("Unable to create message session. Will try again in 30 seconds", "error", err.Error()) + + select { + case <-ctx.Done(): + return fmt.Errorf("context cancelled: %w", ctx.Err()) + case <-time.After(30 * time.Second): + } + } + + statistics, err := json.Marshal(session.Statistics) + if err != nil { + return fmt.Errorf("failed to marshal statistics: %w", err) + } + l.logger.Info("Current runner scale set statistics.", "statistics", string(statistics)) + + l.session = session + + return nil +} + +func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMessage, error) { + l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID) + msg, err := l.client.GetMessage( + ctx, + l.session.MessageQueueURL, + l.session.MessageQueueAccessToken, + l.lastMessageID, + l.maxRunners, + ) + if err == nil { // if NO error + return msg, nil + } + + expiredError := &scaleset.MessageQueueTokenExpiredError{} + if !errors.As(err, &expiredError) { + return nil, fmt.Errorf("failed to get next message: %w", err) + } + + if err := l.refreshSession(ctx); err != nil { + return nil, err + } + + l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID) + + msg, err = l.client.GetMessage( + ctx, + l.session.MessageQueueURL, + l.session.MessageQueueAccessToken, + l.lastMessageID, + l.maxRunners, + ) + if err != nil { // if error + return nil, fmt.Errorf("failed to get next message after message session refresh: %w", err) + } + + return msg, nil +} + +func (l *Listener) deleteLastMessage(ctx context.Context) error { + l.logger.Info("Deleting last message", "lastMessageID", l.lastMessageID) + err := l.client.DeleteMessage( + ctx, + l.session.MessageQueueURL, + l.session.MessageQueueAccessToken, + l.lastMessageID, + ) + if err == nil { // if NO error + return nil + } + + expiredError := &scaleset.MessageQueueTokenExpiredError{} + if !errors.As(err, &expiredError) { + return fmt.Errorf("failed to delete last message: %w", err) + } + + if err := l.refreshSession(ctx); err != nil { + return err + } + + err = l.client.DeleteMessage( + ctx, + l.session.MessageQueueURL, + l.session.MessageQueueAccessToken, + l.lastMessageID, + ) + if err != nil { + return fmt.Errorf("failed to delete last message after message session refresh: %w", err) + } + + return nil +} + +type parsedMessage struct { + statistics *scaleset.RunnerScaleSetStatistic + jobsStarted []*scaleset.JobStarted + jobsCompleted []*scaleset.JobCompleted +} + +func (l *Listener) parseMessage(ctx context.Context, msg *scaleset.RunnerScaleSetMessage) (*parsedMessage, error) { + if msg.MessageType != "RunnerScaleSetJobMessages" { + l.logger.Info("Skipping message", "messageType", msg.MessageType) + return nil, fmt.Errorf("invalid message type: %s", msg.MessageType) + } + + l.logger.Info("Processing message", "messageId", msg.MessageID, "messageType", msg.MessageType) + if msg.Statistics == nil { + return nil, fmt.Errorf("invalid message: statistics is nil") + } + + l.logger.Info("New runner scale set statistics.", "statistics", msg.Statistics) + + var batchedMessages []json.RawMessage + if len(msg.Body) > 0 { + if err := json.Unmarshal([]byte(msg.Body), &batchedMessages); err != nil { + return nil, fmt.Errorf("failed to unmarshal batched messages: %w", err) + } + } + + parsedMsg := &parsedMessage{ + statistics: msg.Statistics, + } + + for _, msg := range batchedMessages { + var messageType scaleset.JobMessageType + if err := json.Unmarshal(msg, &messageType); err != nil { + return nil, fmt.Errorf("failed to decode job message type: %w", err) + } + + switch messageType.MessageType { + case scaleset.MessageTypeJobAssigned: + var jobAssigned scaleset.JobAssigned + if err := json.Unmarshal(msg, &jobAssigned); err != nil { + return nil, fmt.Errorf("failed to decode job assigned: %w", err) + } + + l.logger.Info("Job assigned message received", "jobId", jobAssigned.JobID) + + case scaleset.MessageTypeJobStarted: + var jobStarted scaleset.JobStarted + if err := json.Unmarshal(msg, &jobStarted); err != nil { + return nil, fmt.Errorf("could not decode job started message. %w", err) + } + l.logger.Info("Job started message received.", "JobID", jobStarted.JobID, "RunnerId", jobStarted.RunnerID) + parsedMsg.jobsStarted = append(parsedMsg.jobsStarted, &jobStarted) + + case scaleset.MessageTypeJobCompleted: + var jobCompleted scaleset.JobCompleted + if err := json.Unmarshal(msg, &jobCompleted); err != nil { + return nil, fmt.Errorf("failed to decode job completed: %w", err) + } + + l.logger.Info( + "Job completed message received.", + "JobID", jobCompleted.JobID, + "Result", jobCompleted.Result, + "RunnerId", jobCompleted.RunnerID, + "RunnerName", jobCompleted.RunnerName, + ) + parsedMsg.jobsCompleted = append(parsedMsg.jobsCompleted, &jobCompleted) + + default: + l.logger.Info("unknown job message type.", "messageType", messageType.MessageType) + } + } + + return parsedMsg, ctx.Err() +} + +func (l *Listener) refreshSession(ctx context.Context) error { + l.logger.Info("Message queue token is expired during GetNextMessage, refreshing...") + session, err := l.client.RefreshMessageSession( + ctx, + l.session.RunnerScaleSet.ID, + l.session.SessionID, + ) + if err != nil { + return fmt.Errorf("refresh message session failed. %w", err) + } + + l.session = session + return nil +} + +func (l *Listener) deleteMessageSession() error { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + l.logger.Info("Deleting message session") + + if err := l.client.DeleteMessageSession(ctx, l.session.RunnerScaleSet.ID, l.session.SessionID); err != nil { + return fmt.Errorf("failed to delete message session: %w", err) + } + + return nil +} diff --git a/types.go b/types.go index 5cbc126..462564e 100644 --- a/types.go +++ b/types.go @@ -6,6 +6,17 @@ import ( "github.com/google/uuid" ) +const DefaultRunnerGroup = "default" + +type MessageType string + +// message types +const ( + MessageTypeJobAssigned MessageType = "JobAssigned" + MessageTypeJobStarted MessageType = "JobStarted" + MessageTypeJobCompleted MessageType = "JobCompleted" +) + type Int64List struct { Count int `json:"count"` Value []int64 `json:"value"` @@ -34,7 +45,7 @@ type JobCompleted struct { } type JobMessageType struct { - MessageType string `json:"messageType"` + MessageType MessageType `json:"messageType"` } type JobMessageBase struct { @@ -60,7 +71,7 @@ type Label struct { } type RunnerGroup struct { - ID int64 `json:"id"` + ID int `json:"id"` Name string `json:"name"` Size int64 `json:"size"` IsDefault bool `json:"isDefaultGroup"` @@ -101,7 +112,7 @@ type runnerScaleSetsResponse struct { } type RunnerScaleSetSession struct { - SessionID *uuid.UUID `json:"sessionId,omitempty"` + SessionID uuid.UUID `json:"sessionId,omitempty"` OwnerName string `json:"ownerName,omitempty"` RunnerScaleSet *RunnerScaleSet `json:"runnerScaleSet,omitempty"` MessageQueueURL string `json:"messageQueueUrl,omitempty"`