Add provisionedThroughput for hyperdisk

sunnylovestiramisu · sunnylovestiramisu · commit 6cf846be4c96 · 2022-12-28T10:53:56.000-08:00
diff --git a/README.md b/README.md
@@ -65,7 +65,7 @@ See Github [Issues](https://github.com/kubernetes-sigs/gcp-compute-persistent-di
 | 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. |
-
+| provisioned-throughput-on-create  | string (int64 format). Values typically between 1 and 7,124 mb per second |               | Indicates how much throughput to provision for the disk. See the [hyperdisk documentation](TBD) for details, including valid ranges for throughput. |
 
 ### Topology
 
diff --git a/pkg/common/parameters.go b/pkg/common/parameters.go
@@ -23,11 +23,12 @@ import (
 
 const (
 	// Parameters for StorageClass
-	ParameterKeyType                    = "type"
-	ParameterKeyReplicationType         = "replication-type"
-	ParameterKeyDiskEncryptionKmsKey    = "disk-encryption-kms-key"
-	ParameterKeyLabels                  = "labels"
-	ParameterKeyProvisionedIOPSOnCreate = "provisioned-iops-on-create"
+	ParameterKeyType                          = "type"
+	ParameterKeyReplicationType               = "replication-type"
+	ParameterKeyDiskEncryptionKmsKey          = "disk-encryption-kms-key"
+	ParameterKeyLabels                        = "labels"
+	ParameterKeyProvisionedIOPSOnCreate       = "provisioned-iops-on-create"
+	ParameterKeyProvisionedThroughputOnCreate = "provisioned-throughput-on-create"
 
 	// Parameters for VolumeSnapshotClass
 	ParameterKeyStorageLocations = "storage-locations"
@@ -84,6 +85,9 @@ type DiskParameters struct {
 	// Values: {int64}
 	// Default: none
 	ProvisionedIOPSOnCreate int64
+	// Values: {int64}
+	// Default: none
+	ProvisionedThroughputOnCreate int64
 }
 
 // SnapshotParameters contains normalized and defaulted parameters for snapshots
@@ -153,6 +157,12 @@ func ExtractAndDefaultParameters(parameters map[string]string, driverName string
 				return p, fmt.Errorf("parameters contain invalid provisionedIOPSOnCreate parameter: %w", err)
 			}
 			p.ProvisionedIOPSOnCreate = paramProvisionedIOPSOnCreate
+		case ParameterKeyProvisionedThroughputOnCreate:
+			paramProvisionedThroughputOnCreate, err := ConvertMiBStringToInt64(v)
+			if err != nil {
+				return p, fmt.Errorf("parameters contain invalid provisionedThroughputOnCreate parameter: %w", err)
+			}
+			p.ProvisionedThroughputOnCreate = paramProvisionedThroughputOnCreate
 		default:
 			return p, fmt.Errorf("parameters contains invalid option %q", k)
 		}
diff --git a/pkg/common/parameters_test.go b/pkg/common/parameters_test.go
@@ -90,6 +90,22 @@ func TestExtractAndDefaultParameters(t *testing.T) {
 				ProvisionedIOPSOnCreate: 10000,
 			},
 		},
+		{
+			name:       "values from parameters, checking hyperdisk-throughput",
+			parameters: map[string]string{ParameterKeyType: "hyperdisk-throughput", ParameterKeyReplicationType: "none", ParameterKeyDiskEncryptionKmsKey: "foo/key", ParameterKeyLabels: "key1=value1,key2=value2", ParameterKeyProvisionedThroughputOnCreate: "1000Mi"},
+			labels:     map[string]string{},
+			expectParams: DiskParameters{
+				DiskType:             "hyperdisk-throughput",
+				ReplicationType:      "none",
+				DiskEncryptionKMSKey: "foo/key",
+				Tags:                 map[string]string{},
+				Labels: map[string]string{
+					"key1": "value1",
+					"key2": "value2",
+				},
+				ProvisionedIOPSOnCreate: 1000,
+			},
+		},
 		{
 			name:       "values from parameters, checking balanced pd",
 			parameters: map[string]string{ParameterKeyType: "pd-balanced", ParameterKeyReplicationType: "regional-pd", ParameterKeyDiskEncryptionKmsKey: "foo/key"},
diff --git a/pkg/common/utils.go b/pkg/common/utils.go
@@ -261,3 +261,14 @@ func ConvertGiBStringToInt64(str string) (int64, error) {
 	quantity := resource.MustParse(str)
 	return volumehelpers.RoundUpToGiB(quantity)
 }
+
+// ConvertMiBStringToInt64 converts a GiB string to int64
+func ConvertMiBStringToInt64(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)
+	}
+	quantity := resource.MustParse(str)
+	return volumehelpers.RoundUpToMiB(quantity)
+}
diff --git a/pkg/common/utils_test.go b/pkg/common/utils_test.go
@@ -667,3 +667,93 @@ func TestConvertGiBStringToInt64(t *testing.T) {
 		})
 	}
 }
+
+func TestConvertMiBStringToInt64(t *testing.T) {
+	tests := []struct {
+		desc        string
+		inputStr    string
+		expInt64    int64
+		expectError bool
+	}{
+		{
+			"valid number string",
+			"10000",
+			1,
+			false,
+		},
+		{
+			"round Ki to MiB",
+			"1000Ki",
+			1,
+			false,
+		},
+		{
+			"round k to MiB",
+			"1000k",
+			1,
+			false,
+		},
+		{
+			"round Mi to MiB",
+			"1000Mi",
+			1000,
+			false,
+		},
+		{
+			"round M to MiB",
+			"1000M",
+			954,
+			false,
+		},
+		{
+			"round G to MiB",
+			"1000G",
+			953675,
+			false,
+		},
+		{
+			"round Gi to MiB",
+			"10000Gi",
+			10240000,
+			false,
+		},
+		{
+			"round decimal to MiB",
+			"1.2Gi",
+			1229,
+			false,
+		},
+		{
+			"round big value to MiB",
+			"8191Pi",
+			8795019280384,
+			false,
+		},
+		{
+			"invalid empty string",
+			"",
+			10000,
+			true,
+		},
+		{
+			"invalid string",
+			"ew%65",
+			10000,
+			true,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			actualInt64, err := ConvertMiBStringToInt64(tc.inputStr)
+			if err != nil && !tc.expectError {
+				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 converting string to int64 %s; expect an error", tc.inputStr)
+			}
+			if err == nil && actualInt64 != tc.expInt64 {
+				t.Errorf("Got %d for converting string to int64; expect %d", actualInt64, tc.expInt64)
+			}
+		})
+	}
+}
diff --git a/pkg/gce-cloud-provider/compute/fake-gce.go b/pkg/gce-cloud-provider/compute/fake-gce.go
@@ -196,14 +196,15 @@ 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,
-		ProvisionedIops: params.ProvisionedIOPSOnCreate,
+		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,
+		ProvisionedThroughput: params.ProvisionedThroughputOnCreate,
 	}
 
 	if snapshotID != "" {
diff --git a/pkg/gce-cloud-provider/compute/gce-compute.go b/pkg/gce-cloud-provider/compute/gce-compute.go
@@ -429,12 +429,13 @@ 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,
-		ProvisionedIops: params.ProvisionedIOPSOnCreate,
+		Name:                  volKey.Name,
+		SizeGb:                common.BytesToGbRoundUp(capBytes),
+		Description:           description,
+		Type:                  cloud.GetDiskTypeURI(cloud.project, volKey, params.DiskType),
+		Labels:                params.Labels,
+		ProvisionedIops:       params.ProvisionedIOPSOnCreate,
+		ProvisionedThroughput: params.ProvisionedThroughputOnCreate,
 	}
 	if snapshotID != "" {
 		_, snapshotType, _, err := common.SnapshotIDToProjectKey(snapshotID)
diff --git a/pkg/gce-pd-csi-driver/controller_test.go b/pkg/gce-pd-csi-driver/controller_test.go
@@ -817,6 +817,31 @@ func TestCreateVolumeArguments(t *testing.T) {
 			},
 			expErrCode: codes.InvalidArgument,
 		},
+		{
+			name: "success with provisionedThroughput parameter",
+			req: &csi.CreateVolumeRequest{
+				Name:               name,
+				CapacityRange:      stdCapRange,
+				VolumeCapabilities: stdVolCaps,
+				Parameters:         map[string]string{"labels": "key1=value1,key2=value2", "provisioned-throughput-on-create": "10000"},
+			},
+			expVol: &csi.Volume{
+				CapacityBytes:      common.GbToBytes(20),
+				VolumeId:           testVolumeID,
+				VolumeContext:      nil,
+				AccessibleTopology: stdTopology,
+			},
+		},
+		{
+			name: "fail with malformed provisionedThroughput parameter",
+			req: &csi.CreateVolumeRequest{
+				Name:               name,
+				CapacityRange:      stdCapRange,
+				VolumeCapabilities: stdVolCaps,
+				Parameters:         map[string]string{"labels": "key1=value1,key2=value2", "provisioned-throughput-on-create": "dsfo3"},
+			},
+			expErrCode: codes.InvalidArgument,
+		},
 	}
 
 	// Run test cases
