Cleanup pools on dlite termination (#150)

command/harness/delegate/delegate.go

@@ -137,23 +137,15 @@ func (c *delegateCommand) run(*kingpin.ParseContext) error {
 		return runnerServer.ListenAndServe(ctx)
 	})
 
+	g.Go(func() error {
+		return harness.Cleanup(ctx, &c.env, c.poolManager)
+	})
+
 	waitErr := g.Wait()
 	if waitErr != nil {
 		logrus.WithError(waitErr).
 			Errorln("shutting down the server")
 	}
-
-	// lets remove any old instances.
-	if !env.Settings.ReusePool {
-		cleanErr := c.poolManager.CleanPools(context.Background(), true, true)
-		if cleanErr != nil {
-			logrus.WithError(cleanErr).
-				Errorln("delegate: unable to clean pools")
-		} else {
-			logrus.Infoln("delegate: pools cleaned")
-		}
-	}
-
 	return waitErr
 }
 

command/harness/dlite/dlite.go

@@ -11,11 +11,13 @@ import (
 	"github.com/drone-runners/drone-runner-aws/store/database"
 	loghistory "github.com/drone/runner-go/logger/history"
 	"github.com/drone/signal"
-	"github.com/joho/godotenv"
-	"github.com/sirupsen/logrus"
 	"github.com/wings-software/dlite/delegate"
 	"github.com/wings-software/dlite/poller"
 	"github.com/wings-software/dlite/router"
+
+	"github.com/joho/godotenv"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
 	"gopkg.in/alecthomas/kingpin.v2"
 )
 
@@ -103,22 +105,26 @@ func (c *dliteCommand) run(*kingpin.ParseContext) error {
 	hook := loghistory.New()
 	logrus.AddHook(hook)
 
-	// TODO (Vistaar): Add support for tags based on available pools
-	err = c.startPoller(ctx, []string{})
-	if err != nil {
-		logrus.WithError(err).Error("could not start poller")
-		return err
-	}
-
-	// lets remove any old instances.
-	if !env.Settings.ReusePool {
-		cleanErr := c.poolManager.CleanPools(context.Background(), true, true)
-		if cleanErr != nil {
-			logrus.WithError(cleanErr).Errorln("dlite: unable to clean pools")
-		} else {
-			logrus.Infoln("dlite: pools cleaned")
+	var g errgroup.Group
+	g.Go(func() error {
+		// TODO (Vistaar): Add support for tags based on available pools
+		err = c.startPoller(ctx, []string{})
+		if err != nil {
+			logrus.WithError(err).Error("could not start poller")
+			return err
 		}
-	}
+		return nil
+	})
 
+	g.Go(func() error {
+		return harness.Cleanup(ctx, &c.env, c.poolManager)
+	})
+
+	waitErr := g.Wait()
+	if waitErr != nil {
+		logrus.WithError(waitErr).
+			Errorln("shutting down dlite")
+		return waitErr
+	}
 	return nil
 }

command/harness/pool.go

@@ -65,3 +65,21 @@ func SetupPool(ctx context.Context, env *config.EnvConfig, poolManager *drivers.
 	logrus.Infoln("pool created")
 	return nil
 }
+
+func Cleanup(ctx context.Context, env *config.EnvConfig, poolManager *drivers.Manager) error {
+	if env.Settings.ReusePool {
+		return nil
+	}
+
+	<-ctx.Done()
+	// clean up pool on termination
+	cleanErr := poolManager.CleanPools(context.Background(), true, true)
+
+	if cleanErr != nil {
+		logrus.WithError(cleanErr).
+			Errorln("unable to clean pools")
+	} else {
+		logrus.Infoln("pools cleaned")
+	}
+	return cleanErr
+}

internal/drivers/amazon/driver.go

@@ -169,7 +169,7 @@ func checkIngressRules(ctx context.Context, client *ec2.EC2, groupID string) err
 		}
 	}
 	if !found {
-		return fmt.Errorf("security group %s does not have the correct ingress rules There is no rule for port %d", *securityGroup.GroupName, lehelper.LiteEnginePort)
+		return fmt.Errorf("security group %s does not have the correct ingress rules. There is no rule for port %d", *securityGroup.GroupName, lehelper.LiteEnginePort)
 	}
 	return nil
 }
@@ -336,7 +336,7 @@ func (p *config) Create(ctx context.Context, opts *types.InstanceCreateOpts) (in
 		WithField("time", fmt.Sprintf("%.2fs", time.Since(startTime).Seconds())).
 		Debugln("amazon: [provision] complete")
 
-	return // nolint:nakedret
+	return //nolint:nakedret
 }
 
 // Destroy destroys the server AWS EC2 instances.

internal/drivers/digitalocean/driver.go

@@ -236,7 +236,7 @@ func newClient(ctx context.Context, pat string) *godo.Client {
 	)
 }
 
-//	take a slice of ssh keys and return a slice of godo.DropletCreateSSHKey
+// take a slice of ssh keys and return a slice of godo.DropletCreateSSHKey
 func createSSHKeys(sshKeys []string) []godo.DropletCreateSSHKey {
 	var keys []godo.DropletCreateSSHKey
 	for _, key := range sshKeys {