feat: [PIPE-23578]: Refactor runner to reuse VM tasks (#518)

drone-runners · Nov 21, 2024 · 87f3402 · 87f3402
1 parent df827d0
commit 87f3402
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/command/harness/delegate/delegate.go b/command/harness/delegate/delegate.go
@@ -239,7 +239,8 @@ func (c *delegateCommand) handleDestroy(w http.ResponseWriter, r *http.Request)
 	req.Context.TaskID = rs.CorrelationID
 
 	ctx := r.Context()
-	err := harness.HandleDestroy(ctx, req, c.stageOwnerStore, &c.env, c.poolManager, c.metrics)
+	err := harness.HandleDestroy(ctx, req, c.stageOwnerStore, c.env.LiteEngine.EnableMock,
+		c.env.LiteEngine.MockStepTimeoutSecs, c.poolManager, c.metrics)
 	if err != nil {
 		logrus.WithField("stage_runtime_id", req.StageRuntimeID).WithField("task_id", rs.CorrelationID).WithError(err).Error("could not destroy VM")
 		writeError(w, err)

diff --git a/command/harness/destroy.go b/command/harness/destroy.go
@@ -6,12 +6,11 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/drone-runners/drone-runner-aws/command/config"
-	"github.com/drone-runners/drone-runner-aws/command/harness/storage"
 	"github.com/drone-runners/drone-runner-aws/app/drivers"
 	"github.com/drone-runners/drone-runner-aws/app/lehelper"
 	"github.com/drone-runners/drone-runner-aws/app/oshelp"
 	ierrors "github.com/drone-runners/drone-runner-aws/app/types"
+	"github.com/drone-runners/drone-runner-aws/command/harness/storage"
 	"github.com/drone-runners/drone-runner-aws/metric"
 	"github.com/drone-runners/drone-runner-aws/store"
 	"github.com/drone-runners/drone-runner-aws/types"
@@ -35,7 +34,15 @@ type VMCleanupRequest struct {
 	StorageCleanupType storage.CleanupType `json:"storage_cleanup_type,omitempty"`
 }
 
-func HandleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerStore, env *config.EnvConfig, poolManager drivers.IManager, metrics *metric.Metrics) error {
+func HandleDestroy(
+	ctx context.Context,
+	r *VMCleanupRequest,
+	s store.StageOwnerStore,
+	enableMock bool, // only used for scale testing
+	mockTimeout int, // only used for scale testing
+	poolManager drivers.IManager,
+	metrics *metric.Metrics,
+) error {
 	if r.StageRuntimeID == "" {
 		return ierrors.NewBadRequestError("mandatory field 'stage_runtime_id' in the request body is empty")
 	}
@@ -72,7 +79,7 @@ func HandleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerS
 			}
 			return ctx.Err()
 		case <-timer.C:
-			_, err := handleDestroy(ctx, r, s, env, poolManager, metrics, cnt, logr)
+			_, err := handleDestroy(ctx, r, s, enableMock, mockTimeout, poolManager, metrics, cnt, logr)
 			if err != nil {
 				if lastErr == nil || (lastErr.Error() != err.Error()) {
 					logr.WithError(err).Errorln("could not destroy VM")
@@ -89,7 +96,7 @@ func HandleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerS
 	}
 }
 
-func handleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerStore, env *config.EnvConfig,
+func handleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerStore, enableMock bool, mockTimeout int,
 	poolManager drivers.IManager, metrics *metric.Metrics, retryCount int, logr *logrus.Entry) (*types.Instance, error) {
 	logr = logr.WithField("retry_count", retryCount)
 	entity, err := s.Find(ctx, r.StageRuntimeID)
@@ -116,7 +123,7 @@ func handleDestroy(ctx context.Context, r *VMCleanupRequest, s store.StageOwnerS
 		WithField("instance_name", inst.Name)
 
 	logr.Traceln("invoking lite engine cleanup")
-	client, err := lehelper.GetClient(inst, poolManager.GetTLSServerName(), inst.Port, env.LiteEngine.EnableMock, env.LiteEngine.MockStepTimeoutSecs)
+	client, err := lehelper.GetClient(inst, poolManager.GetTLSServerName(), inst.Port, enableMock, mockTimeout)
 	if err != nil {
 		logr.WithError(err).Errorln("could not create lite engine client for invoking cleanup")
 	} else {

diff --git a/command/harness/dlite/cleanup.go b/command/harness/dlite/cleanup.go
@@ -47,7 +47,8 @@ func (t *VMCleanupTask) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if !req.Distributed {
 		harness.GetCtxState().Delete(req.StageRuntimeID)
 	}
-	err = harness.HandleDestroy(ctx, req, poolManager.GetStageOwnerStore(), &t.c.env, poolManager, t.c.metrics)
+	err = harness.HandleDestroy(ctx, req, poolManager.GetStageOwnerStore(),
+		t.c.env.LiteEngine.EnableMock, t.c.env.LiteEngine.MockStepTimeoutSecs, poolManager, t.c.metrics)
 	if err != nil {
 		t.c.metrics.ErrorCount.WithLabelValues(accountID, strconv.FormatBool(req.Distributed)).Inc()
 		logr.WithError(err).WithField("account_id", accountID).Error("could not destroy VM")