fix(plans): lock the plan before completing a target-met step

internal/config/store_postgres_managed_target_met_integration_test.go

@@ -0,0 +1,94 @@
+//go:build integration
+
+package config
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/require"
+)
+
+// managedTargetMetRefusal names what a refusal may carry: the
+// unavailable-root sentinel, or the duplicate-root conflict its resolution is
+// built on. An unrelated failure must not pass for a refusal.
+func managedTargetMetRefusal(err error) bool {
+	return errors.Is(err, ErrExecutionNotInExpectedStatus) || errors.Is(err, ErrPurchasePlanConflict)
+}
+
+func TestManagedTargetMetRefusesBeforeWrite(t *testing.T) {
+	for _, which := range []string{"missing", "foreign", "finished", "wrong step", "occupied", "scheduled", "duplicate", "legacy foreign", "legacy other step", "beside root", "other step root"} {
+		t.Run(which, func(t *testing.T) {
+			ctx := t.Context()
+			f := managedClockFixture(ctx, t, 1)
+			ids := []string{f.planID}
+			step, rootID := 2, ""
+			if which != "missing" {
+				rootID = f.row(ctx, t, 2, nil, "pending").ExecutionID
+			}
+			switch which {
+			case "foreign", "legacy foreign":
+				if which == "legacy foreign" {
+					_, err := f.store.db.Exec(ctx, `UPDATE purchase_plans SET ramp_schedule=jsonb_set(ramp_schedule,'{dates_managed}','false') WHERE id=$1`, f.planID)
+					require.NoError(t, err)
+				}
+				other, err := f.store.GetPurchasePlan(ctx, f.planID)
+				require.NoError(t, err)
+				other.ID, other.Name = uuid.NewString(), "foreign target-met fixture"
+				other.RampSchedule.CurrentStep, other.RampSchedule.DatesManaged = 0, false
+				require.NoError(t, f.store.CreatePurchasePlan(ctx, other))
+				ids = append(ids, other.ID)
+				rootID = uuid.NewString()
+				require.NoError(t, f.store.SavePurchaseExecution(ctx, &PurchaseExecution{PlanID: other.ID, ExecutionID: rootID,
+					Status: "pending", StepNumber: 2, ScheduledDate: other.RampSchedule.GetNextPurchaseDate()}))
+				var plans, rows int
+				require.NoError(t, f.store.db.QueryRow(ctx, `SELECT
+					(SELECT count(*) FROM purchase_plans WHERE id::text=ANY($1::text[])),
+					(SELECT count(*) FROM purchase_executions WHERE plan_id::text=ANY($1::text[]))`, ids).Scan(&plans, &rows))
+				require.Equal(t, 2, plans)
+				require.Equal(t, 2, rows)
+			case "legacy other step":
+				_, err := f.store.db.Exec(ctx, `UPDATE purchase_plans SET ramp_schedule=jsonb_set(ramp_schedule,'{dates_managed}','false') WHERE id=$1`, f.planID)
+				require.NoError(t, err)
+				rootID = f.row(ctx, t, 3, nil, "pending").ExecutionID
+			case "beside root":
+				rootID = ""
+			case "other step root":
+				step, rootID = 3, f.row(ctx, t, 3, nil, "pending").ExecutionID
+			case "finished":
+				_, err := f.store.db.Exec(ctx, `UPDATE purchase_plans SET ramp_schedule=jsonb_set(ramp_schedule,'{total_steps}','1') WHERE id=$1`, f.planID)
+				require.NoError(t, err)
+			case "wrong step":
+				step = 3
+			case "occupied":
+				require.NoError(t, f.store.SavePurchaseExecution(ctx, &PurchaseExecution{PlanID: f.planID, ExecutionID: rootID,
+					Status: "pending", StepNumber: 2, Recommendations: []RecommendationRecord{baselineRec(uuid.NewString(), "m5.large", 1)}}))
+			case "scheduled":
+				_, err := f.store.db.Exec(ctx, `UPDATE purchase_executions SET status='scheduled',scheduled_execution_at=NOW()+INTERVAL '48 days' WHERE execution_id=$1`, rootID)
+				require.NoError(t, err)
+			case "duplicate":
+				f.row(ctx, t, 2, nil, "pending")
+			}
+			query := `SELECT jsonb_build_object('plans',
+				(SELECT jsonb_agg(to_jsonb(p) ORDER BY p.id) FROM purchase_plans p WHERE p.id::text=ANY($1::text[])),
+				'executions',(SELECT jsonb_agg(to_jsonb(e) ORDER BY e.execution_id) FROM purchase_executions e WHERE e.plan_id::text=ANY($1::text[])))`
+			snapshot := func(row pgx.Row) []byte {
+				var data []byte
+				require.NoError(t, row.Scan(&data))
+				return data
+			}
+			before := snapshot(f.store.db.QueryRow(ctx, query, ids))
+			err := f.store.WithTx(ctx, func(tx pgx.Tx) error {
+				refused := f.store.CompleteTargetMetStepTx(ctx, tx, f.planID, step, rootID, "fixture target met")
+				require.Equal(t, before, snapshot(tx.QueryRow(ctx, query, ids)), "managed target-met must refuse before any write")
+				require.True(t, managedTargetMetRefusal(refused), "managed target-met must refuse an unavailable root, got %v", refused)
+				return refused
+			})
+			require.Error(t, err)
+			require.Equal(t, before, snapshot(f.store.db.QueryRow(ctx, query, ids)))
+			require.Equal(t, 1, f.currentStep(ctx, t), "a refused managed target-met must leave the ramp where it was")
+		})
+	}
+}

internal/config/store_postgres_plan_schedule_clock.go

@@ -93,6 +93,34 @@ func (s *PostgresStore) managedPlanNextRootTx(ctx context.Context, tx pgx.Tx, pl
 	return root, nil
 }
 
+func (s *PostgresStore) lockTargetMetPlanTx(ctx context.Context, tx pgx.Tx, planID, rootID string) (*PurchasePlan, error) {
+	plan, err := s.LockPurchasePlanTx(ctx, tx, planID)
+	if err != nil {
+		return nil, err
+	}
+	if plan == nil {
+		return nil, fmt.Errorf("%w: purchase plan %s", ErrNotFound, planID)
+	}
+	if err := s.requireManagedTargetMetRootTx(ctx, tx, plan, rootID); err != nil {
+		return nil, err
+	}
+	return plan, nil
+}
+
+func (s *PostgresStore) requireManagedTargetMetRootTx(ctx context.Context, tx pgx.Tx, plan *PurchasePlan, rootID string) error {
+	if !plan.RampSchedule.DatesManaged {
+		return nil
+	}
+	root, err := s.ResolveManagedPlanRootTx(ctx, tx, plan)
+	if err != nil {
+		return err
+	}
+	if root == nil || root.ExecutionID != rootID {
+		return fmt.Errorf("%w: managed target-met root is unavailable", ErrExecutionNotInExpectedStatus)
+	}
+	return nil
+}
+
 func managedPlanNextRoot(plan *PurchasePlan, rows []PurchaseExecution) (*PurchaseExecution, error) {
 	var root *PurchaseExecution
 	blocked := false

internal/config/store_postgres_ramp_target_met.go

@@ -36,6 +36,10 @@ func (s *PostgresStore) CompleteTargetMetStepTx(ctx context.Context, tx pgx.Tx,
 	if stepNumber <= 0 {
 		return fmt.Errorf("cannot complete ramp step %d of plan %s: step numbers are 1-based", stepNumber, planID)
 	}
+	plan, err := s.lockTargetMetPlanTx(ctx, tx, planID, rootID)
+	if err != nil {
+		return err
+	}
 	var unknown int
 	if err := tx.QueryRow(ctx, `
 		SELECT count(*) FROM purchase_executions
@@ -60,23 +64,16 @@ func (s *PostgresStore) CompleteTargetMetStepTx(ctx context.Context, tx pgx.Tx,
 		tag, err := tx.Exec(ctx, `
 			UPDATE purchase_executions
 			   SET status = $2, completed_at = $3, error = $4, pre_approval_skip_reason = $5
-			 WHERE execution_id = $1 AND status IN ('pending', 'notified')
+			 WHERE execution_id = $1 AND plan_id = $6 AND step_number = $7 AND status IN ('pending', 'notified')
 			   AND CASE WHEN jsonb_typeof(recommendations) = 'array' THEN jsonb_array_length(recommendations) ELSE 0 END = 0`,
-			rootID, StatusCompleted, now, note, PreApprovalSkipRampTargetMet)
+			rootID, StatusCompleted, now, note, PreApprovalSkipRampTargetMet, planID, stepNumber)
 		if err != nil {
 			return fmt.Errorf("failed to complete target-met root row %s: %w", rootID, err)
 		}
 		if tag.RowsAffected() != 1 {
 			return fmt.Errorf("%w: root row %s", ErrExecutionNotInExpectedStatus, rootID)
 		}
 	}
-	plan, err := s.LockPurchasePlanTx(ctx, tx, planID)
-	if err != nil {
-		return err
-	}
-	if plan == nil {
-		return fmt.Errorf("%w: purchase plan %s", ErrNotFound, planID)
-	}
 	if err := advanceRampStep(ctx, tx, plan, stepNumber); err != nil {
 		return err
 	}