diff --git a/workers/scaleset/autoscale_test.go b/workers/scaleset/autoscale_test.go new file mode 100644 index 000000000..ad07ebce2 --- /dev/null +++ b/workers/scaleset/autoscale_test.go @@ -0,0 +1,140 @@ +//go:build testing + +package scaleset + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "sync/atomic" + "testing" + "time" + + "github.com/google/go-github/v84/github" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + commonParams "github.com/cloudbase/garm-provider-common/params" + "github.com/cloudbase/garm/cache" + storeMocks "github.com/cloudbase/garm/database/common/mocks" + "github.com/cloudbase/garm/params" + runnerMocks "github.com/cloudbase/garm/runner/common/mocks" +) + +func TestAutoscaleRequiresFreshRunnerState(t *testing.T) { + require.NoError(t, registerTestLocker()) + for _, direction := range []string{"up", "down"} { + t.Run(direction, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + entityID := "refresh-" + direction + entity := params.ForgeEntity{ID: entityID, EntityType: params.ForgeEntityTypeRepository, Owner: "owner", Name: "repo"} + cache.SetGithubToolsCache(entity, nil) + defer cache.DeleteGithubToolsCache(entityID) + + var jitRequests, removals, localCreates, localUpdates atomic.Int32 + encodedJIT := base64.StdEncoding.EncodeToString([]byte(`{"key":"value"}`)) + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case runnerRegistrationPath: + _, _ = fmt.Fprintf(w, `{"url":%q,"token":"eyJhbGciOiJub25lIn0.eyJleHAiOjQxNDk5MzYwMDB9."}`, server.URL) + case "/_apis/runtime/runnerscalesets/42/generatejitconfig": + assert.Equal(t, http.MethodPost, r.Method) + jitRequests.Add(1) + _, _ = fmt.Fprintf(w, `{"runner":{"id":99},"encodedJITConfig":%q}`, encodedJIT) + case "/_apis/distributedtask/pools/0/agents/1": + assert.Equal(t, http.MethodDelete, r.Method) + removals.Add(1) + w.WriteHeader(http.StatusNoContent) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + } + })) + defer server.Close() + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + gh := runnerMocks.NewGithubClient(t) + gh.EXPECT().GetEntity().Return(entity).Maybe() + gh.EXPECT().GithubBaseURL().Return(baseURL).Maybe() + gh.EXPECT().CreateEntityRegistrationToken(mock.Anything).Return(&github.RegistrationToken{ + Token: github.Ptr("registration-token"), ExpiresAt: &github.Timestamp{Time: time.Now().Add(time.Hour)}, + }, nil, nil).Maybe() + cache.SetGithubClient(entityID, gh) + defer cache.DeleteGithubClient(entityID) + + failedRefresh := make(chan struct{}) + freshRefresh := make(chan struct{}) + idle := params.Instance{ID: "idle", Name: "idle", AgentID: 1, Status: commonParams.InstanceRunning, RunnerStatus: params.RunnerIdle} + fresh := []params.Instance{{ID: "deleted", Name: "deleted", Status: commonParams.InstanceDeleted}} + if direction == "down" { + fresh = append(fresh, idle) + } + store := storeMocks.NewStore(t) + store.EXPECT().ListScaleSetInstances(mock.Anything, uint(4), false). + Run(func(context.Context, uint, bool) { close(failedRefresh) }).Return(nil, errors.New("database unavailable")).Once() + store.EXPECT().ListScaleSetInstances(mock.Anything, uint(4), false). + Run(func(context.Context, uint, bool) { close(freshRefresh) }).Return(fresh, nil).Once() + // Cleanup failure must not prevent scaling from the successful refresh. + store.EXPECT().DeleteInstanceByName(mock.Anything, "deleted").Return(errors.New("row cleanup failed")).Once() + store.EXPECT().ControllerInfo().Return(params.ControllerInfo{}, nil).Maybe() + store.EXPECT().CreateScaleSetInstance(mock.Anything, uint(4), mock.Anything). + RunAndReturn(func(_ context.Context, _ uint, p params.CreateInstanceParams) (params.Instance, error) { + localCreates.Add(1) + return params.Instance{ID: "new", Name: p.Name, Status: commonParams.InstancePendingCreate}, nil + }).Maybe() + store.EXPECT().UpdateInstance(mock.Anything, "idle", params.UpdateInstanceParams{Status: commonParams.InstancePendingDelete}). + RunAndReturn(func(_ context.Context, _ string, _ params.UpdateInstanceParams) (params.Instance, error) { + localUpdates.Add(1) + return params.Instance{ID: "idle", Name: "idle", Status: commonParams.InstancePendingDelete}, nil + }).Maybe() + w := &Worker{ + ctx: ctx, store: store, consumerID: entityID, quit: make(chan struct{}), + scaleSet: params.ScaleSet{ID: 4, RepoID: entityID, ScaleSetID: 42, Enabled: true, MaxRunners: 1}, + runners: map[string]params.Instance{}, + } + if direction == "up" { + w.scaleSet.MinIdleRunners = 1 + } else { + w.runners[idle.ID] = idle + } + done := make(chan struct{}) + go func() { + defer close(done) + w.handleAutoScale() + }() + defer func() { cancel(); <-done }() + for tick, signal := range []chan struct{}{failedRefresh, freshRefresh} { + select { + case <-signal: + case <-time.After(15 * time.Second): + t.Fatal("autoscale did not reconcile") + } + // The list call runs under mux; acquiring it waits for the whole tick. + w.mux.Lock() + if tick == 0 { + assert.Zero(t, jitRequests.Load(), "created a JIT runner from stale state") + assert.Zero(t, removals.Load(), "removed a runner from stale state") + assert.Zero(t, localCreates.Load()) + assert.Zero(t, localUpdates.Load()) + } + w.mux.Unlock() + } + if direction == "up" { + assert.EqualValues(t, 1, jitRequests.Load()) + assert.EqualValues(t, 1, localCreates.Load()) + assert.Zero(t, removals.Load()) + } else { + assert.EqualValues(t, 1, removals.Load()) + assert.EqualValues(t, 1, localUpdates.Load()) + assert.Zero(t, jitRequests.Load()) + } + }) + } +} diff --git a/workers/scaleset/scaleset.go b/workers/scaleset/scaleset.go index fd66b53a1..7c0cd118a 100644 --- a/workers/scaleset/scaleset.go +++ b/workers/scaleset/scaleset.go @@ -174,6 +174,30 @@ func (w *Worker) ensureScaleSetInGitHub() error { return nil } +func (w *Worker) markScaleSetCreated() error { + if w.scaleSet.State == params.ScaleSetCreated { + return nil + } + + entity, err := w.scaleSet.GetEntity() + if err != nil { + return fmt.Errorf("getting scale set entity: %w", err) + } + state := params.ScaleSetCreated + updated, err := w.store.UpdateEntityScaleSet( + w.ctx, + entity, + w.scaleSet.ID, + params.UpdateScaleSetParams{State: &state}, + nil, + ) + if err != nil { + return fmt.Errorf("updating scale set state: %w", err) + } + w.scaleSet = updated + return nil +} + func (w *Worker) Stop() error { slog.DebugContext(w.ctx, "stopping scale set worker", "scale_set", w.consumerID) w.mux.Lock() @@ -329,6 +353,10 @@ func (w *Worker) Start() (err error) { return fmt.Errorf("failed to ensure scale set: %w", err) } + if err := w.markScaleSetCreated(); err != nil { + return fmt.Errorf("failed to mark scale set as created: %w", err) + } + consumer, err := watcher.RegisterConsumer( w.ctx, w.consumerID, watcher.WithAny( @@ -792,6 +820,27 @@ func (w *Worker) handleInstanceCleanup(instance params.Instance) error { return nil } +func (w *Worker) reconcileRunners() (refreshed bool, err error) { + instances, err := w.store.ListScaleSetInstances(w.ctx, w.scaleSet.ID, false) + if err != nil { + return false, fmt.Errorf("listing scale set instances: %w", err) + } + + runners := make(map[string]params.Instance, len(instances)) + for _, instance := range instances { + runners[instance.ID] = instance + } + w.runners = runners + + var cleanupErrs []error + for _, instance := range instances { + if err := w.handleInstanceCleanup(instance); err != nil { + cleanupErrs = append(cleanupErrs, err) + } + } + return true, errors.Join(cleanupErrs...) +} + func (w *Worker) handleInstanceEntityEvent(event dbCommon.ChangePayload) { instance, ok := event.Payload.(params.Instance) if !ok { @@ -970,7 +1019,8 @@ func (w *Worker) handleScaleUp() { return } - if w.targetRunners() <= w.runnerCount() { + runnersToAdd := w.runnersToAdd() + if runnersToAdd == 0 { slog.DebugContext(w.ctx, "target is less than or equal to current; not scaling up") return } @@ -986,7 +1036,7 @@ func (w *Worker) handleScaleUp() { slog.ErrorContext(w.ctx, "error getting scale set client", "error", err) return } - for i := w.runnerCount(); i < w.targetRunners(); i++ { + for range runnersToAdd { newRunnerName := strings.ToLower(fmt.Sprintf("%s-%s", w.scaleSet.GetRunnerPrefix(), util.NewID())) jitConfig, err := scaleSetCli.GenerateJitRunnerConfig(w.ctx, newRunnerName, w.scaleSet.ScaleSetID) if err != nil { @@ -1056,19 +1106,8 @@ func (w *Worker) handleScaleDown() { slog.ErrorContext(w.ctx, "error getting scale set client", "error", err) return } - // Runners already on their way out satisfy part of the delta before any - // new victim is picked. Counting them during the removal loop instead - // would make the outcome depend on map iteration order: an idle runner - // could be removed while a runner already being deleted would have - // covered the delta. + // runnerCount already excludes runners being deleted. removed := 0 - for _, runner := range w.runners { - switch runner.Status { - case commonParams.InstancePendingDelete, commonParams.InstancePendingForceDelete, - commonParams.InstanceDeleting: - removed++ - } - } for _, runner := range w.runners { slog.InfoContext(w.ctx, "considering runners for removal", "delta", delta, "removed", removed) @@ -1139,12 +1178,7 @@ func (w *Worker) handleScaleDown() { locking.Unlock(runner.Name, false) removed++ case commonParams.InstancePendingDelete, commonParams.InstancePendingForceDelete, - commonParams.InstanceDeleting: - // Already counted against the delta before the loop. - continue - case commonParams.InstanceDeleted: - // Provider resource is gone and the record is excluded from - // runnerCount(); it is not part of the delta. + commonParams.InstanceDeleting, commonParams.InstanceDeleted: continue default: slog.WarnContext(w.ctx, "runner is not in a valid state; skipping", "runner_name", runner.Name, "runner_status", runner.Status) @@ -1163,15 +1197,13 @@ func (w *Worker) targetRunners() int { return int(targetRunners) } -// runnerCount returns the number of runners that still occupy capacity in -// the provider. Runners on their way out (pending_delete, deleting) are -// still physically present in the IaaS, so they count; only runners whose -// provider resource is confirmed gone (deleted) are excluded — those records -// merely await database cleanup. +// runnerCount excludes runners that are being deleted from usable capacity. func (w *Worker) runnerCount() int { count := 0 for _, runner := range w.runners { - if runner.Status == commonParams.InstanceDeleted { + switch runner.Status { + case commonParams.InstancePendingDelete, commonParams.InstancePendingForceDelete, + commonParams.InstanceDeleting, commonParams.InstanceDeleted: continue } count++ @@ -1179,6 +1211,29 @@ func (w *Worker) runnerCount() int { return count } +func (w *Worker) providerSlotsInUse() uint { + var count uint + for _, runner := range w.runners { + if runner.Status != commonParams.InstanceDeleted { + count++ + } + } + return count +} + +func (w *Worker) runnersToAdd() int { + runnerDeficit := max(w.targetRunners()-w.runnerCount(), 0) + providerSlots := w.providerSlotsInUse() + if providerSlots >= w.scaleSet.MaxRunners { + return 0 + } + availableSlots := w.scaleSet.MaxRunners - providerSlots + if availableSlots < uint(runnerDeficit) { + return int(availableSlots) + } + return runnerDeficit +} + func (w *Worker) handleAutoScale() { ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() @@ -1210,10 +1265,13 @@ func (w *Worker) handleAutoScale() { return case <-ticker.C: w.mux.Lock() - for _, instance := range w.runners { - if err := w.handleInstanceCleanup(instance); err != nil { - slog.ErrorContext(w.ctx, "error cleaning up instance", "instance_id", instance.ID, "error", err) - } + refreshed, err := w.reconcileRunners() + if err != nil { + slog.ErrorContext(w.ctx, "error reconciling scale set instances", "error", err) + } + if !refreshed { + w.mux.Unlock() + continue } if w.runnerCount() == w.targetRunners() { diff --git a/workers/scaleset/scaleset_test.go b/workers/scaleset/scaleset_test.go new file mode 100644 index 000000000..ebe6fceff --- /dev/null +++ b/workers/scaleset/scaleset_test.go @@ -0,0 +1,391 @@ +// Copyright 2026 Cloudbase Solutions SRL +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may +// not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +// License for the specific language governing permissions and limitations +// under the License. + +//go:build testing + +package scaleset + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/go-github/v84/github" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + commonParams "github.com/cloudbase/garm-provider-common/params" + "github.com/cloudbase/garm/cache" + storeMocks "github.com/cloudbase/garm/database/common/mocks" + "github.com/cloudbase/garm/locking" + "github.com/cloudbase/garm/params" + runnerMocks "github.com/cloudbase/garm/runner/common/mocks" +) + +const runnerRegistrationPath = "/actions/runner-registration" + +var registerTestLocker = sync.OnceValue(func() error { + locker, err := locking.NewLocalLocker(context.Background(), nil) + if err != nil { + return err + } + return locking.RegisterLocker(locker) +}) + +func TestDeletingInstancesDoNotCountAsRunnersButConsumeProviderSlots(t *testing.T) { + for _, status := range []commonParams.InstanceStatus{ + commonParams.InstancePendingDelete, + commonParams.InstancePendingForceDelete, + commonParams.InstanceDeleting, + } { + t.Run(string(status), func(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{MinIdleRunners: 1, MaxRunners: 1}, + runners: map[string]params.Instance{ + "runner": {Status: status}, + }, + } + + assert.Zero(t, w.runnerCount()) + assert.Zero(t, w.runnersToAdd()) + + w.runners["runner"] = params.Instance{Status: commonParams.InstanceDeleted} + assert.Equal(t, 1, w.runnersToAdd()) + }) + } +} + +func TestDeletingInstancesDoNotIncreaseScaleDownDelta(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{DesiredRunnerCount: 1, MaxRunners: 1}, + runners: map[string]params.Instance{ + "running-1": {Status: commonParams.InstanceRunning}, + "running-2": {Status: commonParams.InstanceRunning}, + "deleting": {Status: commonParams.InstanceDeleting}, + }, + } + + assert.Equal(t, 1, w.runnerCount()-w.targetRunners()) +} + +func TestScaleDownRemovesExcessIdleRunnerDespitePendingDeletion(t *testing.T) { + const entityID = "scale-down-repo" + ctx := context.Background() + require.NoError(t, registerTestLocker()) + + var removals atomic.Int32 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case runnerRegistrationPath: + _, _ = fmt.Fprintf(w, `{"url":%q,"token":"eyJhbGciOiJub25lIn0.eyJleHAiOjQxNDk5MzYwMDB9."}`, server.URL) + case "/_apis/distributedtask/pools/0/agents/1", "/_apis/distributedtask/pools/0/agents/2": + assert.Equal(t, http.MethodDelete, r.Method) + removals.Add(1) + w.WriteHeader(http.StatusNoContent) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + entity := params.ForgeEntity{ID: entityID, EntityType: params.ForgeEntityTypeRepository, Owner: "owner", Name: "repo"} + cli := runnerMocks.NewGithubClient(t) + cli.EXPECT().GetEntity().Return(entity).Maybe() + cli.EXPECT().GithubBaseURL().Return(baseURL).Maybe() + cli.EXPECT().CreateEntityRegistrationToken(mock.Anything).Return(&github.RegistrationToken{ + Token: github.Ptr("registration-token"), ExpiresAt: &github.Timestamp{Time: time.Now().Add(time.Hour)}, + }, nil, nil).Maybe() + cache.SetGithubClient(entityID, cli) + t.Cleanup(func() { cache.DeleteGithubClient(entityID) }) + store := storeMocks.NewStore(t) + store.EXPECT().UpdateInstance(mock.Anything, mock.Anything, params.UpdateInstanceParams{Status: commonParams.InstancePendingDelete}). + RunAndReturn(func(_ context.Context, name string, _ params.UpdateInstanceParams) (params.Instance, error) { + return params.Instance{ID: name, Name: name, Status: commonParams.InstancePendingDelete}, nil + }).Once() + w := &Worker{ + ctx: ctx, store: store, consumerID: "scale-down-test", + scaleSet: params.ScaleSet{RepoID: entityID, DesiredRunnerCount: 1, MaxRunners: 3}, + runners: map[string]params.Instance{ + "idle-1": {ID: "idle-1", Name: "idle-1", AgentID: 1, Status: commonParams.InstanceRunning, RunnerStatus: params.RunnerIdle}, + "idle-2": {ID: "idle-2", Name: "idle-2", AgentID: 2, Status: commonParams.InstanceRunning, RunnerStatus: params.RunnerIdle}, + "deleting": {ID: "deleting", Name: "deleting", Status: commonParams.InstanceDeleting}, + }, + } + w.handleScaleDown() + assert.EqualValues(t, 1, removals.Load()) + assert.Equal(t, 1, w.runnerCount()) +} + +func TestPendingCreatePreventsDuplicateReplacement(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{MinIdleRunners: 1, MaxRunners: 1}, + runners: map[string]params.Instance{ + "deleted": {Status: commonParams.InstanceDeleted}, + "replacement": {Status: commonParams.InstancePendingCreate}, + }, + } + + assert.Equal(t, 1, w.runnerCount()) + assert.Equal(t, w.targetRunners(), w.runnerCount()) + assert.Zero(t, w.runnersToAdd()) +} + +func TestRunnersToAddHandlesLargeMaximum(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{MinIdleRunners: 1, MaxRunners: ^uint(0)}, + runners: make(map[string]params.Instance), + } + + assert.Equal(t, 1, w.runnersToAdd()) +} + +func TestRunnersToAddCapsDeficitAtRemainingProviderSlots(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{MinIdleRunners: 5, MaxRunners: 5}, + runners: map[string]params.Instance{ + "running": {Status: commonParams.InstanceRunning}, + "pending": {Status: commonParams.InstancePendingDelete}, + "pending-force": {Status: commonParams.InstancePendingForceDelete}, + "deleting": {Status: commonParams.InstanceDeleting}, + }, + } + assert.Equal(t, 1, w.runnerCount()) + assert.EqualValues(t, 4, w.providerSlotsInUse()) + assert.Equal(t, 4, w.targetRunners()-w.runnerCount()) + assert.Equal(t, 1, w.runnersToAdd()) +} + +func TestTargetRunnersHonorsMaximum(t *testing.T) { + w := &Worker{ + scaleSet: params.ScaleSet{ + MinIdleRunners: 2, + DesiredRunnerCount: 3, + MaxRunners: 4, + }, + } + + assert.Equal(t, 4, w.targetRunners()) +} + +func TestReconcileRunnersCleansDeletedRowsAndReplacesStaleCache(t *testing.T) { + store := storeMocks.NewStore(t) + store.EXPECT(). + ListScaleSetInstances(mock.Anything, uint(4), false). + Return([]params.Instance{ + {ID: "deleted", Name: "deleted", Status: commonParams.InstanceDeleted, ScaleSetID: 4}, + {ID: "replacement", Name: "replacement", Status: commonParams.InstancePendingCreate, ScaleSetID: 4}, + }, nil) + store.EXPECT(). + DeleteInstanceByName(mock.Anything, "deleted"). + Return(nil) + + w := &Worker{ + ctx: context.Background(), + store: store, + scaleSet: params.ScaleSet{ID: 4, MinIdleRunners: 1, MaxRunners: 1}, + runners: map[string]params.Instance{ + "stale": {ID: "stale", Name: "stale", Status: commonParams.InstanceRunning, ScaleSetID: 4}, + }, + } + + refreshed, err := w.reconcileRunners() + assert.NoError(t, err) + assert.True(t, refreshed) + assert.Equal(t, 1, w.runnerCount()) + assert.Len(t, w.runners, 1) + assert.Contains(t, w.runners, "replacement") +} + +func TestReconcileRunnersContinuesAfterCleanupFailure(t *testing.T) { + store := storeMocks.NewStore(t) + store.EXPECT(). + ListScaleSetInstances(mock.Anything, uint(4), false). + Return([]params.Instance{ + {ID: "failed", Name: "failed", Status: commonParams.InstanceDeleted, ScaleSetID: 4}, + {ID: "cleaned", Name: "cleaned", Status: commonParams.InstanceDeleted, ScaleSetID: 4}, + }, nil) + store.EXPECT(). + DeleteInstanceByName(mock.Anything, "failed"). + Return(errors.New("database error")) + store.EXPECT(). + DeleteInstanceByName(mock.Anything, "cleaned"). + Return(nil) + + w := &Worker{ + ctx: context.Background(), + store: store, + scaleSet: params.ScaleSet{ID: 4}, + } + + refreshed, err := w.reconcileRunners() + assert.True(t, refreshed) + assert.ErrorContains(t, err, "deleting instance failed") + assert.NotContains(t, w.runners, "cleaned") +} + +func TestReconcileDeletedRunnerCreatesOneReplacement(t *testing.T) { + const ( + entityID = "repo-id" + scaleSetID = 4 + ) + + var registrationRequests atomic.Int32 + var jitRequests atomic.Int32 + encodedJIT := base64.StdEncoding.EncodeToString([]byte(`{"key":"value"}`)) + + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case runnerRegistrationPath: + registrationRequests.Add(1) + assert.Equal(t, http.MethodPost, r.Method) + _, _ = fmt.Fprintf(w, `{"url":%q,"token":"eyJhbGciOiJub25lIn0.eyJleHAiOjQxNDk5MzYwMDB9."}`, server.URL) + case "/_apis/runtime/runnerscalesets/42/generatejitconfig": + jitRequests.Add(1) + assert.Equal(t, http.MethodPost, r.Method) + _, _ = fmt.Fprintf(w, `{"runner":{"id":99},"encodedJITConfig":%q}`, encodedJIT) + default: + http.NotFound(w, r) + t.Errorf("unexpected GitHub request: %s %s", r.Method, r.URL.Path) + } + })) + t.Cleanup(server.Close) + + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + entity := params.ForgeEntity{ + ID: entityID, + EntityType: params.ForgeEntityTypeRepository, + Owner: "owner", + Name: "repo", + } + githubClient := runnerMocks.NewGithubClient(t) + githubClient.EXPECT(). + CreateEntityRegistrationToken(mock.Anything). + Return(&github.RegistrationToken{ + Token: github.Ptr("registration-token"), + ExpiresAt: &github.Timestamp{Time: time.Now().Add(time.Hour)}, + }, nil, nil). + Once() + githubClient.EXPECT().GithubBaseURL().Return(baseURL).Once() + githubClient.EXPECT().GetEntity().Return(entity).Maybe() + cache.SetGithubClient(entityID, githubClient) + t.Cleanup(func() { cache.DeleteGithubClient(entityID) }) + + store := storeMocks.NewStore(t) + store.EXPECT(). + ListScaleSetInstances(mock.Anything, uint(scaleSetID), false). + Return([]params.Instance{{ + ID: "deleted", + Name: "deleted", + Status: commonParams.InstanceDeleted, + RunnerStatus: params.RunnerIdle, + ScaleSetID: scaleSetID, + }}, nil). + Once() + store.EXPECT().DeleteInstanceByName(mock.Anything, "deleted").Return(nil).Once() + store.EXPECT().ControllerInfo().Return(params.ControllerInfo{ + CallbackURL: "https://garm.example/callback", + MetadataURL: "https://garm.example/metadata", + }, nil).Once() + store.EXPECT(). + CreateScaleSetInstance( + mock.Anything, + uint(scaleSetID), + mock.MatchedBy(func(create params.CreateInstanceParams) bool { + return create.Status == commonParams.InstancePendingCreate && + create.RunnerStatus == params.RunnerPending && + create.CallbackURL == "https://garm.example/callback" && + create.MetadataURL == "https://garm.example/metadata" && + create.AgentID == 99 && + create.JitConfiguration["key"] == "value" + }), + ). + Return(params.Instance{ + ID: "replacement", + Name: "replacement", + Status: commonParams.InstancePendingCreate, + RunnerStatus: params.RunnerPending, + ScaleSetID: scaleSetID, + }, nil). + Once() + + w := &Worker{ + ctx: context.Background(), + store: store, + scaleSet: params.ScaleSet{ + ID: scaleSetID, + ScaleSetID: 42, + RepoID: entityID, + Enabled: true, + MinIdleRunners: 1, + MaxRunners: 1, + }, + runners: map[string]params.Instance{ + "stale": {ID: "stale", Status: commonParams.InstanceRunning}, + }, + } + + refreshed, err := w.reconcileRunners() + require.NoError(t, err) + require.True(t, refreshed) + w.handleScaleUp() + w.handleScaleUp() + + assert.EqualValues(t, 1, registrationRequests.Load()) + assert.EqualValues(t, 1, jitRequests.Load()) + assert.Len(t, w.runners, 1) + assert.Equal(t, commonParams.InstancePendingCreate, w.runners["replacement"].Status) +} + +func TestMarkScaleSetCreated(t *testing.T) { + store := storeMocks.NewStore(t) + entity := params.ForgeEntity{ID: "repo-id", EntityType: params.ForgeEntityTypeRepository} + store.EXPECT(). + UpdateEntityScaleSet( + mock.Anything, + entity, + uint(4), + mock.MatchedBy(func(update params.UpdateScaleSetParams) bool { + return update.State != nil && *update.State == params.ScaleSetCreated + }), + mock.Anything, + ). + Return(params.ScaleSet{ID: 4, RepoID: entity.ID, State: params.ScaleSetCreated}, nil) + + w := &Worker{ + ctx: context.Background(), + store: store, + scaleSet: params.ScaleSet{ + ID: 4, + RepoID: entity.ID, + State: params.ScaleSetPendingCreate, + }, + } + + assert.NoError(t, w.markScaleSetCreated()) + assert.Equal(t, params.ScaleSetCreated, w.scaleSet.State) +}