From 372537a032f6a392e42e6ec66a55fbe681fa561b Mon Sep 17 00:00:00 2001 From: Amelia Crate <95060558+a-crate@users.noreply.github.com> Date: Fri, 13 Oct 2023 23:16:52 +0000 Subject: [PATCH] Add a wait for disk quota step (#847) * Add a wait for disk quota step When finalizing workflows with wait for X quota steps, copying workflow dependencies from the create step to the quota step to ensure that creation is done as close as possible to finding quota availability Abstract wait for X quota steps into one function Add quotas with identical metrics and regions together when waiting on them - Unit test this behavior Wait for SSD quota in storageperf disk setup. * Update fixtures.go * gofmt --- imagetest/fixtures.go | 36 +++++++--- imagetest/fixtures_test.go | 78 +++++++++++++++++----- imagetest/test_suites/storageperf/setup.go | 13 ++-- imagetest/testworkflow.go | 34 +++++++--- 4 files changed, 122 insertions(+), 39 deletions(-) diff --git a/imagetest/fixtures.go b/imagetest/fixtures.go index b68052a17..16dd135e5 100644 --- a/imagetest/fixtures.go +++ b/imagetest/fixtures.go @@ -27,13 +27,14 @@ import ( ) const ( - waitForVMQuotaStepName = "wait-for-vm-quota" - createVMsStepName = "create-vms" - createDisksStepName = "create-disks" - createNetworkStepName = "create-networks" - createFirewallStepName = "create-firewalls" - createSubnetworkStepName = "create-sub-networks" - successMatch = "FINISHED-TEST" + waitForVMQuotaStepName = "wait-for-vm-quota" + createVMsStepName = "create-vms" + createDisksStepName = "create-disks" + waitForDisksQuotaStepName = "wait-for-disk-quota" + createNetworkStepName = "create-networks" + createFirewallStepName = "create-firewalls" + createSubnetworkStepName = "create-sub-networks" + successMatch = "FINISHED-TEST" // ShouldRebootDuringTest is a local map key to indicate that the // test will reboot and relies on results from the second boot. ShouldRebootDuringTest = "shouldRebootDuringTest" @@ -85,15 +86,32 @@ func (t *TestWorkflow) LockProject() { // WaitForVMQuota appends a list of quotas to the wait for vm quota step. Quotas with a blank region will be populated with the region corresponding to the workflow zone. func (t *TestWorkflow) WaitForVMQuota(qa *daisy.QuotaAvailable) error { - step, ok := t.wf.Steps[waitForVMQuotaStepName] + return t.waitForQuotaStep(qa, waitForVMQuotaStepName) +} + +// WaitForDisksQuota appends a list of quotas to the wait for disk quota step. Quotas with a blank region will be populated with the region corresponding to the workflow zone. +func (t *TestWorkflow) WaitForDisksQuota(qa *daisy.QuotaAvailable) error { + return t.waitForQuotaStep(qa, waitForDisksQuotaStepName) +} + +func (t *TestWorkflow) waitForQuotaStep(qa *daisy.QuotaAvailable, stepname string) error { + step, ok := t.wf.Steps[stepname] if !ok { var err error - step, err = t.wf.NewStep(waitForVMQuotaStepName) + step, err = t.wf.NewStep(stepname) if err != nil { return err } step.WaitForAvailableQuotas = &daisy.WaitForAvailableQuotas{Interval: "90s"} } + // If the step is already waiting for this quota, add the number of units to + // the existing quota. + for _, q := range step.WaitForAvailableQuotas.Quotas { + if q.Metric == qa.Metric && q.Region == qa.Region { + q.Units = q.Units + qa.Units + return nil + } + } step.WaitForAvailableQuotas.Quotas = append(step.WaitForAvailableQuotas.Quotas, qa) return nil } diff --git a/imagetest/fixtures_test.go b/imagetest/fixtures_test.go index 1314a244e..0263558b4 100644 --- a/imagetest/fixtures_test.go +++ b/imagetest/fixtures_test.go @@ -229,24 +229,66 @@ func TestEnableSecureBoot(t *testing.T) { } } -// TestWaitForVMQuota tests that WaitForVMQuota successfully appends a quota to -// the wait for vm quota step. -func TestWaitForVMQuota(t *testing.T) { - twf, err := NewTestWorkflow("name", "image", "30m", "x86", "arm64") - if err != nil { - t.Errorf("failed to create test workflow: %v", err) - } - quota := &daisy.QuotaAvailable{Metric: "test", Units: 1, Region: "us-central1"} - err = twf.WaitForVMQuota(quota) - if err != nil { - t.Errorf("failed to append quota: %v", err) - } - quotaStep, ok := twf.wf.Steps[waitForVMQuotaStepName] - if !ok { - t.Errorf("Could not find wait for vm quota step") - } - if quotaStep.WaitForAvailableQuotas.Quotas[0] != quota { - t.Errorf("An unexpected quota was appended") +// TestWaitForQuotaStep tests that quotas are successfully appended to the wait +// step. +func TestWaitForQuotaStep(t *testing.T) { + testcases := []struct { + name string + input []*daisy.QuotaAvailable + output []*daisy.QuotaAvailable + }{ + { + name: "single quota", + input: []*daisy.QuotaAvailable{{Metric: "test", Units: 1, Region: "us-central1"}}, + output: []*daisy.QuotaAvailable{{Metric: "test", Units: 1, Region: "us-central1"}}, + }, + { + name: "two independent quotas", + input: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test1", Units: 1, Region: "us-west1"}}, + output: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test1", Units: 1, Region: "us-west1"}}, + }, + { + name: "two quotas same region", + input: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test1", Units: 1, Region: "us-central1"}}, + output: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test1", Units: 1, Region: "us-central1"}}, + }, + { + name: "two quotas same metric", + input: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test2", Units: 1, Region: "us-west1"}}, + output: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test2", Units: 1, Region: "us-west1"}}, + }, + { + name: "two identical quotas", + input: []*daisy.QuotaAvailable{{Metric: "test2", Units: 2, Region: "us-central1"}, {Metric: "test2", Units: 1, Region: "us-central1"}}, + output: []*daisy.QuotaAvailable{{Metric: "test2", Units: 3, Region: "us-central1"}}, + }, + } + for _, tc := range testcases { + t.Run(tc.name, func(t *testing.T) { + twf, err := NewTestWorkflow("name", "image", "30m", "x86", "arm64") + if err != nil { + t.Errorf("failed to create test workflow: %v", err) + } + for _, quota := range tc.input { + err = twf.waitForQuotaStep(quota, tc.name) + if err != nil { + t.Errorf("failed to append quota: %v", err) + } + } + quotaStep, ok := twf.wf.Steps[tc.name] + if !ok { + t.Errorf("Could not find wait for vm quota step") + } + if len(quotaStep.WaitForAvailableQuotas.Quotas) != len(tc.output) { + t.Errorf("unexpected output length from WaitForVMQuota, got %d want %d", len(quotaStep.WaitForAvailableQuotas.Quotas), len(tc.output)) + } + for i := range tc.output { + q := quotaStep.WaitForAvailableQuotas.Quotas[i] + if q.Metric != tc.output[i].Metric || q.Units != tc.output[i].Units || q.Region != tc.output[i].Region { + t.Errorf("unexpected quota at position %d\ngot %v\nwant %v", i, *q, *tc.output[i]) + } + } + }) } } diff --git a/imagetest/test_suites/storageperf/setup.go b/imagetest/test_suites/storageperf/setup.go index f233dd492..4ef02556b 100644 --- a/imagetest/test_suites/storageperf/setup.go +++ b/imagetest/test_suites/storageperf/setup.go @@ -106,11 +106,16 @@ func TestSetup(t *imagetest.TestWorkflow) error { continue } + region := tc.zone + if len(region) > 2 { + region = region[:len(region)-2] + } + if err := t.WaitForDisksQuota(&daisy.QuotaAvailable{Metric: "SSD_TOTAL_GB", Units: bootdiskSizeGB + mountdiskSizeGB, Region: region}); err != nil { + return err + } if tc.cpuMetric != "" { - quota := &daisy.QuotaAvailable{Metric: tc.cpuMetric, Region: tc.zone} - if len(quota.Region) > 2 { - quota.Region = quota.Region[:len(quota.Region)-2] - } + quota := &daisy.QuotaAvailable{Metric: tc.cpuMetric, Region: region} + i, err := strconv.ParseFloat(regexp.MustCompile("-[0-9]+$").FindString(tc.machineType)[1:], 64) if err != nil { return err diff --git a/imagetest/testworkflow.go b/imagetest/testworkflow.go index 322dd5356..eb133a406 100644 --- a/imagetest/testworkflow.go +++ b/imagetest/testworkflow.go @@ -412,23 +412,41 @@ func finalizeWorkflows(ctx context.Context, tests []*TestWorkflow, zone, gcsPref twf.wf.Zone = zone - vmquotaStep, vmquotaStepOk := twf.wf.Steps[waitForVMQuotaStepName] - if vmquotaStepOk { - for _, q := range vmquotaStep.WaitForAvailableQuotas.Quotas { + // Process quota steps and associated creation steps. + for quotaStepName, createStepName := range map[string]string{ + waitForVMQuotaStepName: createVMsStepName, + waitForDisksQuotaStepName: createDisksStepName, + } { + quotaStep, ok := twf.wf.Steps[quotaStepName] + if !ok { + continue + } + for _, q := range quotaStep.WaitForAvailableQuotas.Quotas { // Populate empty regions with best guess from the zone if q.Region == "" { q.Region = twf.wf.Zone[:len(twf.wf.Zone)-2] } } + createStep, ok := twf.wf.Steps[createStepName] + if !ok { + continue + } + // Fix dependencies. Create steps should depend on the quota step, and quota steps should inherit all other dependencies. + for _, dep := range twf.wf.Dependencies[createStepName] { + dStep, ok := twf.wf.Steps[dep] + if ok { + if err := twf.wf.AddDependency(quotaStep, dStep); err != nil { + return err + } + } + } + if err := twf.wf.AddDependency(createStep, quotaStep); err != nil { + return err + } } createVMsStep, ok := twf.wf.Steps[createVMsStepName] if ok { - if vmquotaStepOk { - if err := twf.wf.AddDependency(createVMsStep, vmquotaStep); err != nil { - return err - } - } for _, vm := range createVMsStep.CreateInstances.Instances { if vm.MachineType != "" { log.Printf("VM %s machine type set to %s for test %s\n", vm.Name, vm.MachineType, twf.Name)