From 803818162c226742c6dae8d5c5db26ed3012087c Mon Sep 17 00:00:00 2001 From: Tingluo Huang Date: Fri, 27 Jan 2023 09:27:52 -0500 Subject: [PATCH] Allow update runner group for AutoScalingRunnerSet (#2216) --- .../autoscalingrunnerset_controller.go | 58 +++++++++++++++++-- .../autoscalingrunnerset_controller_test.go | 9 ++- github/actions/client.go | 42 ++++++++++++++ .../actions/client_runner_scale_set_test.go | 50 ++++++++++++++++ github/actions/fake/client.go | 21 +++++++ github/actions/mock_ActionsService.go | 23 ++++++++ 6 files changed, 196 insertions(+), 7 deletions(-) diff --git a/controllers/actions.github.com/autoscalingrunnerset_controller.go b/controllers/actions.github.com/autoscalingrunnerset_controller.go index a268d560..bae05cdf 100644 --- a/controllers/actions.github.com/autoscalingrunnerset_controller.go +++ b/controllers/actions.github.com/autoscalingrunnerset_controller.go @@ -21,6 +21,7 @@ import ( "fmt" "sort" "strconv" + "strings" "github.com/actions/actions-runner-controller/apis/actions.github.com/v1alpha1" "github.com/actions/actions-runner-controller/github/actions" @@ -47,6 +48,7 @@ const ( LabelKeyAutoScaleRunnerSetName = "auto-scale-runner-set-name" autoscalingRunnerSetFinalizerName = "autoscalingrunnerset.actions.github.com/finalizer" runnerScaleSetIdKey = "runner-scale-set-id" + runnerScaleSetRunnerGroupNameKey = "runner-scale-set-runner-group-name" // scaleSetListenerLabel is the key of pod.meta.labels to label // that the pod is a listener application @@ -150,6 +152,13 @@ func (r *AutoscalingRunnerSetReconciler) Reconcile(ctx context.Context, req ctrl return r.createRunnerScaleSet(ctx, autoscalingRunnerSet, log) } + // Make sure the runner group of the scale set is up to date + currentRunnerGroupName, ok := autoscalingRunnerSet.Annotations[runnerScaleSetRunnerGroupNameKey] + if !ok || !strings.EqualFold(currentRunnerGroupName, autoscalingRunnerSet.Spec.RunnerGroup) { + log.Info("AutoScalingRunnerSet runner group changed. Updating the runner scale set.") + return r.updateRunnerScaleSetRunnerGroup(ctx, autoscalingRunnerSet, log) + } + secret := new(corev1.Secret) if err := r.Get(ctx, types.NamespacedName{Namespace: autoscalingRunnerSet.Namespace, Name: autoscalingRunnerSet.Spec.GitHubConfigSecret}, secret); err != nil { log.Error(err, "Failed to find GitHub config secret.", @@ -299,8 +308,8 @@ func (r *AutoscalingRunnerSetReconciler) createRunnerScaleSet(ctx context.Contex return ctrl.Result{}, err } + runnerGroupId := 1 if runnerScaleSet == nil { - runnerGroupId := 1 if len(autoscalingRunnerSet.Spec.RunnerGroup) > 0 { runnerGroup, err := actionsClient.GetRunnerGroupByName(ctx, autoscalingRunnerSet.Spec.RunnerGroup) if err != nil { @@ -338,10 +347,12 @@ func (r *AutoscalingRunnerSetReconciler) createRunnerScaleSet(ctx context.Contex autoscalingRunnerSet.Annotations = map[string]string{} } - autoscalingRunnerSet.Annotations[runnerScaleSetIdKey] = strconv.Itoa(runnerScaleSet.Id) - logger.Info("Adding runner scale set ID as an annotation") - if err := r.Update(ctx, autoscalingRunnerSet); err != nil { - logger.Error(err, "Failed to add runner scale set ID") + logger.Info("Adding runner scale set ID and runner group name as an annotation") + if err = patch(ctx, r.Client, autoscalingRunnerSet, func(obj *v1alpha1.AutoscalingRunnerSet) { + obj.Annotations[runnerScaleSetIdKey] = strconv.Itoa(runnerScaleSet.Id) + obj.Annotations[runnerScaleSetRunnerGroupNameKey] = runnerScaleSet.RunnerGroupName + }); err != nil { + logger.Error(err, "Failed to add runner scale set ID and runner group name as an annotation") return ctrl.Result{}, err } @@ -349,6 +360,43 @@ func (r *AutoscalingRunnerSetReconciler) createRunnerScaleSet(ctx context.Contex return ctrl.Result{}, nil } +func (r *AutoscalingRunnerSetReconciler) updateRunnerScaleSetRunnerGroup(ctx context.Context, autoscalingRunnerSet *v1alpha1.AutoscalingRunnerSet, logger logr.Logger) (ctrl.Result, error) { + runnerScaleSetId, err := strconv.Atoi(autoscalingRunnerSet.Annotations[runnerScaleSetIdKey]) + if err != nil { + logger.Error(err, "Failed to parse runner scale set ID") + return ctrl.Result{}, err + } + + actionsClient, err := r.actionsClientFor(ctx, autoscalingRunnerSet) + if err != nil { + logger.Error(err, "Failed to initialize Actions service client for updating a existing runner scale set") + return ctrl.Result{}, err + } + + runnerGroup, err := actionsClient.GetRunnerGroupByName(ctx, autoscalingRunnerSet.Spec.RunnerGroup) + if err != nil { + logger.Error(err, "Failed to get runner group by name", "runnerGroup", autoscalingRunnerSet.Spec.RunnerGroup) + return ctrl.Result{}, err + } + + updatedRunnerScaleSet, err := actionsClient.UpdateRunnerScaleSet(ctx, runnerScaleSetId, &actions.RunnerScaleSet{Name: autoscalingRunnerSet.Name, RunnerGroupId: int(runnerGroup.ID)}) + if err != nil { + logger.Error(err, "Failed to update runner scale set", "runnerScaleSetId", runnerScaleSetId) + return ctrl.Result{}, err + } + + logger.Info("Updating runner scale set runner group name as an annotation") + if err := patch(ctx, r.Client, autoscalingRunnerSet, func(obj *v1alpha1.AutoscalingRunnerSet) { + obj.Annotations[runnerScaleSetRunnerGroupNameKey] = updatedRunnerScaleSet.RunnerGroupName + }); err != nil { + logger.Error(err, "Failed to update runner group name annotation") + return ctrl.Result{}, err + } + + logger.Info("Updated runner scale set with match runner group", "runnerGroup", updatedRunnerScaleSet.RunnerGroupName) + return ctrl.Result{}, nil +} + func (r *AutoscalingRunnerSetReconciler) createEphemeralRunnerSet(ctx context.Context, autoscalingRunnerSet *v1alpha1.AutoscalingRunnerSet, log logr.Logger) (ctrl.Result, error) { desiredRunnerSet, err := r.resourceBuilder.newEphemeralRunnerSet(autoscalingRunnerSet) if err != nil { diff --git a/controllers/actions.github.com/autoscalingrunnerset_controller_test.go b/controllers/actions.github.com/autoscalingrunnerset_controller_test.go index ca8def02..d281d8cc 100644 --- a/controllers/actions.github.com/autoscalingrunnerset_controller_test.go +++ b/controllers/actions.github.com/autoscalingrunnerset_controller_test.go @@ -83,6 +83,7 @@ var _ = Describe("Test AutoScalingRunnerSet controller", func() { GitHubConfigSecret: configSecret.Name, MaxRunners: &max, MinRunners: &min, + RunnerGroup: "testgroup", Template: corev1.PodTemplateSpec{ Spec: corev1.PodSpec{ Containers: []corev1.Container{ @@ -144,10 +145,14 @@ var _ = Describe("Test AutoScalingRunnerSet controller", func() { return "", nil } - return created.Annotations[runnerScaleSetIdKey], nil + if _, ok := created.Annotations[runnerScaleSetRunnerGroupNameKey]; !ok { + return "", nil + } + + return fmt.Sprintf("%s_%s", created.Annotations[runnerScaleSetIdKey], created.Annotations[runnerScaleSetRunnerGroupNameKey]), nil }, autoscalingRunnerSetTestTimeout, - autoscalingRunnerSetTestInterval).Should(BeEquivalentTo("1"), "RunnerScaleSet should be created/fetched and update the AutoScalingRunnerSet's annotation") + autoscalingRunnerSetTestInterval).Should(BeEquivalentTo("1_testgroup"), "RunnerScaleSet should be created/fetched and update the AutoScalingRunnerSet's annotation") // Check if ephemeral runner set is created Eventually( diff --git a/github/actions/client.go b/github/actions/client.go index 48ab4a21..fdadc70b 100644 --- a/github/actions/client.go +++ b/github/actions/client.go @@ -36,6 +36,7 @@ type ActionsService interface { GetRunnerScaleSetById(ctx context.Context, runnerScaleSetId int) (*RunnerScaleSet, error) GetRunnerGroupByName(ctx context.Context, runnerGroup string) (*RunnerGroup, error) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) + UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*RunnerScaleSetSession, error) DeleteMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) error @@ -350,6 +351,47 @@ func (c *Client) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *Runne return createdRunnerScaleSet, nil } +func (c *Client) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) { + u := fmt.Sprintf("%s/%s/%d?api-version=6.0-preview", *c.ActionsServiceURL, scaleSetEndpoint, runnerScaleSetId) + + if err := c.refreshTokenIfNeeded(ctx); err != nil { + return nil, fmt.Errorf("failed to refresh admin token if needed: %w", err) + } + + body, err := json.Marshal(runnerScaleSet) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPatch, u, bytes.NewBuffer(body)) + if err != nil { + return nil, err + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *c.ActionsServiceAdminToken)) + + if c.userAgent != "" { + req.Header.Set("User-Agent", c.userAgent) + } + + resp, err := c.Do(req) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, ParseActionsErrorFromResponse(resp) + } + + var updatedRunnerScaleSet *RunnerScaleSet + err = unmarshalBody(resp, &updatedRunnerScaleSet) + if err != nil { + return nil, err + } + return updatedRunnerScaleSet, nil +} + func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetId int) error { u := fmt.Sprintf("%s/%s/%d?api-version=6.0-preview", *c.ActionsServiceURL, scaleSetEndpoint, runnerScaleSetId) diff --git a/github/actions/client_runner_scale_set_test.go b/github/actions/client_runner_scale_set_test.go index e9d86e2d..980b9846 100644 --- a/github/actions/client_runner_scale_set_test.go +++ b/github/actions/client_runner_scale_set_test.go @@ -336,3 +336,53 @@ func TestCreateRunnerScaleSet(t *testing.T) { assert.Equalf(t, actualRetry, expectedRetry, "A retry was expected after the first request but got: %v", actualRetry) }) } + +func TestUpdateRunnerScaleSet(t *testing.T) { + ctx := context.Background() + auth := &actions.ActionsAuth{ + Token: "token", + } + + scaleSetCreationDateTime := time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC) + runnerScaleSet := actions.RunnerScaleSet{Id: 1, Name: "ScaleSet", RunnerGroupId: 1, RunnerGroupName: "group", CreatedOn: scaleSetCreationDateTime, RunnerSetting: actions.RunnerSetting{}} + + t.Run("Update runner scale set", func(t *testing.T) { + want := &runnerScaleSet + rsl, err := json.Marshal(want) + require.NoError(t, err) + server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Write(rsl) + })) + + client, err := actions.NewClient(ctx, server.configURLForOrg("my-org"), auth) + require.NoError(t, err) + + got, err := client.UpdateRunnerScaleSet(ctx, 1, &actions.RunnerScaleSet{RunnerGroupId: 1}) + require.NoError(t, err) + assert.Equal(t, want, got) + }) + + t.Run("UpdateRunnerScaleSet calls correct url", func(t *testing.T) { + rsl, err := json.Marshal(&runnerScaleSet) + require.NoError(t, err) + url := url.URL{} + method := "" + server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(rsl) + url = *r.URL + method = r.Method + })) + + client, err := actions.NewClient(ctx, server.configURLForOrg("my-org"), auth) + require.NoError(t, err) + + _, err = client.UpdateRunnerScaleSet(ctx, 1, &runnerScaleSet) + require.NoError(t, err) + + u := url.String() + expectedUrl := "/_apis/runtime/runnerscalesets/1?api-version=6.0-preview" + assert.Equal(t, expectedUrl, u) + + assert.Equal(t, "PATCH", method) + }) +} diff --git a/github/actions/fake/client.go b/github/actions/fake/client.go index fc2b75fb..1983180e 100644 --- a/github/actions/fake/client.go +++ b/github/actions/fake/client.go @@ -36,6 +36,18 @@ var defaultRunnerScaleSet = &actions.RunnerScaleSet{ Statistics: nil, } +var defaultUpdatedRunnerScaleSet = &actions.RunnerScaleSet{ + Id: 1, + Name: "testset", + RunnerGroupId: 2, + RunnerGroupName: "testgroup", + Labels: []actions.Label{{Type: "test", Name: "test"}}, + RunnerSetting: actions.RunnerSetting{}, + CreatedOn: time.Now(), + RunnerJitConfigUrl: "test.test.test", + Statistics: nil, +} + var defaultRunnerGroup = &actions.RunnerGroup{ ID: 1, Name: "testgroup", @@ -107,6 +119,10 @@ type FakeClient struct { *actions.RunnerScaleSet err error } + updateRunnerScaleSetResult struct { + *actions.RunnerScaleSet + err error + } createMessageSessionResult struct { *actions.RunnerScaleSetSession err error @@ -164,6 +180,7 @@ func (f *FakeClient) applyDefaults() { f.getRunnerScaleSetByIdResult.RunnerScaleSet = defaultRunnerScaleSet f.getRunnerGroupByNameResult.RunnerGroup = defaultRunnerGroup f.createRunnerScaleSetResult.RunnerScaleSet = defaultRunnerScaleSet + f.updateRunnerScaleSetResult.RunnerScaleSet = defaultUpdatedRunnerScaleSet f.createMessageSessionResult.RunnerScaleSetSession = defaultRunnerScaleSetSession f.refreshMessageSessionResult.RunnerScaleSetSession = defaultRunnerScaleSetSession f.acquireJobsResult.ids = []int64{1} @@ -190,6 +207,10 @@ func (f *FakeClient) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *a return f.createRunnerScaleSetResult.RunnerScaleSet, f.createRunnerScaleSetResult.err } +func (f *FakeClient) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *actions.RunnerScaleSet) (*actions.RunnerScaleSet, error) { + return f.updateRunnerScaleSetResult.RunnerScaleSet, f.updateRunnerScaleSetResult.err +} + func (f *FakeClient) CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*actions.RunnerScaleSetSession, error) { return f.createMessageSessionResult.RunnerScaleSetSession, f.createMessageSessionResult.err } diff --git a/github/actions/mock_ActionsService.go b/github/actions/mock_ActionsService.go index 341dc513..d227e8be 100644 --- a/github/actions/mock_ActionsService.go +++ b/github/actions/mock_ActionsService.go @@ -332,6 +332,29 @@ func (_m *MockActionsService) RemoveRunner(ctx context.Context, runnerId int64) return r0 } +// UpdateRunnerScaleSet provides a mock function with given fields: ctx, runnerScaleSetId, runnerScaleSet +func (_m *MockActionsService) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) { + ret := _m.Called(ctx, runnerScaleSetId, runnerScaleSet) + + var r0 *RunnerScaleSet + if rf, ok := ret.Get(0).(func(context.Context, int, *RunnerScaleSet) *RunnerScaleSet); ok { + r0 = rf(ctx, runnerScaleSetId, runnerScaleSet) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*RunnerScaleSet) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, int, *RunnerScaleSet) error); ok { + r1 = rf(ctx, runnerScaleSetId, runnerScaleSet) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + type mockConstructorTestingTNewMockActionsService interface { mock.TestingT Cleanup(func())