Add support for multi-zone provisioning

Author: pwschuurman

pkg/common/parameters.go

@@ -33,6 +33,7 @@ const (
 	ParameterKeyEnableConfidentialCompute     = "enable-confidential-storage"
 	ParameterKeyStoragePools                  = "storage-pools"
 	ParameterKeyResourceTags                  = "resource-tags"
+	ParameterKeyEnableMultiZoneProvisioning   = "enable-multi-zone-provisioning"
 
 	// Parameters for VolumeSnapshotClass
 	ParameterKeyStorageLocations = "storage-locations"
@@ -102,6 +103,9 @@ type DiskParameters struct {
 	// Values: {map[string]string}
 	// Default: ""
 	ResourceTags map[string]string
+	// Values: {bool}
+	// Default: false
+	MultiZoneProvisioning bool
 }
 
 // SnapshotParameters contains normalized and defaulted parameters for snapshots
@@ -121,11 +125,17 @@ type StoragePool struct {
 	ResourceName string
 }
 
+type ParameterProcessor struct {
+	DriverName         string
+	EnableStoragePools bool
+	EnableMultiZone    bool
+}
+
 // 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
 // parameters, any matching extraVolumeLabels will be overridden.
-func ExtractAndDefaultParameters(parameters map[string]string, driverName string, extraVolumeLabels map[string]string, enableStoragePools bool, extraTags map[string]string) (DiskParameters, error) {
+func (pp *ParameterProcessor) ExtractAndDefaultParameters(parameters map[string]string, extraVolumeLabels map[string]string, extraTags map[string]string) (DiskParameters, error) {
 	p := DiskParameters{
 		DiskType:             "pd-standard",           // Default
 		ReplicationType:      replicationTypeNone,     // Default
@@ -210,24 +220,37 @@ func ExtractAndDefaultParameters(parameters map[string]string, driverName string
 
 			p.EnableConfidentialCompute = paramEnableConfidentialCompute
 		case ParameterKeyStoragePools:
-			if !enableStoragePools {
+			if !pp.EnableStoragePools {
 				return p, fmt.Errorf("parameters contains invalid option %q", ParameterKeyStoragePools)
 			}
 			storagePools, err := ParseStoragePools(v)
 			if err != nil {
-				return p, fmt.Errorf("parameters contain invalid value for %s parameter: %w", ParameterKeyStoragePools, err)
+				return p, fmt.Errorf("parameters contains invalid value for %s parameter %q: %w", ParameterKeyStoragePools, v, err)
 			}
 			p.StoragePools = storagePools
 		case ParameterKeyResourceTags:
 			if err := extractResourceTagsParameter(v, p.ResourceTags); err != nil {
 				return p, err
 			}
+		case ParameterKeyEnableMultiZoneProvisioning:
+			if !pp.EnableMultiZone {
+				return p, fmt.Errorf("parameters contains invalid option %q", ParameterKeyEnableMultiZoneProvisioning)
+			}
+			paramEnableMultiZoneProvisioning, err := ConvertStringToBool(v)
+			if err != nil {
+				return p, fmt.Errorf("parameters contain invalid value for %s parameter: %w", ParameterKeyEnableMultiZoneProvisioning, err)
+			}
+
+			p.MultiZoneProvisioning = paramEnableMultiZoneProvisioning
+			if paramEnableMultiZoneProvisioning {
+				p.Labels[MultiZoneLabel] = "true"
+			}
 		default:
 			return p, fmt.Errorf("parameters contains invalid option %q", k)
 		}
 	}
 	if len(p.Tags) > 0 {
-		p.Tags[tagKeyCreatedBy] = driverName
+		p.Tags[tagKeyCreatedBy] = pp.DriverName
 	}
 	return p, nil
 }

pkg/common/parameters_test.go

@@ -29,6 +29,7 @@ func TestExtractAndDefaultParameters(t *testing.T) {
 		parameters         map[string]string
 		labels             map[string]string
 		enableStoragePools bool
+		enableMultiZone    bool
 		extraTags          map[string]string
 		expectParams       DiskParameters
 		expectErr          bool
@@ -350,19 +351,61 @@ func TestExtractAndDefaultParameters(t *testing.T) {
 			labels:             map[string]string{},
 			expectErr:          true,
 		},
+		{
+			name:            "multi-zone-enable parameters, multi-zone label is set, multi-zone feature enabled",
+			parameters:      map[string]string{ParameterKeyType: "hyperdisk-ml", ParameterKeyEnableMultiZoneProvisioning: "true"},
+			labels:          map[string]string{MultiZoneLabel: "true"},
+			enableMultiZone: true,
+			expectParams: DiskParameters{
+				DiskType:              "hyperdisk-ml",
+				ReplicationType:       "none",
+				Tags:                  map[string]string{},
+				Labels:                map[string]string{MultiZoneLabel: "true"},
+				ResourceTags:          map[string]string{},
+				MultiZoneProvisioning: true,
+			},
+		},
+		{
+			name:            "multi-zone-enable parameters, multi-zone label is false, multi-zone feature enabled",
+			parameters:      map[string]string{ParameterKeyType: "hyperdisk-ml", ParameterKeyEnableMultiZoneProvisioning: "false"},
+			enableMultiZone: true,
+			expectParams: DiskParameters{
+				DiskType:        "hyperdisk-ml",
+				ReplicationType: "none",
+				Tags:            map[string]string{},
+				ResourceTags:    map[string]string{},
+				Labels:          map[string]string{},
+			},
+		},
+		{
+			name:            "multi-zone-enable parameters, invalid value, multi-zone feature enabled",
+			parameters:      map[string]string{ParameterKeyType: "hyperdisk-ml", ParameterKeyEnableMultiZoneProvisioning: "unknown"},
+			enableMultiZone: true,
+			expectErr:       true,
+		},
+		{
+			name:       "multi-zone-enable parameters, multi-zone label is set, multi-zone feature disabled",
+			parameters: map[string]string{ParameterKeyType: "hyperdisk-ml", ParameterKeyEnableMultiZoneProvisioning: "true"},
+			expectErr:  true,
+		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			p, err := ExtractAndDefaultParameters(tc.parameters, "testDriver", tc.labels, tc.enableStoragePools, tc.extraTags)
+			pp := ParameterProcessor{
+				DriverName:         "testDriver",
+				EnableStoragePools: tc.enableStoragePools,
+				EnableMultiZone:    tc.enableMultiZone,
+			}
+			p, err := pp.ExtractAndDefaultParameters(tc.parameters, tc.labels, tc.extraTags)
 			if gotErr := err != nil; gotErr != tc.expectErr {
 				t.Fatalf("ExtractAndDefaultParameters(%+v) = %v; expectedErr: %v", tc.parameters, err, tc.expectErr)
 			}
 			if err != nil {
 				return
 			}
 
-			if diff := cmp.Diff(p, tc.expectParams); diff != "" {
+			if diff := cmp.Diff(tc.expectParams, p); diff != "" {
 				t.Errorf("ExtractAndDefaultParameters(%+v): -want, +got \n%s", tc.parameters, diff)
 			}
 		})

pkg/common/utils.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
@@ -73,6 +74,10 @@ const (
 	// Full or partial URL of the machine type resource, in the format:
 	//   zones/zone/machineTypes/machine-type
 	machineTypePattern = "zones/[^/]+/machineTypes/([^/]+)$"
+
+	// Full or partial URL of the zone resource, in the format:
+	//   projects/{project}/zones/{zone}
+	zoneURIPattern = "projects/[^/]+/zones/([^/]+)$"
 )
 
 var (
@@ -85,6 +90,8 @@ var (
 
 	storagePoolFieldsRegex = regexp.MustCompile(`^projects/([^/]+)/zones/([^/]+)/storagePools/([^/]+)$`)
 
+	zoneURIRegex = regexp.MustCompile(zoneURIPattern)
+
 	// userErrorCodeMap tells how API error types are translated to error codes.
 	userErrorCodeMap = map[int]codes.Code{
 		http.StatusForbidden:       codes.PermissionDenied,
@@ -97,6 +104,8 @@ var (
 	regexParent = regexp.MustCompile(`(^[1-9][0-9]{0,31}$)|(^[a-z][a-z0-9-]{4,28}[a-z0-9]$)`)
 	regexKey    = regexp.MustCompile(`^[a-zA-Z0-9]([0-9A-Za-z_.-]{0,61}[a-zA-Z0-9])?$`)
 	regexValue  = regexp.MustCompile(`^[a-zA-Z0-9]([0-9A-Za-z_.@%=+:,*#&()\[\]{}\-\s]{0,61}[a-zA-Z0-9])?$`)
+
+	csiRetryableErrorCodes = []codes.Code{codes.Canceled, codes.DeadlineExceeded, codes.Unavailable, codes.Aborted, codes.ResourceExhausted}
 )
 
 func BytesToGbRoundDown(bytes int64) int64 {
@@ -545,9 +554,37 @@ func isGoogleAPIError(err error) (codes.Code, error) {
 	return codes.Unknown, fmt.Errorf("googleapi.Error %w does not map to any known errors", err)
 }
 
-func LoggedError(msg string, err error) error {
+func loggedErrorForCode(msg string, code codes.Code, err error) error {
 	klog.Errorf(msg+"%v", err.Error())
-	return status.Errorf(CodeForError(err), msg+"%v", err.Error())
+	return status.Errorf(code, msg+"%v", err.Error())
+}
+
+func LoggedError(msg string, err error) error {
+	return loggedErrorForCode(msg, CodeForError(err), err)
+}
+
+// NewCombinedError tries to return an appropriate wrapped error that captures
+// useful information as an error code
+// If there are multiple errors, it extracts the first "retryable" error
+// as interpreted by the CSI sidecar.
+func NewCombinedError(msg string, errs []error) error {
+	// If there is only one error, return it as the single error code
+	if len(errs) == 1 {
+		LoggedError(msg, errs[0])
+	}
+
+	for _, err := range errs {
+		code := CodeForError(err)
+		if slices.Contains(csiRetryableErrorCodes, code) {
+			// Return this as a TemporaryError to lock-in the retryable code
+			// This will invoke the "existing" error code check in CodeForError
+			return NewTemporaryError(code, fmt.Errorf("%s: %w", msg, err))
+		}
+	}
+
+	// None of these error codes were retryable. Just return a combined error
+	// The first matching error (based on our CodeForError) logic will be returned.
+	return LoggedError(msg, errors.Join(errs...))
 }
 
 func isValidDiskEncryptionKmsKey(DiskEncryptionKmsKey string) bool {
@@ -556,6 +593,14 @@ func isValidDiskEncryptionKmsKey(DiskEncryptionKmsKey string) bool {
 	return kmsKeyPattern.MatchString(DiskEncryptionKmsKey)
 }
 
+func ParseZoneFromURI(zoneURI string) (string, error) {
+	zoneMatch := zoneURIRegex.FindStringSubmatch(zoneURI)
+	if zoneMatch == nil {
+		return "", fmt.Errorf("failed to parse zone URI. Expected projects/{project}/zones/{zone}. Got: %s", zoneURI)
+	}
+	return zoneMatch[1], nil
+}
+
 // ParseStoragePools returns an error if none of the given storagePools
 // (delimited by a comma) are in the format
 // projects/project/zones/zone/storagePools/storagePool.

pkg/common/utils_test.go

@@ -1648,3 +1648,89 @@ func TestUnorderedSlicesEqual(t *testing.T) {
 		})
 	}
 }
+
+func TestParseZoneFromURI(t *testing.T) {
+	testcases := []struct {
+		name      string
+		zoneURI   string
+		wantZone  string
+		expectErr bool
+	}{
+		{
+			name:     "ParseZoneFromURI_FullURI",
+			zoneURI:  "https://www.googleapis.com/compute/v1/projects/psch-gke-dev/zones/us-east4-a",
+			wantZone: "us-east4-a",
+		},
+		{
+			name:     "ParseZoneFromURI_ProjectZoneString",
+			zoneURI:  "projects/psch-gke-dev/zones/us-east4-a",
+			wantZone: "us-east4-a",
+		},
+		{
+			name:      "ParseZoneFromURI_Malformed",
+			zoneURI:   "projects/psch-gke-dev/regions/us-east4",
+			expectErr: true,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			gotZone, err := ParseZoneFromURI(tc.zoneURI)
+			if err != nil && !tc.expectErr {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if err == nil && tc.expectErr {
+				t.Fatalf("Expected err, but none was returned. Zone result: %v", gotZone)
+			}
+			if gotZone != tc.wantZone {
+				t.Errorf("ParseZoneFromURI(%v): got %v, want %v", tc.zoneURI, gotZone, tc.wantZone)
+			}
+		})
+	}
+}
+
+func TestNewCombinedError(t *testing.T) {
+	testcases := []struct {
+		name     string
+		errors   []error
+		wantCode codes.Code
+	}{
+		{
+			name:     "single generic error",
+			errors:   []error{fmt.Errorf("my internal error")},
+			wantCode: codes.Internal,
+		},
+		{
+			name:     "single retryable error",
+			errors:   []error{&googleapi.Error{Code: http.StatusTooManyRequests, Message: "Resource Exhausted"}},
+			wantCode: codes.ResourceExhausted,
+		},
+		{
+			name:     "multi generic error",
+			errors:   []error{fmt.Errorf("my internal error"), fmt.Errorf("my other internal error")},
+			wantCode: codes.Internal,
+		},
+		{
+			name:     "multi retryable error",
+			errors:   []error{fmt.Errorf("my internal error"), &googleapi.Error{Code: http.StatusTooManyRequests, Message: "Resource Exhausted"}},
+			wantCode: codes.ResourceExhausted,
+		},
+		{
+			name:     "multi retryable error",
+			errors:   []error{fmt.Errorf("my internal error"), &googleapi.Error{Code: http.StatusGatewayTimeout, Message: "connection reset by peer"}, fmt.Errorf("my other internal error")},
+			wantCode: codes.Unavailable,
+		},
+		{
+			name:     "multi retryable error",
+			errors:   []error{fmt.Errorf("The disk resource is already being used"), &googleapi.Error{Code: http.StatusGatewayTimeout, Message: "connection reset by peer"}},
+			wantCode: codes.Unavailable,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			gotCode := CodeForError(NewCombinedError("message", tc.errors))
+			if gotCode != tc.wantCode {
+				t.Errorf("NewCombinedError(%v): got %v, want %v", tc.errors, gotCode, tc.wantCode)
+			}
+		})
+	}
+}

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

@@ -256,3 +256,14 @@ func (d *CloudDisk) GetLabels() map[string]string {
 		return nil
 	}
 }
+
+func (d *CloudDisk) GetAccessMode() string {
+	switch {
+	case d.disk != nil:
+		return d.disk.AccessMode
+	case d.betaDisk != nil:
+		return d.betaDisk.AccessMode
+	default:
+		return ""
+	}
+}

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

@@ -148,3 +148,50 @@ func TestGetLabels(t *testing.T) {
 		}
 	}
 }
+
+func TestGetAccessMode(t *testing.T) {
+	testCases := []struct {
+		name           string
+		cloudDisk      *CloudDisk
+		wantAccessMode string
+	}{
+		{
+			name: "v1 disk accessMode",
+			cloudDisk: &CloudDisk{
+				disk: &computev1.Disk{
+					AccessMode: "READ_WRITE_SINGLE",
+				},
+			},
+			wantAccessMode: "READ_WRITE_SINGLE",
+		},
+		{
+			name: "beta disk accessMode",
+			cloudDisk: &CloudDisk{
+				betaDisk: &computebeta.Disk{
+					AccessMode: "READ_ONLY_MANY",
+				},
+			},
+			wantAccessMode: "READ_ONLY_MANY",
+		},
+		{
+			name: "unset disk accessMode",
+			cloudDisk: &CloudDisk{
+				betaDisk: &computebeta.Disk{},
+			},
+			wantAccessMode: "",
+		},
+		{
+			name:           "unset disk",
+			cloudDisk:      &CloudDisk{},
+			wantAccessMode: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Logf("Running test: %v", tc.name)
+		gotAccessMode := tc.cloudDisk.GetAccessMode()
+		if gotAccessMode != tc.wantAccessMode {
+			t.Errorf("GetAccessMode() got %v, want %v", gotAccessMode, tc.wantAccessMode)
+		}
+	}
+}