diff --git a/cli/exp_scaletest_dynamicparameters.go b/cli/exp_scaletest_dynamicparameters.go index 3c6922a8d4..31b6766ac6 100644 --- a/cli/exp_scaletest_dynamicparameters.go +++ b/cli/exp_scaletest_dynamicparameters.go @@ -27,6 +27,7 @@ const ( func (r *RootCmd) scaletestDynamicParameters() *serpent.Command { var ( templateName string + provisionerTags []string numEvals int64 tracingFlags = &scaletestTracingFlags{} prometheusFlags = &scaletestPrometheusFlags{} @@ -56,6 +57,11 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command { return xerrors.Errorf("template cannot be empty") } + tags, err := ParseProvisionerTags(provisionerTags) + if err != nil { + return err + } + org, err := orgContext.Selected(inv, client) if err != nil { return err @@ -99,7 +105,7 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command { }() tracer := tracerProvider.Tracer(scaletestTracerName) - partitions, err := dynamicparameters.SetupPartitions(ctx, client, org.ID, templateName, numEvals, logger) + partitions, err := dynamicparameters.SetupPartitions(ctx, client, org.ID, templateName, tags, numEvals, logger) if err != nil { return xerrors.Errorf("setup dynamic parameters partitions: %w", err) } @@ -160,6 +166,11 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command { Default: "100", Value: serpent.Int64Of(&numEvals), }, + { + Flag: "provisioner-tag", + Description: "Specify a set of tags to target provisioner daemons.", + Value: serpent.StringArrayOf(&provisionerTags), + }, } orgContext.AttachOptions(cmd) output.attach(&cmd.Options) diff --git a/scaletest/dynamicparameters/template.go b/scaletest/dynamicparameters/template.go index 9e7d8bc97f..5faf67e531 100644 --- a/scaletest/dynamicparameters/template.go +++ b/scaletest/dynamicparameters/template.go @@ -168,12 +168,13 @@ type SDKForDynamicParametersSetup interface { // partitioner is an internal struct to hold context and arguments for partition setup // and to provide methods for all sub-steps. type partitioner struct { - ctx context.Context - client SDKForDynamicParametersSetup - orgID uuid.UUID - templateName string - numEvals int64 - logger slog.Logger + ctx context.Context + client SDKForDynamicParametersSetup + orgID uuid.UUID + templateName string + provisionerTags map[string]string + numEvals int64 + logger slog.Logger // for testing clock quartz.Clock @@ -181,17 +182,19 @@ type partitioner struct { func SetupPartitions( ctx context.Context, client SDKForDynamicParametersSetup, - orgID uuid.UUID, templateName string, numEvals int64, + orgID uuid.UUID, templateName string, provisionerTags map[string]string, + numEvals int64, logger slog.Logger, ) ([]Partition, error) { p := &partitioner{ - ctx: ctx, - client: client, - orgID: orgID, - templateName: templateName, - numEvals: numEvals, - logger: logger, - clock: quartz.NewReal(), + ctx: ctx, + client: client, + orgID: orgID, + templateName: templateName, + provisionerTags: provisionerTags, + numEvals: numEvals, + logger: logger, + clock: quartz.NewReal(), } return p.run() } @@ -272,11 +275,12 @@ func (p *partitioner) createTemplateVersion(templateID uuid.UUID) (codersdk.Temp // Create template version versionReq := codersdk.CreateTemplateVersionRequest{ - TemplateID: templateID, - FileID: uploadResp.ID, - Message: "Initial version for scaletest dynamic parameters", - StorageMethod: codersdk.ProvisionerStorageMethodFile, - Provisioner: codersdk.ProvisionerTypeTerraform, + TemplateID: templateID, + FileID: uploadResp.ID, + Message: "Initial version for scaletest dynamic parameters", + StorageMethod: codersdk.ProvisionerStorageMethodFile, + Provisioner: codersdk.ProvisionerTypeTerraform, + ProvisionerTags: p.provisionerTags, } version, err := p.client.CreateTemplateVersion(p.ctx, p.orgID, versionReq) if err != nil { diff --git a/scaletest/dynamicparameters/template_internal_test.go b/scaletest/dynamicparameters/template_internal_test.go index 0b000a4c74..6b1230eeae 100644 --- a/scaletest/dynamicparameters/template_internal_test.go +++ b/scaletest/dynamicparameters/template_internal_test.go @@ -70,6 +70,7 @@ func TestSetupPartitions_TemplateExists(t *testing.T) { t: t, expectedTemplateName: "test-template", expectedOrgID: orgID, + expectedTags: map[string]string{"foo": "bar"}, matchedProvisioners: 1, templateVersionJobStatus: codersdk.ProvisionerJobSucceeded, } @@ -77,13 +78,14 @@ func TestSetupPartitions_TemplateExists(t *testing.T) { trap := mClock.Trap().TickerFunc("waitForTemplateVersionJobs") defer trap.Close() uut := partitioner{ - ctx: ctx, - client: fClient, - orgID: orgID, - templateName: "test-template", - numEvals: 600, - logger: logger, - clock: mClock, + ctx: ctx, + client: fClient, + orgID: orgID, + templateName: "test-template", + provisionerTags: map[string]string{"foo": "bar"}, + numEvals: 600, + logger: logger, + clock: mClock, } var partitions []Partition errCh := make(chan error, 1) @@ -234,6 +236,7 @@ type fakeClient struct { expectedOrgID uuid.UUID templateByNameError error + expectedTags map[string]string matchedProvisioners int templateVersionJobStatus codersdk.ProvisionerJobStatus @@ -270,6 +273,7 @@ func (f *fakeClient) CreateTemplate(ctx context.Context, orgID uuid.UUID, create func (f *fakeClient) CreateTemplateVersion(ctx context.Context, orgID uuid.UUID, createReq codersdk.CreateTemplateVersionRequest) (codersdk.TemplateVersion, error) { f.templateVersionsCount++ + require.Equal(f.t, f.expectedTags, createReq.ProvisionerTags) return codersdk.TemplateVersion{ ID: uuid.New(), Name: f.expectedTemplateName,