Merge pull request #124 from weaveworks/add-az-selector-logic

Add AZ selector logic

cmd/eksctl/create.go

@@ -117,6 +117,10 @@ func doCreateCluster(cfg *eks.ClusterConfig, name string) error {
 		return fmt.Errorf("--region=%s is not supported only %s and %s are supported", cfg.Region, EKS_REGION_US_WEST_2, EKS_REGION_US_EAST_1)
 	}
 
+	if err := ctl.SetAvailabilityZones(); err != nil {
+		return err
+	}
+
 	if err := ctl.LoadSSHPublicKey(); err != nil {
 		return err
 	}

pkg/az/az.go

@@ -0,0 +1,148 @@
+package az
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/pkg/errors"
+)
+
+const (
+	DefaultRequiredAvailabilityZones = 3
+)
+
+// SelectionStrategy provides an interface to allow changing the strategy used to
+// select availaibility zones to use from a list available.
+type SelectionStrategy interface {
+	Select(availableZones []string) []string
+}
+
+// RequiredNumberRandomStrategy selects az zones randomly up to a required amount
+// of zones.
+type RequiredNumberRandomStrategy struct {
+	RequiredAvailabilityZones int
+}
+
+// Select will randomly select az from the supplied list. The number of az's
+// selected will be controlled by RequiredAvailabilityZones.
+func (r *RequiredNumberRandomStrategy) Select(availableZones []string) []string {
+	zones := []string{}
+	for len(zones) < r.RequiredAvailabilityZones {
+		rand := rand.New(rand.NewSource(time.Now().UnixNano()))
+		for _, rn := range rand.Perm(len(availableZones)) {
+			zones = append(zones, availableZones[rn])
+			if len(zones) == r.RequiredAvailabilityZones {
+				break
+			}
+		}
+	}
+	return zones
+}
+
+// NewRequiredNumberRandomStrategy returns a RequiredNumberRandomStrategy that
+// has the number of required zones set to the default (DefaultRequiredAvailabilityZones)
+func NewRequiredNumberRandomStrategy() *RequiredNumberRandomStrategy {
+	return &RequiredNumberRandomStrategy{RequiredAvailabilityZones: DefaultRequiredAvailabilityZones}
+}
+
+// ZoneUsageRule provides an interface to enable rules to determine if a
+// zone should be used.
+type ZoneUsageRule interface {
+	CanUseZone(zone *ec2.AvailabilityZone) bool
+}
+
+// ZonesToAvoidRule can be used to ensure that certain az aren't used. This can be used
+// to avoid zones that are know to be overpopulated.
+type ZonesToAvoidRule struct {
+	zonesToAvoid map[string]bool
+}
+
+// CanUseZone checks if the zupplied zone is in the list of zones to be avoided.
+func (za *ZonesToAvoidRule) CanUseZone(zone *ec2.AvailabilityZone) bool {
+	_, avoidZone := za.zonesToAvoid[*zone.ZoneName]
+	return !avoidZone
+}
+
+// NewZonesToAvoidRule returns a new ZonesToAvoidRule with the supplied
+// zones set to avoid
+func NewZonesToAvoidRule(zonesToAvoid map[string]bool) *ZonesToAvoidRule {
+	return &ZonesToAvoidRule{zonesToAvoid: zonesToAvoid}
+}
+
+// AvailabilityZoneSelector used to select availability zones to use
+type AvailabilityZoneSelector struct {
+	ec2api   ec2iface.EC2API
+	strategy SelectionStrategy
+	rules    []ZoneUsageRule
+}
+
+// NewSelectorWithDefaults create a new AvailabilityZoneSelector with the
+// defaukt selection strategy and usage rules
+func NewSelectorWithDefaults(ec2api ec2iface.EC2API) *AvailabilityZoneSelector {
+	avoidZones := map[string]bool{
+		// well-known over-populated zones
+		"us-east1-a": true,
+		"us-east1-b": true,
+	}
+
+	return &AvailabilityZoneSelector{
+		ec2api:   ec2api,
+		strategy: NewRequiredNumberRandomStrategy(),
+		rules:    []ZoneUsageRule{NewZonesToAvoidRule(avoidZones)},
+	}
+}
+
+// SelectZones returns a list fo az zones to use for the supplied region
+func (a *AvailabilityZoneSelector) SelectZones(regionName string) ([]string, error) {
+	availableZones, err := a.getZonesForRegion(regionName)
+	if err != nil {
+		return nil, err
+	}
+
+	usableZones := a.getUsableZones(availableZones)
+
+	return a.strategy.Select(usableZones), nil
+}
+
+func (a *AvailabilityZoneSelector) getUsableZones(availableZones []*ec2.AvailabilityZone) []string {
+	usableZones := []string{}
+	for _, zone := range availableZones {
+		zoneUsable := true
+		for _, rule := range a.rules {
+			if !rule.CanUseZone(zone) {
+				zoneUsable = false
+				break
+			}
+		}
+		if zoneUsable {
+			usableZones = append(usableZones, *zone.ZoneName)
+		}
+	}
+
+	return usableZones
+}
+
+func (a *AvailabilityZoneSelector) getZonesForRegion(regionName string) ([]*ec2.AvailabilityZone, error) {
+	regionFilter := &ec2.Filter{
+		Name:   aws.String("region-name"),
+		Values: []*string{aws.String(regionName)},
+	}
+	stateFilter := &ec2.Filter{
+		Name:   aws.String("state"),
+		Values: []*string{aws.String(ec2.AvailabilityZoneStateAvailable)},
+	}
+
+	input := &ec2.DescribeAvailabilityZonesInput{
+		Filters: []*ec2.Filter{regionFilter, stateFilter},
+	}
+
+	output, err := a.ec2api.DescribeAvailabilityZones(input)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting availibility zones for %s", regionName)
+	}
+
+	return output.AvailabilityZones, nil
+}

pkg/az/az_suite_test.go

@@ -0,0 +1,13 @@
+package az_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+func TestAZ(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "AZ Suite")
+}