go/gce/autoscaler/autoscaler.go

package autoscaler

import (
	"fmt"
	"sort"
	"sync"

	"go.skia.org/infra/go/gce"
	"go.skia.org/infra/go/util"
	"golang.org/x/oauth2"
)

// Interface useful for mocking.
// TODO(borenet): This doesn't really "auto" scale anything.
type IAutoscaler interface {
	// GetInstanceStatuses returns a map of instance names to booleans
	// indicating whether each instance is online as of the last Update().
	GetInstanceStatuses() map[string]bool

	// GetNamesOfManagedInstances returns names of all instances managed by
	// this autoscaler.
	GetNamesOfManagedInstances() []string

	// GetOnlineInstances returns a slice of all online instance names as of
	// the last Update().
	GetOnlineInstances() []string

	// Start the given instances.
	// Note: This method returns when all instances are in RUNNING state.
	// Does not check to see if they are ready (ssh-able).
	Start([]string) error

	// StartAllInstances starts all instances.
	// Note: This method returns when all instances are in RUNNING state.
	// Does not check to see if they are ready (ssh-able).
	StartAllInstances() error

	// Stop the given instances.
	Stop([]string) error

	// StopAllInstances stops all instances.
	StopAllInstances() error

	// Update determines which instances are online and caches that
	// information.
	Update() error
}

// Autoscaler is a struct used for autoscaling instances in GCE.
type Autoscaler struct {
	g             *gce.GCloud
	instanceNames []string
	instances     map[string]*gce.Instance
	mtx           sync.RWMutex // protects a.online
	online        []bool
}

// Helper function for creating lists of instances. The given range is
// inclusive.
func GetInstanceRange(min, max int, getInstance func(int) *gce.Instance) []*gce.Instance {
	rv := make([]*gce.Instance, 0, max-min+1)
	for i := min; i <= max; i++ {
		rv = append(rv, getInstance(i))
	}
	return rv
}

// Helper function for creating lists of instances.
func GetInstanceSet(intSet string, getInstance func(int) *gce.Instance) ([]*gce.Instance, error) {
	nums, err := util.ParseIntSet(intSet)
	if err != nil {
		return nil, err
	}
	rv := make([]*gce.Instance, 0, len(nums))
	for _, num := range nums {
		rv = append(rv, getInstance(num))
	}
	return rv, nil
}

// NewAutoscaler returns an Autoscaler instance which manages the given GCE
// instances. Automatically calls Update().
func NewAutoscaler(projectId, zone string, ts oauth2.TokenSource, instances []*gce.Instance) (*Autoscaler, error) {
	// Create the GCloud object.
	g, err := gce.NewGCloud(projectId, zone, ts)
	if err != nil {
		return nil, err
	}
	// Create map of instances.
	instanceMap := make(map[string]*gce.Instance, len(instances))
	instanceNames := make([]string, 0, len(instances))
	for _, instance := range instances {
		instanceMap[instance.Name] = instance
		instanceNames = append(instanceNames, instance.Name)
	}
	sort.Strings(instanceNames)
	a := &Autoscaler{
		g:             g,
		instanceNames: instanceNames,
		instances:     instanceMap,
	}
	if err := a.Update(); err != nil {
		return nil, err
	}
	return a, nil
}

// See documentation for IAutoscaler.
func (a *Autoscaler) Update() error {
	online := make([]bool, len(a.instanceNames))
	group := util.NewNamedErrGroup()
	for idx, instance := range a.instanceNames {
		// https://golang.org/doc/faq#closures_and_goroutines
		idx := idx
		instance := instance
		group.Go(instance, func() error {
			isOnline, err := a.g.IsInstanceRunning(instance)
			if err != nil {
				return err
			}
			online[idx] = isOnline
			return nil
		})
	}
	if err := group.Wait(); err != nil {
		return err
	}
	a.mtx.Lock()
	defer a.mtx.Unlock()
	a.online = online
	return nil
}

// See documentation for IAutoscaler.
func (a *Autoscaler) GetInstanceStatuses() map[string]bool {
	a.mtx.RLock()
	defer a.mtx.RUnlock()
	rv := make(map[string]bool, len(a.instanceNames))
	for idx, name := range a.instanceNames {
		rv[name] = a.online[idx]
	}
	return rv
}

// See documentation for IAutoscaler.
func (a *Autoscaler) GetOnlineInstances() []string {
	a.mtx.RLock()
	defer a.mtx.RUnlock()
	onlineInstances := make([]string, 0, len(a.instances))
	for idx, name := range a.instanceNames {
		if a.online[idx] {
			onlineInstances = append(onlineInstances, name)
		}
	}
	return onlineInstances
}

// See documentation for IAutoscaler.
func (a *Autoscaler) GetNamesOfManagedInstances() []string {
	return util.CopyStringSlice(a.instanceNames)
}

// Run the given function in parallel over the given instances.
func (a *Autoscaler) processInstances(instanceNames []string, fn func(*gce.Instance) error) error {
	// Get the requested instances.
	instances := make([]*gce.Instance, 0, len(instanceNames))
	for _, name := range instanceNames {
		instance, ok := a.instances[name]
		if !ok {
			return fmt.Errorf("Unknown instance %q", name)
		}
		instances = append(instances, instance)
	}
	// Start the instances.
	group := util.NewNamedErrGroup()
	for _, instance := range instances {
		instance := instance // https://golang.org/doc/faq#closures_and_goroutines
		group.Go(instance.Name, func() error {
			return fn(instance)
		})
	}
	if err := group.Wait(); err != nil {
		return err
	}
	return nil
}

// See documentation for IAutoscaler.
func (a *Autoscaler) Start(instanceNames []string) error {
	return a.processInstances(instanceNames, func(instance *gce.Instance) error {
		return a.g.StartWithoutReadyCheck(instance)
	})
}

// See documentation for IAutoscaler.
func (a *Autoscaler) StartAllInstances() error {
	return a.Start(a.instanceNames)
}

// See documentation for IAutoscaler.
func (a *Autoscaler) Stop(instanceNames []string) error {
	return a.processInstances(instanceNames, func(instance *gce.Instance) error {
		return a.g.Stop(instance)
	})
}

// See documentation for IAutoscaler.
func (a *Autoscaler) StopAllInstances() error {
	return a.Stop(a.instanceNames)
}

var _ IAutoscaler = &Autoscaler{}