Repository: kubernetes-sigs/aws-ebs-csi-driver
Pull request #2463 (2025-05-06 18:56:13)

Add support for specifying volume initialization rate
#### What type of PR is this?

/kind feature

<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature
-->

#### What is this PR about? / Why do we need it?

This PR introduces a new StorageClass parameter and functionality that allows for requesting a provisioned initialization rate when creating a volume from a snapshot.

#### How was this change tested?

Manually, unit / e2e testing.

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, enter your extended release note in the block below.
-->
```release-note
Introduced new StorageClass parameter: volumeInitializationRate - when creating a volume from a snapshot, this parameter can be used to request a provisioned initialization rate, in MiB/s.
```


k8s-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by:
**Once this PR has been reviewed and has the lgtm label**, please assign [wongma7](https://github.com/wongma7) for approval. For more information see [the Code Review Process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process).

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Faws-ebs-csi-driver).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/kubernetes-sigs/aws-ebs-csi-driver/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["wongma7"]} -->
k8s-ci-robot: @torredil: The following test **failed**, say `/retest` to rerun all failed tests or `/retest-required` to rerun all mandatory failed tests:

Test name | Commit | Details | Required | Rerun command
--- | --- | --- | --- | ---
pull-aws-ebs-csi-driver-verify | b260346049670326f3a075d267b366aafba4cfe7 | [link](https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_aws-ebs-csi-driver/2463/pull-aws-ebs-csi-driver-verify/1919828576962613248) | true | `/test pull-aws-ebs-csi-driver-verify`

[Full PR test history](https://prow.k8s.io/pr-history?org=kubernetes-sigs&repo=aws-ebs-csi-driver&pr=2463). [Your PR dashboard](https://prow.k8s.io/pr?query=is%3Apr%20state%3Aopen%20author%3Atorredil). Please help us cut down on flakes by [linking to](https://git.k8s.io/community/contributors/devel/sig-testing/flaky-tests.md#filing-issues-for-flaky-tests) an [open issue](https://github.com/kubernetes-sigs/aws-ebs-csi-driver/issues?q=is:issue+is:open) when you hit one in your PR.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes-sigs/prow](https://github.com/kubernetes-sigs/prow/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
github-actions[bot]: <!-- pr-coverage -->
## Code Coverage Diff
File | Old Coverage | New Coverage | Delta
---- |:------------:|:------------:|:-----:
github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/driver/controller.go | 83.8% | 84.0% | 0.1

Patch:
diff --git a/docs/parameters.md b/docs/parameters.md
index 5790d6c5..bc930499 100644
--- a/docs/parameters.md
+++ b/docs/parameters.md
@@ -22,6 +22,7 @@ The AWS EBS CSI Driver supports [tagging](tagging.md) through `StorageClass.para
 | "numberOfInodes"             |                                                   |         | The `number-of-inodes` to use when formatting the underlying filesystem. Only supported on linux nodes and with fstype `ext3`, `ext4`.                                                                                                                                                                                                                                                |
 | "ext4BigAlloc"               | true, false                                       | false   | Changes the `ext4` filesystem to use clustered block allocation by enabling the `bigalloc` formatting option. Warning: `bigalloc` may not be fully supported with your node's Linux kernel. Please see our [FAQ](/docs/faq.md).                                                                                                                                                               |
 | "ext4ClusterSize"            |                                                   |         | The cluster size to use when formatting an `ext4` filesystem when the `bigalloc` feature is enabled. Note: The `ext4BigAlloc` parameter must be set to true. See our [FAQ](/docs/faq.md).                                                                                                                                                                                                     |
+| "volumeInitializationRate"   | integer                                           |         |  When creating a volume from a snapshot, this parameter can be used to request a provisioned initialization rate, in MiB/s.                             |
 
 ## Restrictions
 * `gp3` is currently not supported on outposts. Outpost customers need to use a different type for their volumes.
diff --git a/go.mod b/go.mod
index 8a8b0129..8d8509f7 100644
--- a/go.mod
+++ b/go.mod
@@ -4,7 +4,7 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.36.3
 	github.com/aws/aws-sdk-go-v2/config v1.29.14
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30
-	github.com/aws/aws-sdk-go-v2/service/ec2 v1.214.0
+	github.com/aws/aws-sdk-go-v2/service/ec2 v1.215.0
 	github.com/aws/smithy-go v1.22.3
 	github.com/awslabs/volume-modifier-for-k8s v0.5.1
 	github.com/container-storage-interface/spec v1.11.0
diff --git a/go.sum b/go.sum
index 57cf1ece..767bad10 100644
--- a/go.sum
+++ b/go.sum
@@ -28,8 +28,8 @@ github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34 h1:SZwFm17ZUNNg5Np0io
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34/go.mod h1:dFZsC0BLo346mvKQLWmoJxT+Sjp+qcVR1tRVHQGOH9Q=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 h1:bIqFDwgGXXN1Kpp99pDOdKMTTb5d2KyU5X/BZxjOkRo=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo=
-github.com/aws/aws-sdk-go-v2/service/ec2 v1.214.0 h1:cFR69no3REgjwiGNHhSeB8IZ1sGrLAoLqHACMNuN+j4=
-github.com/aws/aws-sdk-go-v2/service/ec2 v1.214.0/go.mod h1:ouvGEfHbLaIlWwpDpOVWPWR+YwO0HDv3vm5tYLq8ImY=
+github.com/aws/aws-sdk-go-v2/service/ec2 v1.215.0 h1:6a5U/gnVIPWWtS2CCdOkrxos3Se8IL9jNMJLyD4BEU8=
+github.com/aws/aws-sdk-go-v2/service/ec2 v1.215.0/go.mod h1:ouvGEfHbLaIlWwpDpOVWPWR+YwO0HDv3vm5tYLq8ImY=
 github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 h1:eAh2A4b5IzM/lum78bZ590jy36+d/aFLgKF/4Vd1xPE=
 github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3/go.mod h1:0yKJC/kb8sAnmlYa6Zs3QVYqaC8ug2AbnNChv5Ox3uA=
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 h1:dM9/92u2F1JbDaGooxTq18wmmFzbJRfXfVfy96/1CXM=
diff --git a/pkg/cloud/cloud.go b/pkg/cloud/cloud.go
index 68167eeb..36998fc0 100644
--- a/pkg/cloud/cloud.go
+++ b/pkg/cloud/cloud.go
@@ -217,8 +217,9 @@ type DiskOptions struct {
 	MultiAttachEnabled     bool
 	// KmsKeyID represents a fully qualified resource name to the key to use for encryption.
 	// example: arn:aws:kms:us-east-1:012345678910:key/abcd1234-a123-456a-a12b-a123b4cd56ef
-	KmsKeyID   string
-	SnapshotID string
+	KmsKeyID                 string
+	SnapshotID               string
+	VolumeInitializationRate int32
 }
 
 // ModifyDiskOptions represents parameters to modify an EBS volume.
@@ -648,6 +649,9 @@ func (c *cloud) CreateDisk(ctx context.Context, volumeName string, diskOptions *
 	if len(snapshotID) > 0 {
 		requestInput.SnapshotId = aws.String(snapshotID)
 	}
+	if diskOptions.VolumeInitializationRate > 0 {
+		requestInput.VolumeInitializationRate = aws.Int32(diskOptions.VolumeInitializationRate)
+	}
 
 	response, err := c.ec2.CreateVolume(ctx, requestInput, func(o *ec2.Options) {
 		o.Retryer = c.rm.createVolumeRetryer
diff --git a/pkg/driver/constants.go b/pkg/driver/constants.go
index 71dfe4f1..2cca8e82 100644
--- a/pkg/driver/constants.go
+++ b/pkg/driver/constants.go
@@ -43,6 +43,9 @@ const (
 	// AllowAutoIOPSPerGBIncreaseKey represents key for allowing automatic increase of IOPS.
 	AllowAutoIOPSPerGBIncreaseKey = "allowautoiopspergbincrease"
 
+	// VolumeInitializationRateKey represents key for volume initialization rate when creating volumes from snapshots.
+	VolumeInitializationRateKey = "volumeinitializationrate"
+
 	// Iops represents key for IOPS for volume.
 	IopsKey = "iops"
 
diff --git a/pkg/driver/controller.go b/pkg/driver/controller.go
index 352cc278..4f531470 100644
--- a/pkg/driver/controller.go
+++ b/pkg/driver/controller.go
@@ -106,16 +106,17 @@ func (d *ControllerService) CreateVolume(ctx context.Context, req *csi.CreateVol
 	defer d.inFlight.Delete(volName)
 
 	var (
-		volumeType             string
-		iopsPerGB              int32
-		allowIOPSPerGBIncrease bool
-		iops                   int32
-		throughput             int32
-		isEncrypted            bool
-		blockExpress           bool
-		kmsKeyID               string
-		scTags                 []string
-		volumeTags             = map[string]string{
+		volumeType               string
+		iopsPerGB                int32
+		allowIOPSPerGBIncrease   bool
+		iops                     int32
+		throughput               int32
+		volumeInitializationRate int32
+		isEncrypted              bool
+		blockExpress             bool
+		kmsKeyID                 string
+		scTags                   []string
+		volumeTags               = map[string]string{
 			cloud.VolumeNameTagKey:   volName,
 			cloud.AwsEbsDriverTagKey: isManagedByDriver,
 		}
@@ -149,6 +150,12 @@ func (d *ControllerService) CreateVolume(ctx context.Context, req *csi.CreateVol
 				return nil, status.Errorf(codes.InvalidArgument, "Could not parse invalid iops: %v", err)
 			}
 			iops = int32(parseIopsKey)
+		case VolumeInitializationRateKey:
+			parseInitRate, parseInitRateErr := strconv.ParseInt(value, 10, 32)
+			if parseInitRateErr != nil {
+				return nil, status.Errorf(codes.InvalidArgument, "Could not parse invalid volumeInitializationRate: %v", err)
+			}
+			volumeInitializationRate = int32(parseInitRate)
 		case ThroughputKey:
 			parseThroughput, parseThroughputErr := strconv.ParseInt(value, 10, 32)
 			if parseThroughputErr != nil {
@@ -312,20 +319,21 @@ func (d *ControllerService) CreateVolume(ctx context.Context, req *csi.CreateVol
 	}
 
 	opts := &cloud.DiskOptions{
-		CapacityBytes:          volSizeBytes,
-		Tags:                   volumeTags,
-		VolumeType:             volumeType,
-		IOPSPerGB:              iopsPerGB,
-		AllowIOPSPerGBIncrease: allowIOPSPerGBIncrease,
-		IOPS:                   iops,
-		Throughput:             throughput,
-		AvailabilityZone:       zone,
-		OutpostArn:             outpostArn,
-		Encrypted:              isEncrypted,
-		BlockExpress:           blockExpress,
-		KmsKeyID:               kmsKeyID,
-		SnapshotID:             snapshotID,
-		MultiAttachEnabled:     multiAttach,
+		CapacityBytes:            volSizeBytes,
+		Tags:                     volumeTags,
+		VolumeType:               volumeType,
+		IOPSPerGB:                iopsPerGB,
+		AllowIOPSPerGBIncrease:   allowIOPSPerGBIncrease,
+		IOPS:                     iops,
+		Throughput:               throughput,
+		AvailabilityZone:         zone,
+		OutpostArn:               outpostArn,
+		Encrypted:                isEncrypted,
+		BlockExpress:             blockExpress,
+		KmsKeyID:                 kmsKeyID,
+		SnapshotID:               snapshotID,
+		MultiAttachEnabled:       multiAttach,
+		VolumeInitializationRate: volumeInitializationRate,
 	}
 
 	disk, err := d.cloud.CreateDisk(ctx, volName, opts)


Test patch:
diff --git a/pkg/cloud/cloud_test.go b/pkg/cloud/cloud_test.go
index 90092474..d9435fc8 100644
--- a/pkg/cloud/cloud_test.go
+++ b/pkg/cloud/cloud_test.go
@@ -1336,6 +1336,40 @@ func TestCreateDisk(t *testing.T) {
 			},
 			expErr: nil,
 		},
+		{
+			name:       "success: create volume from snapshot with initialization rate",
+			volumeName: "vol-test-name",
+			diskOptions: &DiskOptions{
+				CapacityBytes:            util.GiBToBytes(1),
+				Tags:                     map[string]string{VolumeNameTagKey: "vol-test", AwsEbsDriverTagKey: "true"},
+				SnapshotID:               "snapshot-test",
+				VolumeInitializationRate: 200,
+			},
+			expDisk: &Disk{
+				VolumeID:         "vol-test",
+				CapacityGiB:      1,
+				AvailabilityZone: defaultZone,
+			},
+			expCreateVolumeInput: &ec2.CreateVolumeInput{
+				VolumeInitializationRate: aws.Int32(200),
+			},
+			expErr: nil,
+		},
+		{
+			name:       "failure: create volume from snapshot with initialization rate when snapshotID is missing",
+			volumeName: "vol-test-name",
+			diskOptions: &DiskOptions{
+				CapacityBytes:            util.GiBToBytes(1),
+				Tags:                     map[string]string{VolumeNameTagKey: "vol-test", AwsEbsDriverTagKey: "true"},
+				VolumeInitializationRate: 200,
+			},
+			expDisk: nil,
+			expCreateVolumeInput: &ec2.CreateVolumeInput{
+				VolumeInitializationRate: aws.Int32(200),
+			},
+			expCreateVolumeErr: errors.New("InvalidParameterCombination"),
+			expErr:             fmt.Errorf("could not create volume in EC2: %w", errors.New("InvalidParameterCombination")),
+		},
 		{
 			name:       "failure: multi-attach with GP3",
 			volumeName: "vol-test-name",
diff --git a/pkg/driver/controller_test.go b/pkg/driver/controller_test.go
index b80e56f7..2183bad6 100644
--- a/pkg/driver/controller_test.go
+++ b/pkg/driver/controller_test.go
@@ -377,6 +377,101 @@ func TestCreateVolume(t *testing.T) {
 				checkExpectedErrorCode(t, err, codes.AlreadyExists)
 			},
 		},
+		{
+			name: "success with valid initialization rate from snapshot",
+			testFunc: func(t *testing.T) {
+				t.Helper()
+				req := &csi.CreateVolumeRequest{
+					Name:               "vol-test",
+					CapacityRange:      stdCapRange,
+					VolumeCapabilities: stdVolCap,
+					Parameters: map[string]string{
+						VolumeInitializationRateKey: "200",
+					},
+					VolumeContentSource: &csi.VolumeContentSource{
+						Type: &csi.VolumeContentSource_Snapshot{
+							Snapshot: &csi.VolumeContentSource_SnapshotSource{
+								SnapshotId: "snapshot-test",
+							},
+						},
+					},
+				}
+
+				mockDisk := &cloud.Disk{
+					VolumeID:         req.GetName(),
+					AvailabilityZone: expZone,
+					CapacityGiB:      util.BytesToGiB(stdVolSize),
+					SnapshotID:       "snapshot-test",
+				}
+
+				mockCtl := gomock.NewController(t)
+				defer mockCtl.Finish()
+
+				mockCloud := cloud.NewMockCloud(mockCtl)
+				mockCloud.EXPECT().CreateDisk(gomock.Eq(t.Context()), gomock.Eq(req.GetName()),
+					gomock.Any()).DoAndReturn(func(_ context.Context, _ string, diskOptions *cloud.DiskOptions) (*cloud.Disk, error) {
+					if diskOptions.VolumeInitializationRate != 200 {
+						t.Errorf("Expected VolumeInitializationRate to be 200, got %d", diskOptions.VolumeInitializationRate)
+					}
+					return mockDisk, nil
+				})
+
+				awsDriver := ControllerService{
+					cloud:    mockCloud,
+					inFlight: internal.NewInFlight(),
+					options:  &Options{},
+				}
+
+				if _, err := awsDriver.CreateVolume(t.Context(), req); err != nil {
+					t.Fatalf("CreateVolume failed: %v", err)
+				}
+			},
+		},
+		{
+			name: "fail with invalid initialization rate value",
+			testFunc: func(t *testing.T) {
+				t.Helper()
+				req := &csi.CreateVolumeRequest{
+					Name:               "vol-test",
+					CapacityRange:      stdCapRange,
+					VolumeCapabilities: stdVolCap,
+					Parameters: map[string]string{
+						VolumeInitializationRateKey: "invalid",
+					},
+					VolumeContentSource: &csi.VolumeContentSource{
+						Type: &csi.VolumeContentSource_Snapshot{
+							Snapshot: &csi.VolumeContentSource_SnapshotSource{
+								SnapshotId: "snapshot-test",
+							},
+						},
+					},
+				}
+
+				mockCtl := gomock.NewController(t)
+				defer mockCtl.Finish()
+
+				mockCloud := cloud.NewMockCloud(mockCtl)
+
+				awsDriver := ControllerService{
+					cloud:    mockCloud,
+					inFlight: internal.NewInFlight(),
+					options:  &Options{},
+				}
+
+				_, err := awsDriver.CreateVolume(t.Context(), req)
+				if err == nil {
+					t.Fatal("Expected CreateVolume to fail but got no error")
+				}
+
+				srvErr, ok := status.FromError(err)
+				if !ok {
+					t.Fatalf("Could not get error status code from error: %v", srvErr)
+				}
+				if srvErr.Code() != codes.InvalidArgument {
+					t.Fatalf("Expected InvalidArgument but got: %s", srvErr.Code())
+				}
+			},
+		},
 		{
 			name: "fail no name",
 			testFunc: func(t *testing.T) {
diff --git a/tests/e2e/requires_aws_api.go b/tests/e2e/requires_aws_api.go
index c5a56965..ec4846d7 100644
--- a/tests/e2e/requires_aws_api.go
+++ b/tests/e2e/requires_aws_api.go
@@ -17,6 +17,7 @@ package e2e
 import (
 	"context"
 	"fmt"
+	"strconv"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -177,6 +178,79 @@ var _ = Describe("[ebs-csi-e2e] [single-az] [requires-aws-api] Dynamic Provision
 		test.Run(cs, snapshotrcs, ns)
 	})
 
+	It("should create a volume from a snapshot with initialization rate", func() {
+		restoredTagKey := generateTagName()
+		restoredTagVal := "volume-init-rate-test"
+		initRate := "200"
+
+		originalPod := testsuites.PodDetails{
+			Cmd: testsuites.PodCmdWriteToVolume("/mnt/test-1"),
+			Volumes: []testsuites.VolumeDetails{
+				{
+					CreateVolumeParameters: map[string]string{
+						ebscsidriver.VolumeTypeKey: awscloud.VolumeTypeGP3,
+						ebscsidriver.FSTypeKey:     ebscsidriver.FSTypeExt4,
+					},
+					ClaimSize:   driver.MinimumSizeForVolumeType(awscloud.VolumeTypeGP3),
+					VolumeMount: testsuites.DefaultGeneratedVolumeMount,
+				},
+			},
+		}
+
+		restoredPod := testsuites.PodDetails{
+			Cmd: testsuites.PodCmdGrepVolumeData("/mnt/test-1"),
+			Volumes: []testsuites.VolumeDetails{
+				{
+					CreateVolumeParameters: map[string]string{
+						ebscsidriver.VolumeTypeKey:               awscloud.VolumeTypeGP3,
+						ebscsidriver.FSTypeKey:                   ebscsidriver.FSTypeExt4,
+						ebscsidriver.VolumeInitializationRateKey: initRate,
+						ebscsidriver.TagKeyPrefix:                fmt.Sprintf("%s=%s", restoredTagKey, restoredTagVal),
+					},
+					ClaimSize:   driver.MinimumSizeForVolumeType(awscloud.VolumeTypeGP3),
+					VolumeMount: testsuites.DefaultGeneratedVolumeMount,
+				},
+			},
+		}
+
+		test := testsuites.DynamicallyProvisionedVolumeSnapshotTest{
+			CSIDriver:   ebsDriver,
+			Pod:         originalPod,
+			RestoredPod: restoredPod,
+
+			ValidateFunc: func(snapshot *volumesnapshotv1.VolumeSnapshot) {
+				intInitRate, err := strconv.Atoi(initRate)
+				if err != nil {
+					Fail(fmt.Sprintf("failed to parse initRate: %v", err))
+				}
+
+				result, err := ec2Client.DescribeVolumes(context.Background(), &ec2.DescribeVolumesInput{
+					Filters: []types.Filter{
+						{
+							Name:   aws.String("tag:" + restoredTagKey),
+							Values: []string{restoredTagVal},
+						},
+					},
+				})
+				if err != nil {
+					Fail(fmt.Sprintf("failed to describe volume: %v", err))
+				}
+				if len(result.Volumes) != 1 {
+					Fail(fmt.Sprintf("expected 1 volume, got %d", len(result.Volumes)))
+				}
+
+				vol := result.Volumes[0]
+				if vol.VolumeInitializationRate == nil {
+					Fail("VolumeInitializationRate is unexpectedly nil")
+				}
+				if *vol.VolumeInitializationRate != int32(intInitRate) {
+					Fail(fmt.Sprintf("VolumeInitializationRate mismatch: got %d, expected %d", *vol.VolumeInitializationRate, intInitRate))
+				}
+			},
+		}
+		test.Run(cs, snapshotrcs, ns)
+	})
+
 	It("should create a snapshot with FSR enabled", func() {
 		azList, err := ec2Client.DescribeAvailabilityZones(context.Background(), &ec2.DescribeAvailabilityZonesInput{})
 		if err != nil {