cos-customizer: Add -disk-type flag

This new flag allows users to select the disk type of the preload VM.

BUG=b/229419550
TEST=./run_tests.sh

Change-Id: Ifd1908aa794a0474bd3682cc35e344cae9578be9
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/32087
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
diff --git a/src/cmd/cos_customizer/README.md b/src/cmd/cos_customizer/README.md
index e9d772e..80d83d3 100644
--- a/src/cmd/cos_customizer/README.md
+++ b/src/cmd/cos_customizer/README.md
@@ -223,6 +223,9 @@
 image labels present on the source image. The labels specified by the `-labels`
 flag take precedence over labels assigned with this flag.
 
+`-disk-type`: The disk type for the preload VM to use when creating the image.
+Defaults to `pd-standard`.
+
 `-disk-size-gb`: The disk size in GB to use when creating the image.
 This value should never be smaller than 10 (the default size of a COS image).
 If `-oem-size` is set,  the lower limit of `-disk-size-gb` is as shown in the 
diff --git a/src/cmd/cos_customizer/finish_image_build.go b/src/cmd/cos_customizer/finish_image_build.go
index 51232f5..2c46f7f 100644
--- a/src/cmd/cos_customizer/finish_image_build.go
+++ b/src/cmd/cos_customizer/finish_image_build.go
@@ -54,6 +54,7 @@
 	inheritLabels  bool
 	oemSize        string
 	oemFSSize4K    uint64
+	diskType       string
 	diskSize       int
 	timeout        time.Duration
 }
@@ -114,6 +115,7 @@
 	flags.StringVar(&f.oemSize, "oem-size", "", "Size of the new OEM partition, "+
 		"can be a number with unit like 10G, 10M, 10K or 10B, "+
 		"or without unit indicating the number of 512B sectors.")
+	flags.StringVar(&f.diskType, "disk-type", "pd-standard", "The disk type to use when creating the image.")
 	flags.IntVar(&f.diskSize, "disk-size-gb", 0, "The disk size to use when creating the image in GB. Value of '0' "+
 		"indicates the default size.")
 	flags.DurationVar(&f.timeout, "timeout", time.Hour, "Timeout value of the image build process. Must be formatted "+
@@ -166,6 +168,7 @@
 	buildConfig.Project = f.project
 	buildConfig.Zone = f.zone
 	buildConfig.MachineType = f.machineType
+	buildConfig.DiskType = f.diskType
 	buildConfig.DiskSize = f.diskSize
 	buildConfig.Network = f.network
 	buildConfig.Subnet = f.subnet
diff --git a/src/data/build_image.wf.json b/src/data/build_image.wf.json
index 4c96e3d..1ac7a67 100644
--- a/src/data/build_image.wf.json
+++ b/src/data/build_image.wf.json
@@ -6,6 +6,7 @@
     "output_image_family": {"Value": "", "Description": "Family of output image."},
     "output_image_project": {"Required": true, "Description": "Project of output image."},
     "cidata_img": {"Required": true, "Description": "Path to CIDATA vfat image containing cloud-init user-data and the provisioner program. Must be in .tar.gz format."},
+    "disk_type": {"Value": "pd-standard", "Description": "The disk type to use for preloading."},
     "disk_size_gb": {"Value": "10", "Description": "The disk size to use for preloading."},
     "host_maintenance": {"Value": "MIGRATE", "Description": "VM behavior when there is maintenance."},
     "machine_type": {"Required": true, "Description": "Machine type of the preload VM."},
@@ -40,11 +41,13 @@
         {
           "Name": "boot-disk",
           "SourceImage": "${source_image}",
+          "Type": "${disk_type}",
           "SizeGb": "${disk_size_gb}"
         },
         {
           "Name": "cidata-disk",
-          "SourceImage": "cidata"
+          "SourceImage": "cidata",
+          "Type": "${disk_type}"
         }
       ]
     },
diff --git a/src/pkg/config/config.go b/src/pkg/config/config.go
index 57debde..f60752b 100644
--- a/src/pkg/config/config.go
+++ b/src/pkg/config/config.go
@@ -67,6 +67,7 @@
 	Project     string
 	Zone        string
 	MachineType string
+	DiskType    string
 	DiskSize    int
 	GPUType     string
 	Timeout     string
diff --git a/src/pkg/preloader/BUILD.bazel b/src/pkg/preloader/BUILD.bazel
index c7e110a..b6b088a 100644
--- a/src/pkg/preloader/BUILD.bazel
+++ b/src/pkg/preloader/BUILD.bazel
@@ -62,6 +62,7 @@
 
 go_test(
     name = "preloader_test",
+    size = "large",
     srcs = [
         "gcs_test.go",
         "preload_test.go",
diff --git a/src/pkg/preloader/preload.go b/src/pkg/preloader/preload.go
index 1e8973f..cae3ffe 100644
--- a/src/pkg/preloader/preload.go
+++ b/src/pkg/preloader/preload.go
@@ -308,6 +308,8 @@
 		ciDataFile,
 		"-var:machine_type",
 		buildSpec.MachineType,
+		"-var:disk_type",
+		buildSpec.DiskType,
 		"-var:host_maintenance",
 		hostMaintenance,
 		"-var:network",
diff --git a/src/pkg/preloader/preload_test.go b/src/pkg/preloader/preload_test.go
index 2f7f1ee..7de309a 100644
--- a/src/pkg/preloader/preload_test.go
+++ b/src/pkg/preloader/preload_test.go
@@ -316,6 +316,13 @@
 			want:        []string{"-var:cidata_img"},
 		},
 		{
+			testName:    "DiskType",
+			inputImage:  config.NewImage("", ""),
+			outputImage: config.NewImage("", ""),
+			buildConfig: &config.Build{DiskType: "pd-ssd", GCSBucket: "bucket", GCSDir: "dir"},
+			want:        []string{"-var:disk_type", "pd-ssd"},
+		},
+		{
 			testName:    "DiskSize",
 			inputImage:  config.NewImage("", ""),
 			outputImage: config.NewImage("", ""),