Update blockstorage/zone pkg to get Regional persistent disk Availabi…

…lity Zone (#5391)

* Add code to get Regional persistent disk AZone

* Avoid errors

* Address reviews discussed offline

* Address errors

* address review comments

* Avoid errors

* Add comments

* Address review suggestions

pkg/blockstorage/awsebs/awsebs.go

@@ -413,13 +413,15 @@ func (s *ebsStorage) VolumeCreateFromSnapshot(ctx context.Context, snapshot bloc
 		return nil, errors.Errorf("Required volume fields not available, volumeType: %s, Az: %s, VolumeTags: %v", snapshot.Volume.VolumeType, snapshot.Volume.Az, snapshot.Volume.Tags)
 	}
 
-	z, err := zone.FromSourceRegionZone(ctx, s, snapshot.Region, snapshot.Volume.Az)
+	zones, err := zone.FromSourceRegionZone(ctx, s, snapshot.Region, snapshot.Volume.Az)
 	if err != nil {
 		return nil, err
 	}
-
+	if len(zones) != 1 {
+		return nil, errors.Errorf("Length of zone slice should be 1, got %d", len(zones))
+	}
 	cvi := &ec2.CreateVolumeInput{
-		AvailabilityZone: aws.String(z),
+		AvailabilityZone: aws.String(zones[0]),
 		SnapshotId:       aws.String(snapshot.ID),
 		VolumeType:       aws.String(string(snapshot.Volume.VolumeType)),
 	}

pkg/blockstorage/awsebs/zone_test.go

@@ -41,8 +41,7 @@ func (s ZoneSuite) TestZoneWithUnknownNodeZones(c *C) {
 		},
 	} {
 		var t = &ebsTest{}
-		z, err := zone.WithUnknownNodeZones(ctx, t, tc.region, tc.in)
-		c.Assert(err, IsNil)
+		z := zone.WithUnknownNodeZones(ctx, t, tc.region, tc.in, make(map[string]struct{}))
 		c.Assert(z, Not(Equals), "")
 		if tc.out != "" {
 			c.Assert(z, Equals, tc.out)

pkg/blockstorage/gcepd/gcepd.go

@@ -318,10 +318,14 @@ func (s *gpdStorage) VolumeCreateFromSnapshot(ctx context.Context, snapshot bloc
 			tags[tag.Key] = tag.Value
 		}
 	}
-	z, err := zone.FromSourceRegionZone(ctx, s, snapshot.Region, snapshot.Volume.Az)
+	zones, err := zone.FromSourceRegionZone(ctx, s, snapshot.Region, snapshot.Volume.Az)
 	if err != nil {
 		return nil, err
 	}
+	if len(zones) != 1 {
+		return nil, errors.Errorf("Length of zone slice should be 1, got %d", len(zones))
+	}
+
 	createDisk := &compute.Disk{
 		Name:           fmt.Sprintf(volumeNameFmt, uuid.NewV1().String()),
 		SizeGb:         snapshot.Volume.Size,
@@ -330,15 +334,15 @@ func (s *gpdStorage) VolumeCreateFromSnapshot(ctx context.Context, snapshot bloc
 		SourceSnapshot: snap.SelfLink,
 	}
 
-	resp, err := s.service.Disks.Insert(s.project, z, createDisk).Context(ctx).Do()
+	resp, err := s.service.Disks.Insert(s.project, zones[0], createDisk).Context(ctx).Do()
 	if err != nil {
 		return nil, err
 	}
-	if err := s.waitOnOperation(ctx, resp, z); err != nil {
+	if err := s.waitOnOperation(ctx, resp, zones[0]); err != nil {
 		return nil, err
 	}
 
-	return s.VolumeGet(ctx, createDisk.Name, z)
+	return s.VolumeGet(ctx, createDisk.Name, zones[0])
 }
 
 func (s *gpdStorage) SetTags(ctx context.Context, resource interface{}, tags map[string]string) error {

pkg/blockstorage/gcepd/zone_test.go

@@ -41,8 +41,7 @@ func (s ZoneSuite) TestZoneWithUnknownNodeZones(c *C) {
 		},
 	} {
 		var t = &gcpTest{}
-		z, err := zone.WithUnknownNodeZones(ctx, t, tc.region, tc.in)
-		c.Assert(err, IsNil)
+		z := zone.WithUnknownNodeZones(ctx, t, tc.region, tc.in, make(map[string]struct{}))
 		c.Assert(z, Not(Equals), "")
 		if tc.out != "" {
 			c.Assert(z, Equals, tc.out)

pkg/blockstorage/zone/zone.go

@@ -21,26 +21,53 @@ type (
 	}
 )
 
-// FromSourceRegionZone gets the zones from the given region and sourceZome
-func FromSourceRegionZone(ctx context.Context, m Mapper, region string, sourceZone string) (string, error) {
+// FromSourceRegionZone gets the zones from the given region and sourceZones
+// It will return a minimum of 0 and a maximum of zones equal to the length of sourceZones.
+// Depending on the length of the slice returned, the blockstorage providers will decide if
+// a regional volume or a zonal volume should be created.
+func FromSourceRegionZone(ctx context.Context, m Mapper, region string, sourceZones ...string) ([]string, error) {
+	newZones := make(map[string]struct{})
 	cli, err := kube.NewClient()
 	if err == nil {
 		nzs, rs, errzr := nodeZonesAndRegion(ctx, cli)
 		if err != nil {
 			log.Errorf("Ignoring error getting Node availability zones. Error: %+v", errzr)
 		}
 		if len(nzs) != 0 {
-			var z string
-			z, err = getZoneFromKnownNodeZones(ctx, sourceZone, nzs)
-			if err == nil && isZoneValid(ctx, m, z, rs) {
-				return z, nil
+			for _, zone := range sourceZones {
+				var z string
+				// getZoneFromKnownNodeZones() func makes sure that all zones
+				// added to newZones[] are unique and not repeated.
+				z, err = getZoneFromKnownNodeZones(ctx, zone, nzs, newZones)
+				if err == nil && isZoneValid(ctx, m, z, rs) {
+					newZones[z] = struct{}{}
+				}
+				if err != nil {
+					log.Errorf("Ignoring error getting Zone from KnownNodeZones. Error: %+v", err)
+				}
 			}
-			if err != nil {
-				log.Errorf("Ignoring error getting Zone from KnownNodeZones. Error: %+v", err)
+		}
+	}
+	if len(newZones) == 0 {
+		for _, zone := range sourceZones {
+			// WithUnknownNodeZones() func makes sure that all zones
+			// added to newZones[] are unique and not repeated.
+			newZone := WithUnknownNodeZones(ctx, m, region, zone, newZones)
+			if newZone != "" {
+				newZones[newZone] = struct{}{}
 			}
 		}
 	}
-	return WithUnknownNodeZones(ctx, m, region, sourceZone)
+
+	if len(newZones) == 0 {
+		return nil, errors.Errorf("Could not get zone for region %s and sourceZones %s", region, sourceZones)
+	}
+
+	var zones []string
+	for z := range newZones {
+		zones = append(zones, z)
+	}
+	return zones, nil
 }
 
 func isZoneValid(ctx context.Context, m Mapper, zone, region string) bool {
@@ -55,40 +82,59 @@ func isZoneValid(ctx context.Context, m Mapper, zone, region string) bool {
 }
 
 // WithUnknownNodeZones get the zone list  for the region
-func WithUnknownNodeZones(ctx context.Context, m Mapper, region string, sourceZone string) (string, error) {
+func WithUnknownNodeZones(ctx context.Context, m Mapper, region string, sourceZone string, newZones map[string]struct{}) string {
 	// We could not the zones of the nodes, so we return an arbitrary one.
 	zs, err := m.FromRegion(ctx, region)
 	if err != nil || len(zs) == 0 {
 		// If all else fails, we return the original AZ.
 		log.Errorf("Using original AZ. region: %s, Error: %+v", region, err)
-		return sourceZone, nil
+		return sourceZone
 	}
+
 	// We look for a zone with the same suffix.
 	for _, z := range zs {
 		if getZoneSuffixesMatch(z, sourceZone) {
-			return z, nil
+			// check if zone z is already added to newZones
+			if _, ok := newZones[z]; ok {
+				continue
+			}
+			return z
 		}
 	}
+
 	// We return an arbitrary zone in the region.
-	return zs[0], nil
+	for i := range zs {
+		// check if zone zs[i] is already added to newZones
+		if _, ok := newZones[zs[i]]; ok {
+			continue
+		}
+		return zs[i]
+	}
+
+	return ""
 }
 
-func getZoneFromKnownNodeZones(ctx context.Context, sourceZone string, nzs map[string]struct{}) (string, error) {
+func getZoneFromKnownNodeZones(ctx context.Context, sourceZone string, nzs map[string]struct{}, newZones map[string]struct{}) (string, error) {
 	// If the original zone is available, we return that one.
 	if _, ok := nzs[sourceZone]; ok {
 		return sourceZone, nil
 	}
+
 	// If there's an available zone with the zone suffix, we use that one.
 	for nz := range nzs {
 		if getZoneSuffixesMatch(nz, sourceZone) {
+			// check if zone nz is already added to newZones
+			if _, ok := newZones[nz]; ok {
+				continue
+			}
 			return nz, nil
 		}
 	}
 	// If any nodes are available, return an arbitrary one.
-	return consistentZone(sourceZone, nzs)
+	return consistentZone(sourceZone, nzs, newZones)
 }
 
-func consistentZone(sourceZone string, nzs map[string]struct{}) (string, error) {
+func consistentZone(sourceZone string, nzs map[string]struct{}, newZones map[string]struct{}) (string, error) {
 	if len(nzs) == 0 {
 		return "", errors.New("could not restore volume: no zone found")
 	}
@@ -104,6 +150,11 @@ func consistentZone(sourceZone string, nzs map[string]struct{}) (string, error)
 		return "", errors.Errorf("failed to hash source zone %s: %s", sourceZone, err.Error())
 	}
 	i := int(h.Sum32()) % len(nzs)
+
+	// check if zone s[i] is already added to newZones
+	if _, ok := newZones[s[i]]; ok {
+		return "", nil
+	}
 	return s[i], nil
 }
 

pkg/blockstorage/zone/zone_test.go

@@ -51,7 +51,7 @@ func (s ZoneSuite) TestConsistentZone(c *C) {
 			out: "zone1",
 		},
 	} {
-		out, err := consistentZone(tc.sourceZone, tc.nzs)
+		out, err := consistentZone(tc.sourceZone, tc.nzs, make(map[string]struct{}))
 		c.Assert(err, IsNil)
 		c.Assert(out, Equals, tc.out)
 	}