Implement multi-app scaling

superfly · Apr 26, 2024 · 4fdf6a9 · 4fdf6a9
1 parent cc749e1
commit 4fdf6a9
Show file tree

Hide file tree

Showing 9 changed files with 800 additions and 371 deletions.
diff --git a/cmd/fly-autoscaler/eval.go b/cmd/fly-autoscaler/eval.go
@@ -33,7 +33,7 @@ func (c *EvalCommand) Run(ctx context.Context, args []string) (err error) {
 	}
 
 	// Instantiate reconciler and evaluate once.
-	r := fas.NewReconciler(nil)
+	r := fas.NewReconciler()
 	r.MinCreatedMachineN = c.Config.GetMinCreatedMachineN()
 	r.MaxCreatedMachineN = c.Config.GetMaxCreatedMachineN()
 	r.MinStartedMachineN = c.Config.GetMinStartedMachineN()

diff --git a/cmd/fly-autoscaler/main.go b/cmd/fly-autoscaler/main.go
@@ -9,12 +9,14 @@ import (
 	"log/slog"
 	"os"
 	"os/signal"
+	"strconv"
 	"strings"
 	"time"
 
 	fas "github.com/superfly/fly-autoscaler"
 	fasprom "github.com/superfly/fly-autoscaler/prometheus"
 	"github.com/superfly/fly-autoscaler/temporal"
+	fly "github.com/superfly/fly-go"
 	"github.com/superfly/fly-go/flaps"
 	"github.com/superfly/fly-go/tokens"
 	"gopkg.in/yaml.v3"
@@ -106,30 +108,37 @@ func registerConfigPathFlag(fs *flag.FlagSet) *string {
 }
 
 type Config struct {
-	AppName            string        `yaml:"app-name"`
-	Regions            []string      `yaml:"regions"`
-	CreatedMachineN    string        `yaml:"created-machine-count"`
-	MinCreatedMachineN string        `yaml:"min-created-machine-count"`
-	MaxCreatedMachineN string        `yaml:"max-created-machine-count"`
-	StartedMachineN    string        `yaml:"started-machine-count"`
-	MinStartedMachineN string        `yaml:"min-started-machine-count"`
-	MaxStartedMachineN string        `yaml:"max-started-machine-count"`
-	Interval           time.Duration `yaml:"interval"`
-	APIToken           string        `yaml:"api-token"`
-	Verbose            bool          `yaml:"verbose"`
+	AppName                string        `yaml:"app-name"`
+	Org                    string        `yaml:"org"`
+	Regions                []string      `yaml:"regions"`
+	CreatedMachineN        string        `yaml:"created-machine-count"`
+	MinCreatedMachineN     string        `yaml:"min-created-machine-count"`
+	MaxCreatedMachineN     string        `yaml:"max-created-machine-count"`
+	StartedMachineN        string        `yaml:"started-machine-count"`
+	MinStartedMachineN     string        `yaml:"min-started-machine-count"`
+	MaxStartedMachineN     string        `yaml:"max-started-machine-count"`
+	Concurrency            int           `yaml:"concurrency"`
+	Interval               time.Duration `yaml:"interval"`
+	Timeout                time.Duration `yaml:"timeout"`
+	AppListRefreshInterval time.Duration `yaml:"app-list-refresh-interval"`
+	APIToken               string        `yaml:"api-token"`
+	Verbose                bool          `yaml:"verbose"`
 
 	MetricCollectors []*MetricCollectorConfig `yaml:"metric-collectors"`
 }
 
 func NewConfig() *Config {
 	return &Config{
-		Interval: fas.DefaultReconcilerInterval,
+		Concurrency: fas.DefaultConcurrency,
+		Interval:    fas.DefaultReconcileInterval,
+		Timeout:     fas.DefaultReconcileTimeout,
 	}
 }
 
-func NewConfigFromEnv() (*Config, error) {
+func NewConfigFromEnv() (_ *Config, err error) {
 	c := NewConfig()
 	c.AppName = os.Getenv("FAS_APP_NAME")
+	c.Org = os.Getenv("FAS_ORG")
 	c.CreatedMachineN = os.Getenv("FAS_CREATED_MACHINE_COUNT")
 	c.MinCreatedMachineN = os.Getenv("FAS_MIN_CREATED_MACHINE_COUNT")
 	c.MaxCreatedMachineN = os.Getenv("FAS_MAX_CREATED_MACHINE_COUNT")
@@ -142,12 +151,26 @@ func NewConfigFromEnv() (*Config, error) {
 		c.Regions = strings.Split(s, ",")
 	}
 
+	if s := os.Getenv("FAS_CONCURRENCY"); s != "" {
+		if c.Concurrency, err = strconv.Atoi(s); err != nil {
+			return nil, fmt.Errorf("cannot parse FAS_CONCURRENCY as integer: %q", s)
+		}
+	}
+
 	if s := os.Getenv("FAS_INTERVAL"); s != "" {
-		d, err := time.ParseDuration(s)
-		if err != nil {
+		if c.Interval, err = time.ParseDuration(s); err != nil {
 			return nil, fmt.Errorf("cannot parse FAS_INTERVAL as duration: %q", s)
 		}
-		c.Interval = d
+	}
+	if s := os.Getenv("FAS_TIMEOUT"); s != "" {
+		if c.Timeout, err = time.ParseDuration(s); err != nil {
+			return nil, fmt.Errorf("cannot parse FAS_TIMEOUT as duration: %q", s)
+		}
+	}
+	if s := os.Getenv("FAS_APP_LIST_REFRESH_INTERVAL"); s != "" {
+		if c.AppListRefreshInterval, err = time.ParseDuration(s); err != nil {
+			return nil, fmt.Errorf("cannot parse FAS_APP_LIST_REFRESH_INTERVAL as duration: %q", s)
+		}
 	}
 
 	if addr := os.Getenv("FAS_PROMETHEUS_ADDRESS"); addr != "" {
@@ -278,15 +301,28 @@ func (c *Config) validateStartedMachineCount() error {
 	return nil
 }
 
-func (c *Config) NewFlapsClient(ctx context.Context) (*flaps.Client, error) {
+func (c *Config) NewFlyClient(ctx context.Context) (*fly.Client, error) {
 	if c.APIToken == "" {
 		return nil, fmt.Errorf("api token required")
 	}
 
-	return flaps.NewWithOptions(ctx, flaps.NewClientOpts{
-		AppName: c.AppName,
-		Tokens:  tokens.Parse(c.APIToken),
-	})
+	return fly.NewClientFromOptions(fly.ClientOptions{
+		Tokens: tokens.Parse(c.APIToken),
+	}), nil
+}
+
+func (c *Config) NewFlapsClient() (fas.NewFlapsClientFunc, error) {
+	if c.APIToken == "" {
+		return nil, fmt.Errorf("api token required")
+	}
+	tok := tokens.Parse(c.APIToken)
+
+	return func(ctx context.Context, appName string) (fas.FlapsClient, error) {
+		return flaps.NewWithOptions(ctx, flaps.NewClientOpts{
+			AppName: appName,
+			Tokens:  tok,
+		})
+	}, nil
 }
 
 func (c *Config) NewMetricCollectors() ([]fas.MetricCollector, error) {

diff --git a/cmd/fly-autoscaler/serve.go b/cmd/fly-autoscaler/serve.go
@@ -16,17 +16,19 @@ import (
 
 // ServeCommand represents a command run the autoscaler server process.
 type ServeCommand struct {
-	reconciler *fas.Reconciler
-	Config     *Config
+	pool   *fas.ReconcilerPool
+	Config *Config
 }
 
 func NewServeCommand() *ServeCommand {
 	return &ServeCommand{}
 }
 
 func (c *ServeCommand) Close() (err error) {
-	if c.reconciler != nil {
-		c.reconciler.Stop()
+	if c.pool != nil {
+		if err := c.pool.Close(); err != nil {
+			slog.Warn("failed to close reconciler pool", slog.Any("err", err))
+		}
 	}
 	return nil
 }
@@ -43,12 +45,12 @@ func (c *ServeCommand) Run(ctx context.Context, args []string) (err error) {
 		return err
 	}
 
-	// Instantiate client to scale up machines.
-	client, err := c.Config.NewFlapsClient(ctx)
+	// Instantiate clients for access org/apps & for scaling machines.
+	flyClient, err := c.Config.NewFlyClient(ctx)
 	if err != nil {
-		return fmt.Errorf("cannot create flaps client: %w", err)
+		return fmt.Errorf("cannot create fly client: %w", err)
 	}
-	slog.Info("connected to flaps", slog.String("app_name", c.Config.AppName))
+	slog.Info("connected to fly")
 
 	// Instantiate prometheus collector.
 	collectors, err := c.Config.NewMetricCollectors()
@@ -57,47 +59,67 @@ func (c *ServeCommand) Run(ctx context.Context, args []string) (err error) {
 	}
 	slog.Info("metrics collectors initialized", slog.Int("n", len(collectors)))
 
-	// Instantiate and start reconcilation.
-	r := fas.NewReconciler(client)
-	r.MinCreatedMachineN = c.Config.GetMinCreatedMachineN()
-	r.MaxCreatedMachineN = c.Config.GetMaxCreatedMachineN()
-	r.MinStartedMachineN = c.Config.GetMinStartedMachineN()
-	r.MaxStartedMachineN = c.Config.GetMaxStartedMachineN()
-	r.Regions = c.Config.Regions
-	r.Interval = c.Config.Interval
-	r.Collectors = collectors
-	r.RegisterPromMetrics(prometheus.DefaultRegisterer)
-	c.reconciler = r
+	minCreatedMachineN := c.Config.GetMinCreatedMachineN()
+	maxCreatedMachineN := c.Config.GetMaxCreatedMachineN()
+	minStartedMachineN := c.Config.GetMinStartedMachineN()
+	maxStartedMachineN := c.Config.GetMaxStartedMachineN()
+
+	// Instantiate pool.
+	p := fas.NewReconcilerPool(flyClient, c.Config.Concurrency)
+	if p.NewFlapsClient, err = c.Config.NewFlapsClient(); err != nil {
+		return fmt.Errorf("cannot initialize flaps client constructor: %w", err)
+	}
+	p.NewReconciler = func() *fas.Reconciler {
+		r := fas.NewReconciler()
+		r.MinCreatedMachineN = minCreatedMachineN
+		r.MaxCreatedMachineN = maxCreatedMachineN
+		r.MinStartedMachineN = minStartedMachineN
+		r.MaxStartedMachineN = maxStartedMachineN
+		r.Regions = c.Config.Regions
+		r.Collectors = collectors
+		return r
+	}
+	p.AppName = c.Config.AppName
+	p.OrganizationSlug = c.Config.Org
+	p.ReconcileInterval = c.Config.Interval
+	p.ReconcileTimeout = c.Config.Timeout
+	p.AppListRefreshInterval = c.Config.AppListRefreshInterval
+	p.RegisterPromMetrics(prometheus.DefaultRegisterer)
+	c.pool = p
 
 	attrs := []any{
-		slog.Duration("interval", r.Interval),
-		slog.Int("collectors", len(r.Collectors)),
+		slog.Duration("interval", p.ReconcileInterval),
+		slog.Duration("timeout", p.ReconcileTimeout),
+		slog.Duration("appListRefreshInterval", p.AppListRefreshInterval),
+		slog.Int("collectors", len(collectors)),
 	}
 
-	if len(r.Regions) > 0 {
-		attrs = append(attrs, slog.Any("regions", r.Regions))
+	if regions := c.Config.Regions; len(regions) > 0 {
+		attrs = append(attrs, slog.Any("regions", regions))
 	}
 
-	if r.MinCreatedMachineN == r.MaxCreatedMachineN {
-		attrs = append(attrs, slog.String("created", r.MinCreatedMachineN))
-	} else if r.MinCreatedMachineN != "" || r.MaxCreatedMachineN != "" {
+	if minCreatedMachineN == maxCreatedMachineN {
+		attrs = append(attrs, slog.String("created", minCreatedMachineN))
+	} else if minCreatedMachineN != "" || maxCreatedMachineN != "" {
 		attrs = append(attrs, slog.Group("created",
-			slog.String("min", r.MinCreatedMachineN),
-			slog.String("max", r.MaxCreatedMachineN),
+			slog.String("min", minCreatedMachineN),
+			slog.String("max", maxCreatedMachineN),
 		))
 	}
 
-	if r.MinStartedMachineN == r.MaxStartedMachineN {
-		attrs = append(attrs, slog.String("started", r.MinStartedMachineN))
-	} else if r.MinStartedMachineN != "" || r.MaxStartedMachineN != "" {
+	if minStartedMachineN == maxStartedMachineN {
+		attrs = append(attrs, slog.String("started", minStartedMachineN))
+	} else if minStartedMachineN != "" || maxStartedMachineN != "" {
 		attrs = append(attrs, slog.Group("started",
-			slog.String("min", r.MinStartedMachineN),
-			slog.String("max", r.MaxStartedMachineN),
+			slog.String("min", minStartedMachineN),
+			slog.String("max", maxStartedMachineN),
 		))
 	}
 
-	slog.Info("reconciler initialized, beginning loop", attrs...)
-	r.Start()
+	slog.Info("reconciler pool initialized, beginning loop", attrs...)
+	if err := p.Open(); err != nil {
+		return fmt.Errorf("cannot initialize reconciler pool: %w", err)
+	}
 
 	go c.serveMetricsServer(ctx)
 

diff --git a/fas.go b/fas.go
@@ -15,6 +15,13 @@ var (
 	ErrExprInf      = errors.New("expression returned Inf")
 )
 
+var _ FlyClient = (*fly.Client)(nil)
+
+type FlyClient interface {
+	GetOrganizationBySlug(ctx context.Context, slug string) (*fly.Organization, error)
+	GetAppsForOrganization(ctx context.Context, orgID string) ([]fly.App, error)
+}
+
 var _ FlapsClient = (*flaps.Client)(nil)
 
 type FlapsClient interface {
@@ -24,3 +31,5 @@ type FlapsClient interface {
 	Start(ctx context.Context, id, nonce string) (*fly.MachineStartResponse, error)
 	Stop(ctx context.Context, in fly.StopMachineInput, nonce string) error
 }
+
+type NewFlapsClientFunc func(ctx context.Context, appName string) (FlapsClient, error)
diff --git a/mock/fly_client.go b/mock/fly_client.go
@@ -0,0 +1,23 @@
+package mock
+
+import (
+	"context"
+
+	fas "github.com/superfly/fly-autoscaler"
+	"github.com/superfly/fly-go"
+)
+
+var _ fas.FlyClient = (*FlyClient)(nil)
+
+type FlyClient struct {
+	GetOrganizationBySlugFunc  func(ctx context.Context, slug string) (*fly.Organization, error)
+	GetAppsForOrganizationFunc func(ctx context.Context, orgID string) ([]fly.App, error)
+}
+
+func (m *FlyClient) GetOrganizationBySlug(ctx context.Context, slug string) (*fly.Organization, error) {
+	return m.GetOrganizationBySlugFunc(ctx, slug)
+}
+
+func (m *FlyClient) GetAppsForOrganization(ctx context.Context, orgID string) ([]fly.App, error) {
+	return m.GetAppsForOrganizationFunc(ctx, orgID)
+}