fix: add provision/tags to prebuilds scenario (#22294)

cli/exp_scaletest_prebuilds.go

@@ -29,6 +29,7 @@ func (r *RootCmd) scaletestPrebuilds() *serpent.Command {
 		templateVersionJobTimeout time.Duration
 		prebuildWorkspaceTimeout  time.Duration
 		noCleanup                 bool
+		provisionerTags           []string
 
 		tracingFlags    = &scaletestTracingFlags{}
 		timeoutStrategy = &timeoutFlags{}
@@ -111,10 +112,16 @@ func (r *RootCmd) scaletestPrebuilds() *serpent.Command {
 
 			th := harness.NewTestHarness(timeoutStrategy.wrapStrategy(harness.ConcurrentExecutionStrategy{}), cleanupStrategy.toStrategy())
 
+			tags, err := ParseProvisionerTags(provisionerTags)
+			if err != nil {
+				return err
+			}
+
 			for i := range numTemplates {
 				id := strconv.Itoa(int(i))
 				cfg := prebuilds.Config{
 					OrganizationID:            me.OrganizationIDs[0],
+					ProvisionerTags:           tags,
 					NumPresets:                int(numPresets),
 					NumPresetPrebuilds:        int(numPresetPrebuilds),
 					TemplateVersionJobTimeout: templateVersionJobTimeout,
@@ -283,6 +290,11 @@ func (r *RootCmd) scaletestPrebuilds() *serpent.Command {
 			Description: "Skip cleanup (deletion test) and leave resources intact.",
 			Value:       serpent.BoolOf(&noCleanup),
 		},
+		{
+			Flag:        "provisioner-tag",
+			Description: "Specify a set of tags to target provisioner daemons.",
+			Value:       serpent.StringArrayOf(&provisionerTags),
+		},
 	}
 
 	tracingFlags.attach(&cmd.Options)

scaletest/prebuilds/config.go

@@ -13,6 +13,9 @@ import (
 type Config struct {
 	// OrganizationID is the ID of the organization to create the prebuilds in.
 	OrganizationID uuid.UUID `json:"organization_id"`
+	// ProvisionerTags are optional tags used to route template version
+	// provisioning jobs to specific provisioner daemons.
+	ProvisionerTags map[string]string `json:"provisioner_tags"`
 	// NumPresets is the number of presets the template should have.
 	NumPresets int `json:"num_presets"`
 	// NumPresetPrebuilds is the number of prebuilds per preset.

scaletest/prebuilds/run.go

@@ -259,11 +259,12 @@ func (r *Runner) createTemplateVersion(ctx context.Context, templateID uuid.UUID
 	}
 
 	versionReq := codersdk.CreateTemplateVersionRequest{
-		TemplateID:    templateID,
-		FileID:        uploadResp.ID,
-		Message:       "Template version for scaletest prebuilds",
-		StorageMethod: codersdk.ProvisionerStorageMethodFile,
-		Provisioner:   codersdk.ProvisionerTypeTerraform,
+		TemplateID:      templateID,
+		FileID:          uploadResp.ID,
+		Message:         "Template version for scaletest prebuilds",
+		StorageMethod:   codersdk.ProvisionerStorageMethodFile,
+		Provisioner:     codersdk.ProvisionerTypeTerraform,
+		ProvisionerTags: r.cfg.ProvisionerTags,
 	}
 	version, err := r.client.CreateTemplateVersion(ctx, r.cfg.OrganizationID, versionReq)
 	if err != nil {