feat(cli/exp): extend scaletest create-workspaces with --retry option (#11825)

mafredri · web-flow · commit 02124758fbfa · 2024-01-26T11:29:48.000Z
Part of #11801
diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -524,6 +524,7 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd {
 func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 	var (
 		count    int64
+		retry    int64
 		template string
 
 		noCleanup bool
@@ -644,6 +645,7 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 							RichParameterValues: richParameters,
 						},
 						NoWaitForAgents: noWaitForAgents,
+						Retry:           int(retry),
 					},
 					NoCleanup: noCleanup,
 				}
@@ -753,6 +755,13 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 			Description:   "Required: Number of workspaces to create.",
 			Value:         clibase.Int64Of(&count),
 		},
+		{
+			Flag:        "retry",
+			Env:         "CODER_SCALETEST_RETRY",
+			Default:     "0",
+			Description: "Number of tries to create and bring up the workspace.",
+			Value:       clibase.Int64Of(&retry),
+		},
 		{
 			Flag:          "template",
 			FlagShorthand: "t",
diff --git a/scaletest/workspacebuild/config.go b/scaletest/workspacebuild/config.go
@@ -19,6 +19,9 @@ type Config struct {
 	// NoWaitForAgents determines whether the test should wait for the workspace
 	// agents to connect before returning.
 	NoWaitForAgents bool `json:"no_wait_for_agents"`
+	// Retry determines how many times to retry starting a workspace build if it
+	// fails.
+	Retry int `json:"retry"`
 }
 
 func (c Config) Validate() error {
diff --git a/scaletest/workspacebuild/run.go b/scaletest/workspacebuild/run.go
@@ -66,7 +66,25 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
 
 	err = waitForBuild(ctx, logs, r.client, workspace.LatestBuild.ID)
 	if err != nil {
-		return xerrors.Errorf("wait for build: %w", err)
+		for i := 0; i < r.cfg.Retry; i++ {
+			_, _ = fmt.Fprintf(logs, "Retrying build %d/%d...\n", i+1, r.cfg.Retry)
+
+			workspace.LatestBuild, err = r.client.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{
+				Transition:          codersdk.WorkspaceTransitionStart,
+				RichParameterValues: req.RichParameterValues,
+				TemplateVersionID:   req.TemplateVersionID,
+			})
+			if err != nil {
+				return xerrors.Errorf("create workspace build: %w", err)
+			}
+			err = waitForBuild(ctx, logs, r.client, workspace.LatestBuild.ID)
+			if err == nil {
+				break
+			}
+		}
+		if err != nil {
+			return xerrors.Errorf("wait for build: %w", err)
+		}
 	}
 
 	if r.cfg.NoWaitForAgents {
diff --git a/scaletest/workspacebuild/run_test.go b/scaletest/workspacebuild/run_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -230,4 +231,53 @@ func Test_Runner(t *testing.T) {
 		require.Error(t, err)
 		require.ErrorContains(t, err, "test error")
 	})
+
+	t.Run("RetryBuild", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+		client := coderdtest.New(t, &coderdtest.Options{
+			IncludeProvisionerDaemon: true,
+			Logger:                   &logger,
+		})
+		user := coderdtest.CreateFirstUser(t, client)
+
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:         echo.ParseComplete,
+			ProvisionPlan: echo.PlanComplete,
+			ProvisionApply: []*proto.Response{
+				{
+					Type: &proto.Response_Apply{
+						Apply: &proto.ApplyComplete{
+							Error: "test error",
+						},
+					},
+				},
+			},
+		})
+
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+
+		runner := workspacebuild.NewRunner(client, workspacebuild.Config{
+			OrganizationID: user.OrganizationID,
+			UserID:         codersdk.Me,
+			Request: codersdk.CreateWorkspaceRequest{
+				TemplateID: template.ID,
+			},
+			Retry: 1,
+		})
+
+		logs := bytes.NewBuffer(nil)
+		err := runner.Run(ctx, "1", logs)
+		logsStr := logs.String()
+		t.Log("Runner logs:\n\n" + logsStr)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "test error")
+		require.Equal(t, 1, strings.Count(logsStr, "Retrying build"))
+		require.Equal(t, 2*2, strings.Count(logsStr, "test error")) // once per error, once per logged sdk response.
+	})
 }

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,9 @@ type Config struct {`
`19`	`19`	`// NoWaitForAgents determines whether the test should wait for the workspace`
`20`	`20`	`// agents to connect before returning.`
`21`	`21`	NoWaitForAgents bool `json:"no_wait_for_agents"`
	`22`	`+ // Retry determines how many times to retry starting a workspace build if it`
	`23`	`+ // fails.`
	`24`	+ Retry int `json:"retry"`
`22`	`25`	`}`
`23`	`26`
`24`	`27`	`func (c Config) Validate() error {`