Add polling to wait for post plan tasks to complete in remote backend

Author: Maed223

internal/backend/remote/backend_common.go

@@ -47,6 +47,33 @@ func backoff(min, max float64, iter int) time.Duration {
 	return time.Duration(backoff) * time.Millisecond
 }
 
+func (b *Remote) waitForPostPlanTasks(stopCtx, cancelCtx context.Context, op *backendrun.Operation, r *tfe.Run) error {
+	taskStages := make(taskStages, 0)
+	result, err := b.client.Runs.ReadWithOptions(stopCtx, r.ID, &tfe.RunReadOptions{
+		Include: []tfe.RunIncludeOpt{tfe.RunTaskStages},
+	})
+	if err == nil {
+		for _, t := range result.TaskStages {
+			if t != nil {
+				taskStages[t.Stage] = t
+			}
+		}
+	} else {
+		// This error would be expected for older versions of TFE that do not allow
+		// fetching task_stages.
+		if !strings.HasSuffix(err.Error(), "Invalid include parameter") {
+			generalError("Failed to retrieve run", err)
+		}
+	}
+
+	if stage, ok := taskStages[tfe.PostPlan]; ok {
+		if err := b.waitTaskStage(stopCtx, cancelCtx, op, r, stage.ID); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (b *Remote) waitForRun(stopCtx, cancelCtx context.Context, op *backendrun.Operation, opType string, r *tfe.Run, w *tfe.Workspace) (*tfe.Run, error) {
 	started := time.Now()
 	updated := started

internal/backend/remote/backend_plan.go

@@ -411,6 +411,15 @@ in order to capture the filesystem context the remote workspace expects:
 		return r, generalError("Failed to retrieve run", err)
 	}
 
+	// Wait for post plan tasks to complete before proceeding.
+	// Otherwise, in the case of an apply, if they are still running
+	// when we check for whether the run is confirmable the CLI will
+	// uncermoniously exit before the user has a chance to confirm, or for an auto-apply to take place.
+	err = b.waitForPostPlanTasks(stopCtx, cancelCtx, op, r)
+	if err != nil {
+		return r, err
+	}
+
 	// If the run is canceled or errored, we still continue to the
 	// cost-estimation and policy check phases to ensure we render any
 	// results available. In the case of a hard-failed policy check, the

internal/backend/remote/backend_taskStages.go

@@ -0,0 +1,52 @@
+package remote
+
+import (
+	"context"
+	"fmt"
+
+	tfe "github.com/hashicorp/go-tfe"
+	"github.com/hashicorp/terraform/internal/backend/backendrun"
+)
+
+type taskStages map[tfe.Stage]*tfe.TaskStage
+
+const (
+	taskStageBackoffMin = 4000.0
+	taskStageBackoffMax = 12000.0
+)
+
+// waitTaskStage waits for a task stage to complete, only informs the caller if the stage has failed in some way.
+func (b *Remote) waitTaskStage(stopCtx, cancelCtx context.Context, op *backendrun.Operation, r *tfe.Run, stageID string) error {
+	ctx := &IntegrationContext{
+		StopContext:   stopCtx,
+		CancelContext: cancelCtx,
+	}
+	return ctx.Poll(taskStageBackoffMin, taskStageBackoffMax, func(i int) (bool, error) {
+		options := tfe.TaskStageReadOptions{
+			Include: []tfe.TaskStageIncludeOpt{tfe.TaskStageTaskResults, tfe.PolicyEvaluationsTaskResults},
+		}
+		stage, err := b.client.TaskStages.Read(ctx.StopContext, stageID, &options)
+		if err != nil {
+			return false, generalError("Failed to retrieve task stage", err)
+		}
+
+		switch stage.Status {
+		case tfe.TaskStagePending:
+			// Waiting for it to start
+			return true, nil
+		case tfe.TaskStageRunning:
+			// not a terminal status so we continue to poll
+			return true, nil
+		case tfe.TaskStagePassed:
+			return false, nil
+		case tfe.TaskStageCanceled, tfe.TaskStageErrored, tfe.TaskStageFailed:
+			return false, fmt.Errorf("Task Stage '%s': %s.", stage.ID, stage.Status)
+		case tfe.TaskStageAwaitingOverride:
+			return false, fmt.Errorf("Task Stage '%s' awaiting override.", stage.ID, stage.Status)
+		case tfe.TaskStageUnreachable:
+			return false, nil
+		default:
+			return false, fmt.Errorf("Task stage '%s' has invalid status: %s", stage.ID, stage.Status)
+		}
+	})
+}

internal/backend/remote/cloud_integration.go

@@ -0,0 +1,36 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package remote
+
+import (
+	"context"
+	"log"
+	"time"
+)
+
+// IntegrationContext is a set of data that is useful when performing HCP Terraform integration operations
+type IntegrationContext struct {
+	StopContext   context.Context
+	CancelContext context.Context
+}
+
+func (s *IntegrationContext) Poll(backoffMinInterval float64, backoffMaxInterval float64, every func(i int) (bool, error)) error {
+	for i := 0; ; i++ {
+		select {
+		case <-s.StopContext.Done():
+			log.Print("IntegrationContext.Poll: StopContext.Done() called")
+			return s.StopContext.Err()
+		case <-s.CancelContext.Done():
+			log.Print("IntegrationContext.Poll: CancelContext.Done() called")
+			return s.CancelContext.Err()
+		case <-time.After(backoff(backoffMinInterval, backoffMaxInterval, i)):
+			// blocks for a time between min and max
+		}
+
+		cont, err := every(i)
+		if !cont {
+			return err
+		}
+	}
+}