fix bug where volume cloning topology requirements are ignored when chosing the location of the volume

Author: amacaskill

pkg/gce-pd-csi-driver/controller.go

@@ -20,6 +20,7 @@ import (
 	"math/rand"
 	"regexp"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/GoogleCloudPlatform/k8s-cloud-provider/pkg/cloud/meta"
@@ -32,6 +33,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/uuid"
 	"k8s.io/client-go/util/flowcontrol"
 	"k8s.io/klog/v2"
+	"k8s.io/utils/strings/slices"
 
 	"sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/common"
 	gce "sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/gce-cloud-provider/compute"
@@ -98,6 +100,13 @@ type workItem struct {
 	unpublishReq *csi.ControllerUnpublishVolumeRequest
 }
 
+// locationRequirements are additional location topology requirements that must be respected when creating a volume.
+type locationRequirements struct {
+	region          string
+	zone            string
+	replicationType string
+}
+
 var _ csi.ControllerServer = &GCEControllerServer{}
 
 const (
@@ -142,6 +151,40 @@ func isDiskReady(disk *gce.CloudDisk) (bool, error) {
 	return false, nil
 }
 
+// cloningLocationRequirements returns additional location requirements to be applied to the given create volume requests topology.
+// If the CreateVolumeRequest will use volume cloning, location requirements in compliance with the volume cloning limitations
+// will be returned: https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/volume-cloning#limitations.
+func cloningLocationRequirements(req *csi.CreateVolumeRequest, replicationType string) (*locationRequirements, error) {
+	if !useVolumeCloning(req) {
+		return nil, nil
+	}
+	// If we are using volume cloning, this will be set.
+	volSrc := req.VolumeContentSource.GetVolume()
+	volSrcVolID := volSrc.GetVolumeId()
+
+	_, sourceVolKey, err := common.VolumeIDToKey(volSrcVolID)
+	if err != nil {
+		return nil, err
+	}
+	if isZonalKey(sourceVolKey) {
+		region, err := common.GetRegionFromZones([]string{sourceVolKey.Zone})
+		if err != nil {
+			return nil, err
+		}
+		sourceVolKey.Region = region
+	}
+	return &locationRequirements{zone: sourceVolKey.Zone, region: sourceVolKey.Region, replicationType: replicationType}, nil
+}
+
+func isZonalKey(key *meta.Key) bool {
+	return key.Zone != "" && key.Region == ""
+}
+
+// useVolumeCloning returns true if the create volume request should be created with volume cloning.
+func useVolumeCloning(req *csi.CreateVolumeRequest) bool {
+	return req.VolumeContentSource != nil && req.VolumeContentSource.GetVolume() != nil
+}
+
 func (gceCS *GCEControllerServer) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest) (*csi.CreateVolumeResponse, error) {
 	var err error
 	// Validate arguments
@@ -177,12 +220,21 @@ func (gceCS *GCEControllerServer) CreateVolume(ctx context.Context, req *csi.Cre
 	if multiWriter {
 		gceAPIVersion = gce.GCEAPIVersionBeta
 	}
+
+	var locationTopReq *locationRequirements
+	if useVolumeCloning(req) {
+		locationTopReq, err = cloningLocationRequirements(req, params.ReplicationType)
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "failed to get location requirements: %v", err.Error())
+		}
+	}
+
 	// Determine the zone or zones+region of the disk
 	var zones []string
 	var volKey *meta.Key
 	switch params.ReplicationType {
 	case replicationTypeNone:
-		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 1)
+		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 1, locationTopReq)
 		if err != nil {
 			return nil, status.Errorf(codes.InvalidArgument, "CreateVolume failed to pick zones for disk: %v", err.Error())
 		}
@@ -192,7 +244,7 @@ func (gceCS *GCEControllerServer) CreateVolume(ctx context.Context, req *csi.Cre
 		volKey = meta.ZonalKey(name, zones[0])
 
 	case replicationTypeRegionalPD:
-		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 2)
+		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 2, locationTopReq)
 		if err != nil {
 			return nil, status.Errorf(codes.InvalidArgument, "CreateVolume failed to pick zones for disk: %v", err.Error())
 		}
@@ -1358,7 +1410,18 @@ func diskIsAttachedAndCompatible(deviceName string, instance *compute.Instance,
 	return false, nil
 }
 
-func pickZonesFromTopology(top *csi.TopologyRequirement, numZones int) ([]string, error) {
+// pickZonesInRegion will remove any zones that are not in the given region.
+func pickZonesInRegion(region string, zones []string) []string {
+	refinedZones := []string{}
+	for _, zone := range zones {
+		if strings.Contains(zone, region) {
+			refinedZones = append(refinedZones, zone)
+		}
+	}
+	return refinedZones
+}
+
+func pickZonesFromTopology(top *csi.TopologyRequirement, numZones int, locationTopReq *locationRequirements) ([]string, error) {
 	reqZones, err := getZonesFromTopology(top.GetRequisite())
 	if err != nil {
 		return nil, fmt.Errorf("could not get zones from requisite topology: %w", err)
@@ -1368,6 +1431,21 @@ func pickZonesFromTopology(top *csi.TopologyRequirement, numZones int) ([]string
 		return nil, fmt.Errorf("could not get zones from preferred topology: %w", err)
 	}
 
+	if locationTopReq != nil {
+		switch locationTopReq.replicationType {
+		// For zonal -> zonal cloning, the source disk zone must match the destination disk zone.
+		case replicationTypeNone:
+			if slices.Contains(prefZones, locationTopReq.zone) || slices.Contains(reqZones, locationTopReq.zone) {
+				return []string{locationTopReq.zone}, nil
+			}
+			return nil, fmt.Errorf("failed to find zone in preferred zones %v or requisite zones %v that satisfies location requirements %v", prefZones, reqZones, locationTopReq)
+		// for zonal or regional -> regional disk cloning, the source disk region must match the destination disk region.
+		case replicationTypeRegionalPD:
+			prefZones = pickZonesInRegion(locationTopReq.region, prefZones)
+			reqZones = pickZonesInRegion(locationTopReq.region, reqZones)
+		}
+	}
+
 	if numZones <= len(prefZones) {
 		return prefZones[0:numZones], nil
 	} else {
@@ -1426,16 +1504,16 @@ func getZoneFromSegment(seg map[string]string) (string, error) {
 	return zone, nil
 }
 
-func pickZones(ctx context.Context, gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZones int) ([]string, error) {
+func pickZones(ctx context.Context, gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZones int, locationTopReq *locationRequirements) ([]string, error) {
 	var zones []string
 	var err error
 	if top != nil {
-		zones, err = pickZonesFromTopology(top, numZones)
+		zones, err = pickZonesFromTopology(top, numZones, locationTopReq)
 		if err != nil {
 			return nil, fmt.Errorf("failed to pick zones from topology: %w", err)
 		}
 	} else {
-		zones, err = getDefaultZonesInRegion(ctx, gceCS, []string{gceCS.CloudProvider.GetDefaultZone()}, numZones)
+		zones, err = getDefaultZonesInRegion(ctx, gceCS, []string{gceCS.CloudProvider.GetDefaultZone()}, numZones, locationTopReq)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get default %v zones in region: %w", numZones, err)
 		}
@@ -1445,16 +1523,35 @@ func pickZones(ctx context.Context, gceCS *GCEControllerServer, top *csi.Topolog
 	return zones, nil
 }
 
-func getDefaultZonesInRegion(ctx context.Context, gceCS *GCEControllerServer, existingZones []string, numZones int) ([]string, error) {
+func getDefaultZonesInRegion(ctx context.Context, gceCS *GCEControllerServer, existingZones []string, numZones int, locationTopReq *locationRequirements) ([]string, error) {
 	region, err := common.GetRegionFromZones(existingZones)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get region from zones: %w", err)
 	}
 	needToGet := numZones - len(existingZones)
+
 	totZones, err := gceCS.CloudProvider.ListZones(ctx, region)
 	if err != nil {
 		return nil, fmt.Errorf("failed to list zones from cloud provider: %w", err)
 	}
+
+	if locationTopReq != nil {
+		if region != locationTopReq.region {
+			return nil, fmt.Errorf("default region %s does not match source volume region %s", region, locationTopReq.region)
+		}
+		switch locationTopReq.replicationType {
+		// For zonal -> zonal cloning, the source disk zone must match the destination disk zone.
+		case replicationTypeNone:
+			if slices.Contains(totZones, locationTopReq.zone) {
+				return []string{locationTopReq.zone}, nil
+			}
+			return nil, fmt.Errorf("failed to find zone in %v that satisfies location requirements %v", totZones, locationTopReq)
+		// for zonal or regional -> regional disk cloning, the source disk region must match the destination disk region.
+		case replicationTypeRegionalPD:
+			pickZonesInRegion(locationTopReq.region, totZones)
+		}
+	}
+
 	remainingZones := sets.NewString(totZones...).Difference(sets.NewString(existingZones...))
 	l := remainingZones.List()
 	if len(l) < needToGet {