Add dynamic wait for vm hibernate

.drone.yml

@@ -20,7 +20,7 @@ steps:
   - name: golangci-lint
     image: golangci/golangci-lint
     commands:
-      - golangci-lint run --timeout 300s
+      - golangci-lint run --timeout 500s
     volumes:
       - name: cache
         path: /go

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Masterminds/squirrel v1.5.2
 	github.com/aws/aws-sdk-go v1.42.10
 	github.com/buildkite/yaml v2.1.0+incompatible
+	github.com/cenkalti/backoff/v4 v4.1.1
 	github.com/dchest/uniuri v0.0.0-20160212164326-8902c56451e9
 	github.com/drone/drone-go v1.7.1
 	github.com/drone/envsubst v1.0.3
@@ -25,6 +26,7 @@ require (
 	github.com/maragudk/migrate v0.4.1
 	github.com/mattn/go-isatty v0.0.14
 	github.com/mattn/go-sqlite3 v1.14.7
+	github.com/pkg/errors v0.9.1
 	github.com/rs/zerolog v1.26.1
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.6.1
@@ -41,7 +43,6 @@ require (
 	github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect
 	github.com/alecthomas/units v0.0.0-20210927113745-59d0afb8317a // indirect
 	github.com/bmatcuk/doublestar v1.1.1 // indirect
-	github.com/cenkalti/backoff/v4 v4.1.1 // indirect
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/docker/distribution v2.7.1+incompatible // indirect
@@ -58,7 +59,6 @@ require (
 	github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect
 	github.com/natessilva/dag v0.0.0-20180124060714-7194b8dcc5c4 // indirect
 	github.com/opencontainers/go-digest v1.0.0-rc1 // indirect
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	go.opencensus.io v0.22.4 // indirect
 	golang.org/x/crypto v0.0.0-20211215165025-cf75a172585e // indirect

internal/drivers/amazon/pool.go

@@ -55,7 +55,8 @@ func (p *provider) Create(ctx context.Context, opts *types.InstanceCreateOpts) (
 		WithField("pool", opts.PoolName).
 		WithField("region", p.region).
 		WithField("image", p.image).
-		WithField("size", p.size)
+		WithField("size", p.size).
+		WithField("hibernate", p.CanHibernate())
 	var name = fmt.Sprintf(opts.RunnerName+"-"+opts.PoolName+"-%d", time.Now().Unix())
 
 	var tags = map[string]string{

internal/drivers/manager.go

@@ -2,17 +2,19 @@ package drivers
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"runtime/debug"
 	"sort"
 	"sync"
 	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/drone-runners/drone-runner-aws/internal/certs"
 	"github.com/drone-runners/drone-runner-aws/store"
 	"github.com/drone-runners/drone-runner-aws/types"
 	"github.com/drone/runner-go/logger"
+	lehttp "github.com/harness/lite-engine/cli/client"
+	"github.com/pkg/errors"
 
 	"github.com/sirupsen/logrus"
 )
@@ -438,7 +440,7 @@ func (m *Manager) buildPool(ctx context.Context, pool *poolEntry) error {
 				Infoln("build pool: created new instance")
 
 			go func() {
-				herr := m.hibernate(ctx, pool.Name, inst.ID)
+				herr := m.hibernate(context.Background(), pool.Name, inst.ID)
 				if herr != nil {
 					logr.WithError(herr).Errorln("failed to hibernate the vm")
 				}
@@ -540,8 +542,11 @@ func (m *Manager) hibernate(ctx context.Context, poolName, instanceID string) er
 		return nil
 	}
 
-	// TODO: Use health check to determine wait time before hibernating the instance
-	time.Sleep(600 * time.Second) // nolint:gomnd
+	if !pool.Driver.CanHibernate() {
+		return nil
+	}
+
+	m.waitForInstanceConnectivity(ctx, instanceID)
 
 	pool.Lock()
 	defer pool.Unlock()
@@ -577,3 +582,53 @@ func (m *Manager) forEach(ctx context.Context, f func(ctx context.Context, pool
 
 	return nil
 }
+
+func (m *Manager) waitForInstanceConnectivity(ctx context.Context, instanceID string) {
+	bf := backoff.NewExponentialBackOff()
+	for {
+		duration := bf.NextBackOff()
+		if duration == bf.Stop {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			logrus.WithField("instanceID", instanceID).Warnln("hibernate: connectivity check deadline exceeded")
+			return
+		case <-time.After(duration):
+			err := m.checkInstanceConnectivity(ctx, instanceID)
+			if err == nil {
+				return
+			}
+			logrus.WithError(err).WithField("instanceID", instanceID).Traceln("hibernate: instance connectivity check failed")
+		}
+	}
+}
+
+func (m *Manager) checkInstanceConnectivity(ctx context.Context, instanceID string) error {
+	instance, err := m.Find(ctx, instanceID)
+	if err != nil {
+		return errors.Wrap(err, "failed to find the instance in db")
+	}
+
+	if instance.Address == "" {
+		return errors.New("instance has not received IP address")
+	}
+
+	endpoint := fmt.Sprintf("https://%s:9079/", instance.Address)
+	client, err := lehttp.NewHTTPClient(endpoint, m.runnerName, string(instance.CACert), string(instance.TLSCert), string(instance.TLSKey))
+	if err != nil {
+		return errors.Wrap(err, "failed to create client")
+	}
+
+	response, err := client.Health(ctx)
+	if err != nil {
+		return err
+	}
+
+	if !response.OK {
+		return errors.New("health check call failed")
+	}
+
+	return nil
+}

internal/drivers/pool.go

@@ -30,6 +30,7 @@ type Driver interface {
 	Hibernate(ctx context.Context, instanceID, poolName string) error
 	Start(ctx context.Context, instanceID, poolName string) (ipAddress string, err error)
 	Ping(ctx context.Context) error
+	CanHibernate() bool
 
 	RootDir() string
 	ProviderName() string