Allow users to specify "storage-locations" for snapshots.

TeweiLuo · TeweiLuo · commit 31a8f87f161e · 2021-06-23T20:45:19.000Z
diff --git a/docs/kubernetes/user-guides/snapshots.md b/docs/kubernetes/user-guides/snapshots.md
@@ -23,6 +23,14 @@ Kubernetes 1.17+.
     kubectl create -f ./examples/kubernetes/snapshot/default-volumesnapshotclass.yaml
     ```
 
+  To place the snapshot in a [custom storage location](https://cloud.google.com/compute/docs/disks/snapshots#custom_location),
+  edit `volumesnapshotclass-storage-locations.yaml` to change the `storage-locations` parameter to a location of your
+  choice, and then create the `VolumeSnapshotClass`.
+
+    ```console
+    kubectl create -f ./examples/kubernetes/snapshot/volumesnapshotclass-storage-locations.yaml
+    ```
+
 1. Create source PVC
 
     ```console
diff --git a/examples/kubernetes/snapshot/volumesnapshotclass-storage-locations.yaml b/examples/kubernetes/snapshot/volumesnapshotclass-storage-locations.yaml
@@ -0,0 +1,8 @@
+apiVersion: snapshot.storage.k8s.io/v1beta1
+kind: VolumeSnapshotClass
+metadata:
+  name: csi-gce-pd-snapshot-class-with-storage-locations
+parameters:
+  storage-locations: us-east2
+driver: pd.csi.storage.gke.io
+deletionPolicy: Delete
diff --git a/pkg/common/parameters.go b/pkg/common/parameters.go
@@ -22,11 +22,15 @@ import (
 )
 
 const (
+	// Parameters for StorageClass
 	ParameterKeyType                 = "type"
 	ParameterKeyReplicationType      = "replication-type"
 	ParameterKeyDiskEncryptionKmsKey = "disk-encryption-kms-key"
 	ParameterKeyLabels               = "labels"
 
+	// Parameters for VolumeSnapshotClass
+	ParameterKeyStorageLocations = "storage-locations"
+
 	replicationTypeNone = "none"
 
 	// Keys for PV and PVC parameters as reported by external-provisioner
@@ -60,6 +64,11 @@ type DiskParameters struct {
 	Labels map[string]string
 }
 
+// SnapshotParameters contains normalized and defaulted parameters for snapshots
+type SnapshotParameters struct {
+	StorageLocations []string
+}
+
 // ExtractAndDefaultParameters will take the relevant parameters from a map and
 // put them into a well defined struct making sure to default unspecified fields.
 // extraVolumeLabels are added as labels; if there are also labels specified in
@@ -118,3 +127,22 @@ func ExtractAndDefaultParameters(parameters map[string]string, driverName string
 	}
 	return p, nil
 }
+
+func ExtractAndDefaultSnapshotParameters(parameters map[string]string) (SnapshotParameters, error) {
+	p := SnapshotParameters{
+		StorageLocations: []string{},
+	}
+	for k, v := range parameters {
+		switch strings.ToLower(k) {
+		case ParameterKeyStorageLocations:
+			normalizedStorageLocations, err := ProcessStorageLocations(v)
+			if err != nil {
+				return p, err
+			}
+			p.StorageLocations = normalizedStorageLocations
+		default:
+			return p, fmt.Errorf("parameters contains invalid option %q", k)
+		}
+	}
+	return p, nil
+}
diff --git a/pkg/common/parameters_test.go b/pkg/common/parameters_test.go
@@ -164,3 +164,40 @@ func TestExtractAndDefaultParameters(t *testing.T) {
 		})
 	}
 }
+
+// Currently the only parameter is storage-locations, which is already tested
+// in utils_test/TestSnapshotStorageLocations. Here we just test the case where
+// no parameter is set in the snapshot class.
+func TestSnapshotParameters(t *testing.T) {
+	tests := []struct {
+		desc                    string
+		parameters              map[string]string
+		expectedSnapshotParames SnapshotParameters
+	}{
+		{
+			desc:       "valid parameter",
+			parameters: map[string]string{ParameterKeyStorageLocations: "ASIA "},
+			expectedSnapshotParames: SnapshotParameters{
+				StorageLocations: []string{"asia"},
+			},
+		},
+		{
+			desc:       "nil parameter",
+			parameters: nil,
+			expectedSnapshotParames: SnapshotParameters{
+				StorageLocations: []string{},
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			p, err := ExtractAndDefaultSnapshotParameters(tc.parameters)
+			if err != nil {
+				t.Errorf("Got error processing snapshot parameters: %v; expect no error", err)
+			}
+			if !reflect.DeepEqual(p, tc.expectedSnapshotParames) {
+				t.Errorf("Got ExtractAndDefaultSnapshotParameters(%+v) = %+v; expect %+v", tc.parameters, p, tc.expectedSnapshotParames)
+			}
+		})
+	}
+}
diff --git a/pkg/common/utils.go b/pkg/common/utils.go
@@ -50,6 +50,18 @@ const (
 	nodeIDTotalElements = 6
 
 	regionalDeviceNameSuffix = "_regional"
+
+	// Snapshot storage location format
+	// Reference: https://cloud.google.com/storage/docs/locations
+	// Example: us
+	multiRegionalLocationFmt = "^[a-z]+$"
+	// Example: us-east1
+	regionalLocationFmt = "^[a-z]+-[a-z]+[0-9]$"
+)
+
+var (
+	multiRegionalPattern = regexp.MustCompile(multiRegionalLocationFmt)
+	regionalPattern      = regexp.MustCompile(regionalLocationFmt)
 )
 
 func BytesToGbRoundDown(bytes int64) int64 {
@@ -217,3 +229,12 @@ func ConvertLabelsStringToMap(labels string) (map[string]string, error) {
 
 	return labelsMap, nil
 }
+
+// ProcessStorageLocations trims and normalizes storage location to lower letters.
+func ProcessStorageLocations(storageLocations string) ([]string, error) {
+	normalizedLoc := strings.ToLower(strings.TrimSpace(storageLocations))
+	if !multiRegionalPattern.MatchString(normalizedLoc) && !regionalPattern.MatchString(normalizedLoc) {
+		return []string{}, fmt.Errorf("invalid location for snapshot: %q", storageLocations)
+	}
+	return []string{normalizedLoc}, nil
+}
diff --git a/pkg/common/utils_test.go b/pkg/common/utils_test.go
@@ -534,3 +534,46 @@ func TestConvertLabelsStringToMap(t *testing.T) {
 	})
 
 }
+
+func TestSnapshotStorageLocations(t *testing.T) {
+	tests := []struct {
+		desc                        string
+		locationString              string
+		expectedNormalizedLocations []string
+		expectError                 bool
+	}{
+		{
+			"valid multi-region",
+			"   uS ",
+			[]string{"us"},
+			false,
+		},
+		{
+			"valid region",
+			"  US-EAST1",
+			[]string{"us-east1"},
+			false,
+		},
+		{
+			// Zones are not valid bucket/snapshot locations.
+			"single zone",
+			"us-east1a",
+			[]string{},
+			true,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			normalizedLocations, err := ProcessStorageLocations(tc.locationString)
+			if err != nil && !tc.expectError {
+				t.Errorf("Got error %v processing storage locations %q; expect no error", err, tc.locationString)
+			}
+			if err == nil && tc.expectError {
+				t.Errorf("Got no error processing storage locations %q; expect an error", tc.locationString)
+			}
+			if err == nil && !reflect.DeepEqual(normalizedLocations, tc.expectedNormalizedLocations) {
+				t.Errorf("Got %v for normalized storage locations; expect %v", normalizedLocations, tc.expectedNormalizedLocations)
+			}
+		})
+	}
+}
diff --git a/pkg/gce-cloud-provider/compute/fake-gce.go b/pkg/gce-cloud-provider/compute/fake-gce.go
@@ -387,7 +387,7 @@ func (cloud *FakeCloudProvider) GetSnapshot(ctx context.Context, project, snapsh
 	return snapshot, nil
 }
 
-func (cloud *FakeCloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error) {
+func (cloud *FakeCloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error) {
 	if snapshot, ok := cloud.snapshots[snapshotName]; ok {
 		return snapshot, nil
 	}
@@ -398,6 +398,7 @@ func (cloud *FakeCloudProvider) CreateSnapshot(ctx context.Context, project stri
 		CreationTimestamp: Timestamp,
 		Status:            "UPLOADING",
 		SelfLink:          cloud.getGlobalSnapshotURI(project, snapshotName),
+		StorageLocations:  snapshotParams.StorageLocations,
 	}
 	switch volKey.Type() {
 	case meta.Zonal:
@@ -493,11 +494,11 @@ type FakeBlockingCloudProvider struct {
 // Upon starting a CreateSnapshot, it passes a chan 'executeCreateSnapshot' into readyToExecute, then blocks on executeCreateSnapshot.
 // The test calling this function can block on readyToExecute to ensure that the operation has started and
 // allowed the CreateSnapshot to continue by passing a struct into executeCreateSnapshot.
-func (cloud *FakeBlockingCloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error) {
+func (cloud *FakeBlockingCloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error) {
 	executeCreateSnapshot := make(chan struct{})
 	cloud.ReadyToExecute <- executeCreateSnapshot
 	<-executeCreateSnapshot
-	return cloud.FakeCloudProvider.CreateSnapshot(ctx, project, volKey, snapshotName)
+	return cloud.FakeCloudProvider.CreateSnapshot(ctx, project, volKey, snapshotName, snapshotParams)
 }
 
 func notFoundError() *googleapi.Error {
diff --git a/pkg/gce-cloud-provider/compute/gce-compute.go b/pkg/gce-cloud-provider/compute/gce-compute.go
@@ -73,7 +73,7 @@ type GCECompute interface {
 	ListZones(ctx context.Context, region string) ([]string, error)
 	ListSnapshots(ctx context.Context, filter string, maxEntries int64, pageToken string) ([]*computev1.Snapshot, string, error)
 	GetSnapshot(ctx context.Context, project, snapshotName string) (*computev1.Snapshot, error)
-	CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error)
+	CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error)
 	DeleteSnapshot(ctx context.Context, project, snapshotName string) error
 }
 
@@ -787,13 +787,13 @@ func (cloud *CloudProvider) DeleteSnapshot(ctx context.Context, project, snapsho
 	return nil
 }
 
-func (cloud *CloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error) {
+func (cloud *CloudProvider) CreateSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error) {
 	klog.V(5).Infof("Creating snapshot %s for volume %v", snapshotName, volKey)
 	switch volKey.Type() {
 	case meta.Zonal:
-		return cloud.createZonalDiskSnapshot(ctx, project, volKey, snapshotName)
+		return cloud.createZonalDiskSnapshot(ctx, project, volKey, snapshotName, snapshotParams)
 	case meta.Regional:
-		return cloud.createRegionalDiskSnapshot(ctx, project, volKey, snapshotName)
+		return cloud.createRegionalDiskSnapshot(ctx, project, volKey, snapshotName, snapshotParams)
 	default:
 		return nil, fmt.Errorf("could not create snapshot, key was neither zonal nor regional, instead got: %v", volKey.String())
 	}
@@ -864,9 +864,10 @@ func (cloud *CloudProvider) resizeRegionalDisk(ctx context.Context, project stri
 	return requestGb, nil
 }
 
-func (cloud *CloudProvider) createZonalDiskSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error) {
+func (cloud *CloudProvider) createZonalDiskSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error) {
 	snapshotToCreate := &computev1.Snapshot{
-		Name: snapshotName,
+		Name:             snapshotName,
+		StorageLocations: snapshotParams.StorageLocations,
 	}
 
 	_, err := cloud.service.Disks.CreateSnapshot(project, volKey.Zone, volKey.Name, snapshotToCreate).Context(ctx).Do()
@@ -878,9 +879,10 @@ func (cloud *CloudProvider) createZonalDiskSnapshot(ctx context.Context, project
 	return cloud.waitForSnapshotCreation(ctx, project, snapshotName)
 }
 
-func (cloud *CloudProvider) createRegionalDiskSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string) (*computev1.Snapshot, error) {
+func (cloud *CloudProvider) createRegionalDiskSnapshot(ctx context.Context, project string, volKey *meta.Key, snapshotName string, snapshotParams common.SnapshotParameters) (*computev1.Snapshot, error) {
 	snapshotToCreate := &computev1.Snapshot{
-		Name: snapshotName,
+		Name:             snapshotName,
+		StorageLocations: snapshotParams.StorageLocations,
 	}
 
 	_, err := cloud.service.RegionDisks.CreateSnapshot(project, volKey.Region, volKey.Name, snapshotToCreate).Context(ctx).Do()
diff --git a/pkg/gce-pd-csi-driver/controller.go b/pkg/gce-pd-csi-driver/controller.go
@@ -609,7 +609,11 @@ func (gceCS *GCEControllerServer) CreateSnapshot(ctx context.Context, req *csi.C
 			return nil, status.Error(codes.Internal, fmt.Sprintf("Unknown get snapshot error: %v", err))
 		}
 		// If we could not find the snapshot, we create a new one
-		snapshot, err = gceCS.CloudProvider.CreateSnapshot(ctx, project, volKey, req.Name)
+		snapshotParams, err := common.ExtractAndDefaultSnapshotParameters(req.GetParameters())
+		if err != nil {
+			return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("Invalid snapshot parameters: %v", err))
+		}
+		snapshot, err = gceCS.CloudProvider.CreateSnapshot(ctx, project, volKey, req.Name, snapshotParams)
 		if err != nil {
 			if gce.IsGCEError(err, "notFound") {
 				return nil, status.Error(codes.NotFound, fmt.Sprintf("Could not find volume with ID %v: %v", volKey.String(), err))
diff --git a/pkg/gce-pd-csi-driver/controller_test.go b/pkg/gce-pd-csi-driver/controller_test.go
@@ -84,6 +84,7 @@ func TestCreateSnapshotArguments(t *testing.T) {
 			req: &csi.CreateSnapshotRequest{
 				Name:           name,
 				SourceVolumeId: testVolumeID,
+				Parameters:     map[string]string{common.ParameterKeyStorageLocations: " US-WEST2"},
 			},
 			seedDisks: []*gce.CloudDisk{
 				createZonalCloudDisk(name),
@@ -128,6 +129,30 @@ func TestCreateSnapshotArguments(t *testing.T) {
 			},
 			expErrCode: codes.InvalidArgument,
 		},
+		{
+			name: "invalid snapshot parameter key",
+			req: &csi.CreateSnapshotRequest{
+				Name:           name,
+				SourceVolumeId: testVolumeID,
+				Parameters:     map[string]string{"bad-key": ""},
+			},
+			seedDisks: []*gce.CloudDisk{
+				createZonalCloudDisk(name),
+			},
+			expErrCode: codes.InvalidArgument,
+		},
+		{
+			name: "invalid snapshot locations",
+			req: &csi.CreateSnapshotRequest{
+				Name:           name,
+				SourceVolumeId: testVolumeID,
+				Parameters:     map[string]string{common.ParameterKeyStorageLocations: "bad-region"},
+			},
+			seedDisks: []*gce.CloudDisk{
+				createZonalCloudDisk(name),
+			},
+			expErrCode: codes.InvalidArgument,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -848,7 +873,11 @@ func TestCreateVolumeWithVolumeSource(t *testing.T) {
 		}
 
 		if tc.snapshotOnCloud {
-			gceDriver.cs.CloudProvider.CreateSnapshot(context.Background(), tc.project, tc.volKey, name)
+			snapshotParams, err := common.ExtractAndDefaultSnapshotParameters(req.GetParameters())
+			if err != nil {
+				t.Errorf("Got error extracting snapshot parameters: %v", err)
+			}
+			gceDriver.cs.CloudProvider.CreateSnapshot(context.Background(), tc.project, tc.volKey, name, snapshotParams)
 		}
 		resp, err := gceDriver.cs.CreateVolume(context.Background(), req)
 		//check response
