From 24da47d4ff886125d944da1e21750ea02b8a1d0c Mon Sep 17 00:00:00 2001 From: Vistaar Juneja Date: Wed, 20 Nov 2024 14:30:36 -0800 Subject: [PATCH] feat: [PIPE-22742]: Refactor pool management --- command/harness/delegate/delegate.go | 6 +-- command/harness/dlite/dlite.go | 6 +-- command/harness/pool.go | 67 ++++++++++++++++++++++------ 3 files changed, 59 insertions(+), 20 deletions(-) diff --git a/command/harness/delegate/delegate.go b/command/harness/delegate/delegate.go index efdb1403..cc6d2475 100644 --- a/command/harness/delegate/delegate.go +++ b/command/harness/delegate/delegate.go @@ -109,8 +109,8 @@ func (c *delegateCommand) run(*kingpin.ParseContext) error { c.stageOwnerStore = stageOwnerStore c.poolManager = drivers.New(ctx, instanceStore, &c.env) - _, err = harness.SetupPool(ctx, &c.env, c.poolManager, c.poolFile) - defer harness.Cleanup(&c.env, c.poolManager, true, true) //nolint: errcheck + _, err = harness.SetupPoolWithEnv(ctx, &c.env, c.poolManager, c.poolFile) + defer harness.Cleanup(c.env.Settings.ReusePool, c.poolManager, true, true) //nolint: errcheck if err != nil { return err } @@ -134,7 +134,7 @@ func (c *delegateCommand) run(*kingpin.ParseContext) error { g.Go(func() error { <-ctx.Done() - return harness.Cleanup(&c.env, c.poolManager, true, true) + return harness.Cleanup(c.env.Settings.ReusePool, c.poolManager, true, true) }) g.Go(func() error { diff --git a/command/harness/dlite/dlite.go b/command/harness/dlite/dlite.go index 185fea80..2bed2978 100644 --- a/command/harness/dlite/dlite.go +++ b/command/harness/dlite/dlite.go @@ -118,7 +118,7 @@ func (c *dliteCommand) run(*kingpin.ParseContext) error { var poolConfig *config.PoolFile poolConfig, err = c.setupDistributedPool(ctx) - defer harness.Cleanup(&c.env, c.distributedPoolManager, false, true) //nolint: errcheck + defer harness.Cleanup(c.env.Settings.ReusePool, c.distributedPoolManager, false, true) //nolint: errcheck if err != nil { return err } @@ -143,7 +143,7 @@ func (c *dliteCommand) run(*kingpin.ParseContext) error { g.Go(func() error { <-ctx.Done() // delete unused instances for distributed pool - return harness.Cleanup(&c.env, c.distributedPoolManager, false, true) + return harness.Cleanup(c.env.Settings.ReusePool, c.distributedPoolManager, false, true) }) g.Go(func() error { @@ -196,7 +196,7 @@ func (c *dliteCommand) setupDistributedPool(ctx context.Context) (*config.PoolFi c.env.Settings.HarnessTestBinaryURI, c.env.Settings.PluginBinaryURI, c.env.Settings.AutoInjectionBinaryURI)) - poolConfig, err := harness.SetupPool(ctx, &c.env, c.distributedPoolManager, c.poolFile) + poolConfig, err := harness.SetupPoolWithEnv(ctx, &c.env, c.distributedPoolManager, c.poolFile) if err != nil { logrus.WithError(err).Error("could not setup distributed pool") return poolConfig, err diff --git a/command/harness/pool.go b/command/harness/pool.go index 4285eb06..697a19fd 100644 --- a/command/harness/pool.go +++ b/command/harness/pool.go @@ -7,16 +7,23 @@ import ( "github.com/drone-runners/drone-runner-aws/app/drivers" "github.com/drone-runners/drone-runner-aws/app/poolfile" "github.com/drone-runners/drone-runner-aws/command/config" + "github.com/drone-runners/drone-runner-aws/types" "github.com/sirupsen/logrus" ) -func SetupPool(ctx context.Context, env *config.EnvConfig, poolManager drivers.IManager, poolFile string) (*config.PoolFile, error) { - configPool, confErr := poolfile.ConfigPoolFile(poolFile, env) - if confErr != nil { - logrus.WithError(confErr).Fatalln("Unable to load pool file, or use an in memory pool") - } - - pools, err := poolfile.ProcessPool(configPool, env.Runner.Name, env.Passwords()) +// SetupPool sets up a pool of instances given a config pool. +func SetupPool( + ctx context.Context, + configPool *config.PoolFile, + runnerName string, + passwords types.Passwords, + poolManager drivers.IManager, + busyMaxAge int64, + freeMaxAge int64, + purgerTime int64, + reusePool bool, +) (*config.PoolFile, error) { + pools, err := poolfile.ProcessPool(configPool, runnerName, passwords) if err != nil { logrus.WithError(err).Errorln("unable to process pool file") return configPool, err @@ -36,17 +43,17 @@ func SetupPool(ctx context.Context, env *config.EnvConfig, poolManager drivers.I } // setup lifetimes of instances - busyMaxAge := time.Hour * time.Duration(env.Settings.BusyMaxAge) // includes time required to setup an instance - freeMaxAge := time.Hour * time.Duration(env.Settings.FreeMaxAge) - purgerTime := time.Minute * time.Duration(env.Settings.PurgerTime) - err = poolManager.StartInstancePurger(ctx, busyMaxAge, freeMaxAge, purgerTime) + busyMaxAgeDuration := time.Hour * time.Duration(busyMaxAge) // includes time required to setup an instance + freeMaxAgeDuration := time.Hour * time.Duration(freeMaxAge) + purgerDuration := time.Minute * time.Duration(purgerTime) + err = poolManager.StartInstancePurger(ctx, busyMaxAgeDuration, freeMaxAgeDuration, purgerDuration) if err != nil { logrus.WithError(err). Errorln("failed to start instance purger") return configPool, err } // lets remove any old instances. - if !env.Settings.ReusePool { + if !reusePool { cleanErr := poolManager.CleanPools(ctx, true, true) if cleanErr != nil { return configPool, cleanErr @@ -64,8 +71,40 @@ func SetupPool(ctx context.Context, env *config.EnvConfig, poolManager drivers.I return configPool, nil } -func Cleanup(env *config.EnvConfig, poolManager drivers.IManager, destroyBusy, destroyFree bool) error { - if env.Settings.ReusePool { +func SetupPoolWithFile( + ctx context.Context, + poolFilePath string, + poolManager drivers.IManager, + passwords types.Passwords, + runnerName string, + busyAge, + freeAge, + purgerTime int64, + reusePool bool, +) (*config.PoolFile, error) { + configPool, err := config.ParseFile(poolFilePath) + if err != nil { + logrus.WithError(err). + WithField("path", poolFilePath). + Errorln("exec: unable to parse pool file") + return nil, err + } + + return SetupPool(ctx, configPool, runnerName, passwords, poolManager, busyAge, freeAge, purgerTime, reusePool) +} + +func SetupPoolWithEnv(ctx context.Context, env *config.EnvConfig, poolManager drivers.IManager, poolFile string) (*config.PoolFile, error) { + configPool, confErr := poolfile.ConfigPoolFile(poolFile, env) + if confErr != nil { + logrus.WithError(confErr).Fatalln("Unable to load pool file, or use an in memory pool") + } + + return SetupPool(ctx, configPool, env.Runner.Name, env.Passwords(), poolManager, env.Settings.BusyMaxAge, env.Settings.FreeMaxAge, env.Settings.PurgerTime, env.Settings.ReusePool) + +} + +func Cleanup(reusePool bool, poolManager drivers.IManager, destroyBusy, destroyFree bool) error { + if reusePool { return nil }