Merge pull request #1101 from sunnylovestiramisu/hyperdisk

Add provisionedThroughput for hyperdisk-throughput

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
 

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)
 		}

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",
+				},
+				ProvisionedThroughputOnCreate: 1000,
+			},
+		},
 		{
 			name:       "values from parameters, checking balanced pd",
 			parameters: map[string]string{ParameterKeyType: "pd-balanced", ParameterKeyReplicationType: "regional-pd", ParameterKeyDiskEncryptionKmsKey: "foo/key"},

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)
+}

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)
+			}
+		})
+	}
+}

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

@@ -40,6 +40,8 @@ const (
 	cryptoKeyVerDelimiter          = "/cryptoKeyVersions"
 )
 
+var hyperdiskTypes = []string{"hyperdisk-extreme", "hyperdisk-throughput"}
+
 type GCEAPIVersion string
 
 const (
@@ -387,15 +389,15 @@ func convertV1CustomerEncryptionKeyToBeta(v1Key *computev1.CustomerEncryptionKey
 	}
 }
 
-func convertV1DiskToBetaDisk(v1Disk *computev1.Disk) *computebeta.Disk {
+func convertV1DiskToBetaDisk(v1Disk *computev1.Disk, provisionedThroughputOnCreate int64) *computebeta.Disk {
 	var dek *computebeta.CustomerEncryptionKey = nil
 
 	if v1Disk.DiskEncryptionKey != nil {
 		dek = convertV1CustomerEncryptionKeyToBeta(v1Disk.DiskEncryptionKey)
 	}
 
 	// Note: this is an incomplete list. It only includes the fields we use for disk creation.
-	return &computebeta.Disk{
+	betaDisk := &computebeta.Disk{
 		Name:              v1Disk.Name,
 		SizeGb:            v1Disk.SizeGb,
 		Description:       v1Disk.Description,
@@ -404,6 +406,11 @@ func convertV1DiskToBetaDisk(v1Disk *computev1.Disk) *computebeta.Disk {
 		ReplicaZones:      v1Disk.ReplicaZones,
 		DiskEncryptionKey: dek,
 	}
+	if provisionedThroughputOnCreate > 0 {
+		betaDisk.ProvisionedThroughput = provisionedThroughputOnCreate
+	}
+
+	return betaDisk
 }
 
 func (cloud *CloudProvider) insertRegionalDisk(
@@ -465,7 +472,7 @@ func (cloud *CloudProvider) insertRegionalDisk(
 
 	if gceAPIVersion == GCEAPIVersionBeta {
 		var insertOp *computebeta.Operation
-		betaDiskToCreate := convertV1DiskToBetaDisk(diskToCreate)
+		betaDiskToCreate := convertV1DiskToBetaDisk(diskToCreate, 0)
 		betaDiskToCreate.MultiWriter = multiWriter
 		insertOp, err = cloud.betaService.RegionDisks.Insert(project, volKey.Region, betaDiskToCreate).Context(ctx).Do()
 		if insertOp != nil {
@@ -537,7 +544,7 @@ func (cloud *CloudProvider) insertZonalDisk(
 		gceAPIVersion = GCEAPIVersionV1
 	)
 
-	if multiWriter {
+	if multiWriter || containsBetaDiskType(hyperdiskTypes, params.DiskType) {
 		gceAPIVersion = GCEAPIVersionBeta
 	}
 
@@ -575,7 +582,7 @@ func (cloud *CloudProvider) insertZonalDisk(
 
 	if gceAPIVersion == GCEAPIVersionBeta {
 		var insertOp *computebeta.Operation
-		betaDiskToCreate := convertV1DiskToBetaDisk(diskToCreate)
+		betaDiskToCreate := convertV1DiskToBetaDisk(diskToCreate, params.ProvisionedThroughputOnCreate)
 		betaDiskToCreate.MultiWriter = multiWriter
 		insertOp, err = cloud.betaService.Disks.Insert(project, volKey.Zone, betaDiskToCreate).Context(ctx).Do()
 		if insertOp != nil {
@@ -1167,3 +1174,13 @@ func encodeTags(tags map[string]string) (string, error) {
 	}
 	return string(enc), nil
 }
+
+func containsBetaDiskType(betaDiskTypes []string, diskType string) bool {
+	for _, betaDiskType := range betaDiskTypes {
+		if betaDiskType == diskType {
+			return true
+		}
+	}
+
+	return false
+}

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

test/e2e/tests/setup_e2e_test.go

@@ -27,6 +27,7 @@ import (
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 	computealpha "google.golang.org/api/compute/v0.alpha"
+	computebeta "google.golang.org/api/compute/v0.beta"
 	compute "google.golang.org/api/compute/v1"
 	"k8s.io/klog/v2"
 	testutils "sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/test/e2e/utils"
@@ -46,6 +47,7 @@ var (
 	testContexts        = []*remote.TestContext{}
 	computeService      *compute.Service
 	computeAlphaService *computealpha.Service
+	computeBetaService  *computebeta.Service
 	kmsClient           *cloudkms.KeyManagementClient
 )
 
@@ -75,6 +77,9 @@ var _ = BeforeSuite(func() {
 	computeAlphaService, err = remote.GetComputeAlphaClient()
 	Expect(err).To(BeNil())
 
+	computeBetaService, err = remote.GetComputeBetaClient()
+	Expect(err).To(BeNil())
+
 	// Create the KMS client.
 	kmsClient, err = cloudkms.NewKeyManagementClient(context.Background())
 	Expect(err).To(BeNil())

test/remote/instance.go

@@ -28,6 +28,7 @@ import (
 
 	"golang.org/x/oauth2/google"
 	computealpha "google.golang.org/api/compute/v0.alpha"
+	computebeta "google.golang.org/api/compute/v0.beta"
 	compute "google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
 	"k8s.io/apimachinery/pkg/util/uuid"
@@ -71,12 +72,11 @@ func (i *InstanceInfo) GetNodeID() string {
 
 func CreateInstanceInfo(project, instanceArchitecture, instanceZone, name, machineType string, cs *compute.Service) (*InstanceInfo, error) {
 	return &InstanceInfo{
-		project:      project,
-		architecture: instanceArchitecture,
-		zone:         instanceZone,
-		name:         name,
-		machineType:  machineType,
-
+		project:        project,
+		architecture:   instanceArchitecture,
+		zone:           instanceZone,
+		name:           name,
+		machineType:    machineType,
 		computeService: cs,
 	}, nil
 }
@@ -330,6 +330,36 @@ func GetComputeAlphaClient() (*computealpha.Service, error) {
 	return nil, err
 }
 
+func GetComputeBetaClient() (*computebeta.Service, error) {
+	const retries = 10
+	const backoff = time.Second * 6
+
+	klog.V(4).Infof("Getting compute client...")
+
+	// Setup the gce client for provisioning instances
+	// Getting credentials on gce jenkins is flaky, so try a couple times
+	var err error
+	var cs *computebeta.Service
+	for i := 0; i < retries; i++ {
+		if i > 0 {
+			time.Sleep(backoff)
+		}
+
+		var client *http.Client
+		client, err = google.DefaultClient(context.Background(), computebeta.ComputeScope)
+		if err != nil {
+			continue
+		}
+
+		cs, err = computebeta.New(client)
+		if err != nil {
+			continue
+		}
+		return cs, nil
+	}
+	return nil, err
+}
+
 func generateMetadataWithPublicKey(pubKeyFile string) (*compute.Metadata, error) {
 	publicKeyByte, err := ioutil.ReadFile(pubKeyFile)
 	if err != nil {