Add provisionedIops for pd-extreme

Author: sunnylovestiramisu

README.md

@@ -58,12 +58,14 @@ See Github [Issues](https://github.com/kubernetes-sigs/gcp-compute-persistent-di
 
 ### CreateVolume Parameters
 
-| Parameter        | Values                    | Default       | Description                                                                                        |
-|------------------|---------------------------|---------------|----------------------------------------------------------------------------------------------------|
-| type             | Any PD type (see [GCP documentation](https://cloud.google.com/compute/docs/disks#disk-types)), eg `pd-ssd` `pd-balanced` | `pd-standard` | Type allows you to choose between standard Persistent Disks  or Solid State Drive Persistent Disks |
-| replication-type | `none` OR `regional-pd`   | `none`        | Replication type allows you to choose between Zonal Persistent Disks or Regional Persistent Disks  |
-| disk-encryption-kms-key | Fully qualified resource identifier for the key to use to encrypt new disks. | Empty string. | Encrypt disk using Customer Managed Encryption Key (CMEK). See [GKE Docs](https://cloud.google.com/kubernetes-engine/docs/how-to/using-cmek#create_a_cmek_protected_attached_disk) for details. |
-| labels           | `key1=value1,key2=value2` |               | Labels allow you to assign custom [GCE Disk labels](https://cloud.google.com/compute/docs/labeling-resources). |
+| Parameter                   | Values                    | Default       | Description                                                                                        |
+|-----------------------------|---------------------------|---------------|----------------------------------------------------------------------------------------------------|
+| type                        | Any PD type (see [GCP documentation](https://cloud.google.com/compute/docs/disks#disk-types)), eg `pd-ssd` `pd-balanced` | `pd-standard` | Type allows you to choose between standard Persistent Disks  or Solid State Drive Persistent Disks |
+| replication-type            | `none` OR `regional-pd`   | `none`        | Replication type allows you to choose between Zonal Persistent Disks or Regional Persistent Disks  |
+| disk-encryption-kms-key     | Fully qualified resource identifier for the key to use to encrypt new disks. | Empty string. | Encrypt disk using Customer Managed Encryption Key (CMEK). See [GKE Docs](https://cloud.google.com/kubernetes-engine/docs/how-to/using-cmek#create_a_cmek_protected_attached_disk) for details. |
+| labels                      | `key1=value1,key2=value2` |               | Labels allow you to assign custom [GCE Disk labels](https://cloud.google.com/compute/docs/labeling-resources). |
+| provisioned-iops-on-create  | string (int64 format). Values typically between 10,000 and 120,000 |               | Indicates how many IOPS to provision for the disk. See the [Extreme persistent disk documentation](https://cloud.google.com/compute/docs/disks/extreme-persistent-disk) for details, including valid ranges for IOPS. |
+
 
 ### Topology
 

go.mod

@@ -31,6 +31,11 @@ require (
 
 require google.golang.org/protobuf v1.26.0
 
+require (
+	google.golang.org/protobuf v1.26.0
+	k8s.io/cloud-provider v0.22.0
+)
+
 require (
 	github.com/Microsoft/go-winio v0.4.16 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect

go.sum

@@ -1197,6 +1197,7 @@ k8s.io/apiserver v0.22.0/go.mod h1:04kaIEzIQrTGJ5syLppQWvpkLJXQtJECHmae+ZGc/nc=
 k8s.io/cli-runtime v0.22.0/go.mod h1:An6zELQ7udUI0GaXvkuMqyopPA14dIgNqpH8cZu1vig=
 k8s.io/client-go v0.22.0 h1:sD6o9O6tCwUKCENw8v+HFsuAbq2jCu8cWC61/ydwA50=
 k8s.io/client-go v0.22.0/go.mod h1:GUjIuXR5PiEv/RVK5OODUsm6eZk7wtSWZSaSJbpFdGg=
+k8s.io/cloud-provider v0.22.0 h1:eK0swLQ1TZCLefRbgwEo/ZS4ZDo6FkOJDkDIBITshyw=
 k8s.io/cloud-provider v0.22.0/go.mod h1:UsQNOxrStwOXoDfVNgEbKgcQt2BYuHGKobixm0zKTis=
 k8s.io/cluster-bootstrap v0.22.0/go.mod h1:VeZXiGfH+yfnC2KtvkSwNTAqahg6yiCV/szbWpoI+3k=
 k8s.io/code-generator v0.22.0/go.mod h1:eV77Y09IopzeXOJzndrDyCI88UBok2h6WxAlBwpxa+o=

pkg/common/parameters.go

@@ -23,10 +23,11 @@ import (
 
 const (
 	// Parameters for StorageClass
-	ParameterKeyType                 = "type"
-	ParameterKeyReplicationType      = "replication-type"
-	ParameterKeyDiskEncryptionKmsKey = "disk-encryption-kms-key"
-	ParameterKeyLabels               = "labels"
+	ParameterKeyType                    = "type"
+	ParameterKeyReplicationType         = "replication-type"
+	ParameterKeyDiskEncryptionKmsKey    = "disk-encryption-kms-key"
+	ParameterKeyLabels                  = "labels"
+	ParameterKeyProvisionedIOPSOnCreate = "provisioned-iops-on-create"
 
 	// Parameters for VolumeSnapshotClass
 	ParameterKeyStorageLocations = "storage-locations"
@@ -75,6 +76,9 @@ type DiskParameters struct {
 	// Values: {map[string]string}
 	// Default: ""
 	Labels map[string]string
+	// Values: {int64}
+	// Default: none
+	ProvisionedIOPSOnCreate int64
 }
 
 // SnapshotParameters contains normalized and defaulted parameters for snapshots
@@ -134,6 +138,12 @@ func ExtractAndDefaultParameters(parameters map[string]string, driverName string
 			for labelKey, labelValue := range paramLabels {
 				p.Labels[labelKey] = labelValue
 			}
+		case ParameterKeyProvisionedIOPSOnCreate:
+			paramProvisionedIOPSOnCreate, err := ConvertGiBStringToInt64(v)
+			if err != nil {
+				return p, fmt.Errorf("parameters contain invalid provisionedIOPSOnCreate parameter: %w", err)
+			}
+			p.ProvisionedIOPSOnCreate = paramProvisionedIOPSOnCreate
 		default:
 			return p, fmt.Errorf("parameters contains invalid option %q", k)
 		}

pkg/common/parameters_test.go

@@ -74,6 +74,22 @@ func TestExtractAndDefaultParameters(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:       "values from parameters, checking pd-extreme",
+			parameters: map[string]string{ParameterKeyType: "pd-extreme", ParameterKeyReplicationType: "none", ParameterKeyDiskEncryptionKmsKey: "foo/key", ParameterKeyLabels: "key1=value1,key2=value2", ParameterKeyProvisionedIOPSOnCreate: "10000Gi"},
+			labels:     map[string]string{},
+			expectParams: DiskParameters{
+				DiskType:             "pd-extreme",
+				ReplicationType:      "none",
+				DiskEncryptionKMSKey: "foo/key",
+				Tags:                 map[string]string{},
+				Labels: map[string]string{
+					"key1": "value1",
+					"key2": "value2",
+				},
+				ProvisionedIOPSOnCreate: 10000,
+			},
+		},
 		{
 			name:       "values from parameters, checking balanced pd",
 			parameters: map[string]string{ParameterKeyType: "pd-balanced", ParameterKeyReplicationType: "regional-pd", ParameterKeyDiskEncryptionKmsKey: "foo/key"},

pkg/common/utils.go

@@ -22,7 +22,9 @@ import (
 	"strings"
 
 	"github.com/GoogleCloudPlatform/k8s-cloud-provider/pkg/cloud/meta"
+	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/util/sets"
+	volumehelpers "k8s.io/cloud-provider/volume/helpers"
 )
 
 const (
@@ -257,14 +259,13 @@ func ValidateSnapshotType(snapshotType string) error {
 	}
 }
 
-// ParseMachineType returns an extracted machineType from a URL, or empty if not found.
-// machineTypeUrl: Full or partial URL of the machine type resource, in the format:
-//
-//	zones/zone/machineTypes/machine-type
-func ParseMachineType(machineTypeUrl string) (string, error) {
-	machineType := machineTypeRegex.FindStringSubmatch(machineTypeUrl)
-	if machineType == nil {
-		return "", fmt.Errorf("failed to parse machineTypeUrl. Expected suffix: zones/{zone}/machineTypes/{machine-type}. Got: %s", machineTypeUrl)
+// ConvertGiBStringToInt64 converts a GiB string to int64
+func ConvertGiBStringToInt64(str string) (int64, error) {
+	// Verify regex before
+	match, _ := regexp.MatchString("^([+-]?[0-9.]+)([eEinumkKMGTP]*[-+]?[0-9]*)$", str)
+	if !match {
+		return 0, fmt.Errorf("invalid string %s", str)
 	}
-	return machineType[1], nil
+	quantity := resource.MustParse(str)
+	return volumehelpers.RoundUpToGiB(quantity)
 }

pkg/common/utils_test.go

@@ -578,55 +578,91 @@ func TestSnapshotStorageLocations(t *testing.T) {
 	}
 }
 
-func TestParseMachineType(t *testing.T) {
+func TestConvertGiBStringToInt64(t *testing.T) {
 	tests := []struct {
-		desc                string
-		inputMachineTypeUrl string
-		expectedMachineType string
-		expectError         bool
+		desc        string
+		inputStr    string
+		expInt64    int64
+		expectError bool
 	}{
 		{
-			desc:                "full URL machine type",
-			inputMachineTypeUrl: "https://www.googleapis.com/compute/v1/projects/my-project/zones/us-central1-c/machineTypes/c3-highcpu-4",
-			expectedMachineType: "c3-highcpu-4",
+			"valid number string",
+			"10000",
+			1,
+			false,
 		},
 		{
-			desc:                "partial URL machine type",
-			inputMachineTypeUrl: "zones/us-central1-c/machineTypes/n2-standard-4",
-			expectedMachineType: "n2-standard-4",
+			"round Ki to GiB",
+			"1000Ki",
+			1,
+			false,
 		},
 		{
-			desc:                "custom partial URL machine type",
-			inputMachineTypeUrl: "zones/us-central1-c/machineTypes/e2-custom-2-4096",
-			expectedMachineType: "e2-custom-2-4096",
+			"round k to GiB",
+			"1000k",
+			1,
+			false,
 		},
 		{
-			desc:                "incorrect URL",
-			inputMachineTypeUrl: "https://www.googleapis.com/compute/v1/projects/psch-gke-dev/zones/us-central1-c",
-			expectError:         true,
+			"round Mi to GiB",
+			"1000Mi",
+			1,
+			false,
 		},
 		{
-			desc:                "incorrect partial URL",
-			inputMachineTypeUrl: "zones/us-central1-c/machineTypes/",
-			expectError:         true,
+			"round M to GiB",
+			"1000M",
+			1,
+			false,
+		},
+		{
+			"round G to GiB",
+			"1000G",
+			932,
+			false,
+		},
+		{
+			"round Gi to GiB",
+			"10000Gi",
+			10000,
+			false,
 		},
 		{
-			desc:                "missing zone",
-			inputMachineTypeUrl: "zones//machineTypes/n2-standard-4",
-			expectError:         true,
+			"round decimal to GiB",
+			"1.2Gi",
+			2,
+			false,
+		},
+		{
+			"round big value to GiB",
+			"8191Pi",
+			8588886016,
+			false,
+		},
+		{
+			"invalid empty string",
+			"",
+			10000,
+			true,
+		},
+		{
+			"invalid string",
+			"ew%65",
+			10000,
+			true,
 		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {
-			actualMachineFamily, err := ParseMachineType(tc.inputMachineTypeUrl)
+			actualInt64, err := ConvertGiBStringToInt64(tc.inputStr)
 			if err != nil && !tc.expectError {
-				t.Errorf("Got error %v parsing machine type %s; expect no error", err, tc.inputMachineTypeUrl)
+				t.Errorf("Got error %v converting string to int64 %s; expect no error", err, tc.inputStr)
 			}
 			if err == nil && tc.expectError {
-				t.Errorf("Got no error parsing machine type %s; expect an error", tc.inputMachineTypeUrl)
+				t.Errorf("Got no error converting string to int64 %s; expect an error", tc.inputStr)
 			}
-			if err == nil && actualMachineFamily != tc.expectedMachineType {
-				t.Errorf("Got %s parsing machine type; expect %s", actualMachineFamily, tc.expectedMachineType)
+			if err == nil && actualInt64 != tc.expInt64 {
+				t.Errorf("Got %d for converting string to int64; expect %d", actualInt64, tc.expInt64)
 			}
 		})
 	}

pkg/gce-cloud-provider/compute/fake-gce.go

@@ -196,13 +196,14 @@ func (cloud *FakeCloudProvider) InsertDisk(ctx context.Context, project string,
 	}
 
 	computeDisk := &computev1.Disk{
-		Name:         volKey.Name,
-		SizeGb:       common.BytesToGbRoundUp(capBytes),
-		Description:  "Disk created by GCE-PD CSI Driver",
-		Type:         cloud.GetDiskTypeURI(project, volKey, params.DiskType),
-		SourceDiskId: volumeContentSourceVolumeID,
-		Status:       cloud.mockDiskStatus,
-		Labels:       params.Labels,
+		Name:            volKey.Name,
+		SizeGb:          common.BytesToGbRoundUp(capBytes),
+		Description:     "Disk created by GCE-PD CSI Driver",
+		Type:            cloud.GetDiskTypeURI(project, volKey, params.DiskType),
+		SourceDiskId:    volumeContentSourceVolumeID,
+		Status:          cloud.mockDiskStatus,
+		Labels:          params.Labels,
+		ProvisionedIops: params.ProvisionedIOPSOnCreate,
 	}
 
 	if snapshotID != "" {

pkg/gce-cloud-provider/compute/gce-compute.go

@@ -443,11 +443,12 @@ func (cloud *CloudProvider) insertRegionalDisk(
 	}
 
 	diskToCreate := &computev1.Disk{
-		Name:        volKey.Name,
-		SizeGb:      common.BytesToGbRoundUp(capBytes),
-		Description: description,
-		Type:        cloud.GetDiskTypeURI(cloud.project, volKey, params.DiskType),
-		Labels:      params.Labels,
+		Name:            volKey.Name,
+		SizeGb:          common.BytesToGbRoundUp(capBytes),
+		Description:     description,
+		Type:            cloud.GetDiskTypeURI(cloud.project, volKey, params.DiskType),
+		Labels:          params.Labels,
+		ProvisionedIops: params.ProvisionedIOPSOnCreate,
 	}
 	if snapshotID != "" {
 		_, snapshotType, _, err := common.SnapshotIDToProjectKey(snapshotID)

pkg/gce-pd-csi-driver/controller_test.go

@@ -793,6 +793,31 @@ func TestCreateVolumeArguments(t *testing.T) {
 			},
 			expErrCode: codes.InvalidArgument,
 		},
+		{
+			name: "success with provisionedIops parameter",
+			req: &csi.CreateVolumeRequest{
+				Name:               name,
+				CapacityRange:      stdCapRange,
+				VolumeCapabilities: stdVolCaps,
+				Parameters:         map[string]string{"labels": "key1=value1,key2=value2", "provisioned-iops-on-create": "10000"},
+			},
+			expVol: &csi.Volume{
+				CapacityBytes:      common.GbToBytes(20),
+				VolumeId:           testVolumeID,
+				VolumeContext:      nil,
+				AccessibleTopology: stdTopology,
+			},
+		},
+		{
+			name: "fail with malformed provisionedIops parameter",
+			req: &csi.CreateVolumeRequest{
+				Name:               name,
+				CapacityRange:      stdCapRange,
+				VolumeCapabilities: stdVolCaps,
+				Parameters:         map[string]string{"labels": "key1=value1,key2=value2", "provisioned-iops-on-create": "dsfo3"},
+			},
+			expErrCode: codes.InvalidArgument,
+		},
 	}
 
 	// Run test cases