WIP: Lookup available zones for GCP

Author: angelcerveraroldan

mantle/cmd/kola/options.go

@@ -121,7 +121,7 @@ func init() {
 	// gcp-specific options
 	sv(&kola.GCPOptions.Image, "gcp-image", "", "GCP image, full api endpoints names are accepted if resource is in a different project")
 	sv(&kola.GCPOptions.Project, "gcp-project", "fedora-coreos-devel", "GCP project name")
-	sv(&kola.GCPOptions.Zone, "gcp-zone", "us-central1-a", "GCP zone name")
+	sv(&kola.GCPOptions.Zone, "gcp-zone", "us-central1-a", "Preferred GCP zone name, if the resources to this zone are depleated, we will fallback to another zone in the same region")
 	sv(&kola.GCPOptions.MachineType, "gcp-machinetype", "", "GCP machine type")
 	sv(&kola.GCPOptions.DiskType, "gcp-disktype", "", "GCP disk type (default pd-ssd)")
 	sv(&kola.GCPOptions.Network, "gcp-network", "default", "GCP network")

mantle/platform/api/gcloud/api.go

@@ -17,10 +17,13 @@ package gcloud
 
 import (
 	"context"
-	"google.golang.org/api/option"
+	"fmt"
 	"net/http"
+	"regexp"
 	"time"
 
+	"google.golang.org/api/option"
+
 	"github.com/coreos/pkg/capnslog"
 	"google.golang.org/api/compute/v1"
 
@@ -50,6 +53,68 @@ type API struct {
 	client  *http.Client
 	compute *compute.Service
 	options *Options
+	zones   []string
+}
+
+// This regex shuld match all standard (non-AI) zones
+// See: https://docs.cloud.google.com/compute/docs/regions-zones
+var standardZoneRegexp = regexp.MustCompile(`^([a-z]+-[a-z]+\d+)-[a-z]$`)
+
+// zones are in the form "us-central1-a" and the region would be "us-central1"
+// See: https://docs.cloud.google.com/compute/docs/regions-zones
+func extractRegionFromZone(zone string) (string, error) {
+	matches := standardZoneRegexp.FindStringSubmatch(zone)
+	if matches == nil {
+		return "", fmt.Errorf("zone %q does not match expected format {region}-{letter}", zone)
+	}
+	return matches[1], nil
+}
+
+func getAvailableZones(computeService *compute.Service, opts *Options) ([]string, error) {
+	if opts.MachineType == "" {
+		return []string{opts.Zone}, nil
+	}
+
+	list, err := computeService.MachineTypes.AggregatedList(opts.Project).
+		Filter("name=" + opts.MachineType).Do()
+
+	if err != nil {
+		return nil, err
+	}
+
+	targetRegion, err := extractRegionFromZone(opts.Zone)
+	if err != nil {
+		return nil, fmt.Errorf("could not extract region from zone %q: %w", opts.Zone, err)
+	}
+
+	zones := []string{}
+	for _, scopedList := range list.Items {
+		// There should be either 1 or 0 MachineTypes
+		// 0 if this zone does not have the required machine type
+		// 1 if this zone does have the required machine type
+		if len(scopedList.MachineTypes) == 0 {
+			continue
+		}
+		if len(scopedList.MachineTypes) > 1 {
+			plog.Warningf("Unexpected: got %d machine types for filter name=%s", len(scopedList.MachineTypes), opts.MachineType)
+			continue
+		}
+		zone := scopedList.MachineTypes[0].Zone
+		if region, err := extractRegionFromZone(zone); err == nil && region == targetRegion {
+			// If the preferred zone can be used, it should be the first zone that we use,
+			// so we will make add it to the start of the list, rather than the end.
+			if zone == opts.Zone {
+				zones = append([]string{zone}, zones...)
+			} else {
+				zones = append(zones, zone)
+			}
+		}
+	}
+
+	if len(zones) == 0 {
+		return zones, fmt.Errorf("no zones in region %s for machine type %s were found", targetRegion, opts.MachineType)
+	}
+	return zones, nil
 }
 
 func New(opts *Options) (*API, error) {
@@ -83,6 +148,12 @@ func New(opts *Options) (*API, error) {
 		return nil, err
 	}
 
+	zones, err := getAvailableZones(computeService, opts)
+	if err != nil {
+		plog.Warningf("Failed to discover available zones: %v. Falling back to preferred zone (%s) only.", err, opts.Zone)
+		zones = []string{opts.Zone}
+	}
+
 	if opts.ServiceAcct == "" {
 		proj, err := computeService.Projects.Get(opts.Project).Do()
 		if err != nil {
@@ -95,6 +166,7 @@ func New(opts *Options) (*API, error) {
 		client:  client,
 		compute: computeService,
 		options: opts,
+		zones:   zones,
 	}
 
 	return api, nil