Support AZs list (#8)

* add availability-zones to perform an intersection of availabile instance types across multiple zones

* refactor location validation

* refactor location count increment

Author: bwagner5

README.md

@@ -130,11 +130,12 @@ Usage:
   ec2-instance-selector [flags]
 
 Examples:
-ec2-instance-selector --vcpus 4 --region us-east-2 --availability-zone us-east-2b
+ec2-instance-selector --vcpus 4 --region us-east-2 --availability-zones us-east-2b
 ec2-instance-selector --memory-min 4096 --memory-max 8192 --vcpus-min 4 --vcpus-max 8 --region us-east-2
 
 Filter Flags:
-  -z, --availability-zone string          Availability zone or zone id to check only EC2 capacity offered in a specific AZ
+      --availability-zone string          [DEPRECATED] use --availability-zones instead
+  -z, --availability-zones strings        Availability zones or zone ids to check EC2 capacity offered in specific AZs
       --baremetal                         Bare Metal instance types (.metal instances)
   -b, --burst-support                     Burstable instance types
   -a, --cpu-architecture string           CPU architecture [x86_64, i386, or arm64]

cmd/main.go

@@ -56,6 +56,7 @@ const (
 	burstSupport           = "burst-support"
 	hypervisor             = "hypervisor"
 	availabilityZone       = "availability-zone"
+	availabilityZones      = "availability-zones"
 	currentGeneration      = "current-generation"
 	networkInterfaces      = "network-interfaces"
 	networkPerformance     = "network-performance"
@@ -90,7 +91,7 @@ func main() {
 	longUsage := binName + ` is a CLI tool to filter EC2 instance types based on resource criteria. 
 Filtering allows you to select all the instance types that match your application requirements.
 Full docs can be found at github.com/aws/amazon-` + binName
-	examples := fmt.Sprintf(`%s --vcpus 4 --region us-east-2 --availability-zone us-east-2b
+	examples := fmt.Sprintf(`%s --vcpus 4 --region us-east-2 --availability-zones us-east-2b
 %s --memory-min 4096 --memory-max 8192 --vcpus-min 4 --vcpus-max 8 --region us-east-2`, binName, binName)
 
 	runFunc := func(cmd *cobra.Command, args []string) {}
@@ -120,7 +121,8 @@ Full docs can be found at github.com/aws/amazon-` + binName
 	cli.BoolFlag(fpgaSupport, cli.StringMe("f"), nil, "FPGA instance types")
 	cli.BoolFlag(burstSupport, cli.StringMe("b"), nil, "Burstable instance types")
 	cli.StringFlag(hypervisor, nil, nil, "Hypervisor: [xen or nitro]", nil)
-	cli.StringFlag(availabilityZone, cli.StringMe("z"), nil, "Availability zone or zone id to check only EC2 capacity offered in a specific AZ", nil)
+	cli.StringFlag(availabilityZone, nil, nil, "[DEPRECATED] use --availability-zones instead", nil)
+	cli.StringSliceFlag(availabilityZones, cli.StringMe("z"), nil, "Availability zones or zone ids to check EC2 capacity offered in specific AZs")
 	cli.BoolFlag(currentGeneration, nil, nil, "Current generation instance types (explicitly set this to false to not return current generation instance types)")
 	cli.IntMinMaxRangeFlags(networkInterfaces, nil, nil, "Number of network interfaces (ENIs) that can be attached to the instance")
 	cli.IntMinMaxRangeFlags(networkPerformance, nil, nil, "Bandwidth in Gib/s of network performance (Example: 100)")
@@ -164,6 +166,15 @@ Full docs can be found at github.com/aws/amazon-` + binName
 		sessOpts.Profile = *cli.StringMe(flags[profile])
 	}
 
+	if flags[availabilityZone] != nil {
+		log.Printf("You are using a deprecated flag --%s which will be removed in future versions, switch to --%s to avoid issues.\n", availabilityZone, availabilityZones)
+		if flags[availabilityZones] != nil {
+			flags[availabilityZones] = append(*cli.StringSliceMe(flags[availabilityZones]), *cli.StringMe(flags[availabilityZone]))
+		} else {
+			flags[availabilityZones] = []string{*cli.StringMe(flags[availabilityZone])}
+		}
+	}
+
 	sess := session.Must(session.NewSessionWithOptions(sessOpts))
 
 	instanceSelector := selector.New(sess)
@@ -185,7 +196,7 @@ Full docs can be found at github.com/aws/amazon-` + binName
 		Fpga:                   cli.BoolMe(flags[fpgaSupport]),
 		Burstable:              cli.BoolMe(flags[burstSupport]),
 		Region:                 cli.StringMe(flags[region]),
-		AvailabilityZone:       cli.StringMe(flags[availabilityZone]),
+		AvailabilityZones:      cli.StringSliceMe(flags[availabilityZones]),
 		CurrentGeneration:      cli.BoolMe(flags[currentGeneration]),
 		MaxResults:             cli.IntMe(flags[maxResults]),
 		NetworkInterfaces:      cli.IntRangeMe(flags[networkInterfaces]),

pkg/selector/selector.go

@@ -17,7 +17,6 @@ package selector
 import (
 	"fmt"
 	"reflect"
-	"regexp"
 	"sort"
 	"strings"
 
@@ -34,12 +33,6 @@ var (
 )
 
 const (
-	// regionNameRegex Matches strings like: us-east-1 or us-east-2
-	regionNameRegex = `^[a-z]{2,3}\-([a-z]{1,10}\-)?[a-z]{1,10}\-[1-9]`
-	// zoneIDRegex Matches strings like: use1-az1 or use2-az3
-	zoneIDRegex = `^[a-z]{3}[1-9]{1}\-az[1-9]$`
-	// zoneNameRegex Matches strings like: us-east-1a or us-east-2c
-	zoneNameRegex          = `^[a-z]{2,3}\-([a-z]{1,10}\-)?[a-z]{1,10}\-[1-9][a-z]$`
 	locationFilterKey      = "location"
 	zoneIDLocationType     = "availability-zone-id"
 	zoneNameLocationType   = "availability-zone"
@@ -175,13 +168,23 @@ func (itf Selector) rawFilter(filters Filters) ([]*ec2.InstanceTypeInfo, error)
 	if err != nil {
 		return nil, err
 	}
-	var location string
+	var locations []string
+
+	// Support the deprecated singular availabilityZone filter in favor of the plural
 	if filters.AvailabilityZone != nil {
-		location = *filters.AvailabilityZone
+		if filters.AvailabilityZones != nil {
+			*filters.AvailabilityZones = append(*filters.AvailabilityZones, *filters.AvailabilityZone)
+		} else {
+			filters.AvailabilityZones = &[]string{*filters.AvailabilityZone}
+		}
+	}
+
+	if filters.AvailabilityZones != nil {
+		locations = *filters.AvailabilityZones
 	} else if filters.Region != nil {
-		location = *filters.Region
+		locations = []string{*filters.Region}
 	}
-	locationInstanceOfferings, err := itf.RetrieveInstanceTypesSupportedInLocation(location)
+	locationInstanceOfferings, err := itf.RetrieveInstanceTypesSupportedInLocations(locations)
 	if err != nil {
 		return nil, err
 	}
@@ -329,41 +332,68 @@ func (itf Selector) executeFilters(filterToInstanceSpecMapping map[string]filter
 	return true, nil
 }
 
-// RetrieveInstanceTypesSupportedInLocation returns a map of instance type -> AZ or Region for all instance types supported in the location passed in
+// RetrieveInstanceTypesSupportedInLocations returns a map of instance type -> AZ or Region for all instance types supported in the intersected locations passed in
 // The location can be a zone-id (ie. use1-az1), a zone-name (us-east-1a), or a region name (us-east-1).
 // Note that zone names are not necessarily the same across accounts
-func (itf Selector) RetrieveInstanceTypesSupportedInLocation(zone string) (map[string]string, error) {
-	if zone == "" {
+func (itf Selector) RetrieveInstanceTypesSupportedInLocations(locations []string) (map[string]string, error) {
+	if len(locations) == 0 {
 		return nil, nil
 	}
-	availableInstanceTypes := map[string]string{}
-	instanceTypeOfferingsInput := &ec2.DescribeInstanceTypeOfferingsInput{
-		Filters: []*ec2.Filter{
-			{
-				Name:   aws.String(locationFilterKey),
-				Values: []*string{aws.String(zone)},
+	availableInstanceTypes := map[string]int{}
+	for _, location := range locations {
+		instanceTypeOfferingsInput := &ec2.DescribeInstanceTypeOfferingsInput{
+			Filters: []*ec2.Filter{
+				{
+					Name:   aws.String(locationFilterKey),
+					Values: []*string{aws.String(location)},
+				},
 			},
-		},
-	}
-	if isZoneID, _ := regexp.MatchString(zoneIDRegex, zone); isZoneID {
-		instanceTypeOfferingsInput.SetLocationType(zoneIDLocationType)
-	} else if isZoneName, _ := regexp.MatchString(zoneNameRegex, zone); isZoneName {
-		instanceTypeOfferingsInput.SetLocationType(zoneNameLocationType)
-	} else if isRegion, _ := regexp.MatchString(regionNameRegex, zone); isRegion {
-		instanceTypeOfferingsInput.SetLocationType(regionNameLocationType)
-	} else {
-		return nil, fmt.Errorf("The location passed in (%s) is not a valid zone-id, zone-name, or region name", zone)
+		}
+		locationType, err := itf.getLocationType(location)
+		if err != nil {
+			return nil, err
+		}
+		instanceTypeOfferingsInput.SetLocationType(locationType)
+
+		err = itf.EC2.DescribeInstanceTypeOfferingsPages(instanceTypeOfferingsInput, func(page *ec2.DescribeInstanceTypeOfferingsOutput, lastPage bool) bool {
+			for _, instanceType := range page.InstanceTypeOfferings {
+				if i, ok := availableInstanceTypes[*instanceType.InstanceType]; !ok {
+					availableInstanceTypes[*instanceType.InstanceType] = 1
+				} else {
+					availableInstanceTypes[*instanceType.InstanceType] = i + 1
+				}
+			}
+			return true
+		})
+		if err != nil {
+			return nil, fmt.Errorf("Encountered an error when describing instance type offerings: %w", err)
+		}
 	}
-	err := itf.EC2.DescribeInstanceTypeOfferingsPages(instanceTypeOfferingsInput, func(page *ec2.DescribeInstanceTypeOfferingsOutput, lastPage bool) bool {
-		for _, instanceType := range page.InstanceTypeOfferings {
-			availableInstanceTypes[*instanceType.InstanceType] = *instanceType.Location
+	availableInstanceTypesAllLocations := map[string]string{}
+	for instanceType, locationsSupported := range availableInstanceTypes {
+		if locationsSupported == len(locations) {
+			availableInstanceTypesAllLocations[instanceType] = ""
 		}
-		return true
-	})
+	}
+
+	return availableInstanceTypesAllLocations, nil
+}
+
+func (itf Selector) getLocationType(location string) (string, error) {
+	azs, err := itf.EC2.DescribeAvailabilityZones(&ec2.DescribeAvailabilityZonesInput{})
 	if err != nil {
-		return nil, fmt.Errorf("Encountered an error when describing instance type offerings: %w", err)
+		return "", err
+	}
+	for _, zone := range azs.AvailabilityZones {
+		if location == *zone.RegionName {
+			return regionNameLocationType, nil
+		} else if location == *zone.ZoneName {
+			return zoneNameLocationType, nil
+		} else if location == *zone.ZoneId {
+			return zoneIDLocationType, nil
+		}
 	}
-	return availableInstanceTypes, nil
+	return "", fmt.Errorf("The location passed in (%s) is not a valid zone-id, zone-name, or region name", location)
 }
 
 func isSupportedInLocation(instanceOfferings map[string]string, instanceType string) bool {