fix(purchase): preserve ramp progress across stale plan updates

internal/analytics/collector_test.go

@@ -160,6 +160,10 @@ func (m *mockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P
 	return nil
 }
 
+func (m *mockConfigStore) StampPlanNotificationSent(context.Context, string, time.Time) error {
+	return nil
+}
+
 func (m *mockConfigStore) CompletePlanStep(ctx context.Context, planID string, stepNumber int) error {
 	return nil
 }

internal/api/handler_plans.go

@@ -112,14 +112,14 @@ func calculateNextExecutionDate(plan *config.PurchasePlan, now time.Time) *time.
 	return &nextDate
 }
 
-// mapCreatePlanStorageError converts a storage-layer error from createPlan into
-// the appropriate HTTP ClientError. If err is ErrNotFound it returns a 404
-// with notFoundMsg; otherwise it logs the supplied format string at ERROR level
-// and returns a 500 with genericMsg.
+// mapCreatePlanStorageError keeps storage details out of API responses.
 func mapCreatePlanStorageError(err error, notFoundMsg, genericMsg, logFmt string, logArgs ...any) error {
 	if errors.Is(err, config.ErrNotFound) {
 		return NewClientError(http.StatusNotFound, notFoundMsg)
 	}
+	if errors.Is(err, config.ErrPurchasePlanConflict) {
+		return NewClientError(http.StatusConflict, "plan changed; refresh the plan and try again")
+	}
 	logging.Errorf(logFmt, logArgs...)
 	return NewClientError(http.StatusInternalServerError, genericMsg)
 }
@@ -289,7 +289,7 @@ func (h *Handler) updatePlan(ctx context.Context, httpReq *events.LambdaFunction
 
 	// Preserve timestamps from existing plan
 	plan.CreatedAt = existingPlan.CreatedAt
-	plan.UpdatedAt = time.Now()
+	plan.UpdatedAt = existingPlan.UpdatedAt
 
 	// toPurchasePlan always rebuilds the ramp schedule from scratch
 	// (CurrentStep=0, StartDate=now) and recomputes NextExecutionDate,
@@ -309,7 +309,7 @@ func (h *Handler) updatePlan(ctx context.Context, httpReq *events.LambdaFunction
 	}
 
 	if err := h.config.UpdatePurchasePlan(ctx, plan); err != nil {
-		return nil, err
+		return nil, mapCreatePlanStorageError(err, "plan not found", "failed to update plan", "updatePlan: UpdatePurchasePlan failed")
 	}
 
 	return plan, nil
@@ -605,7 +605,6 @@ func (h *Handler) createPurchaseExecutionsTx(ctx context.Context, tx pgx.Tx, pla
 func (h *Handler) updatePlanNextExecutionDateTx(ctx context.Context, tx pgx.Tx, plan *config.PurchasePlan, startDate time.Time) error {
 	if plan.NextExecutionDate == nil || plan.NextExecutionDate.After(startDate) {
 		plan.NextExecutionDate = &startDate
-		plan.UpdatedAt = time.Now()
 		if err := h.config.UpdatePurchasePlanTx(ctx, tx, plan); err != nil {
 			return fmt.Errorf("failed to update plan: %w", err)
 		}
@@ -687,14 +686,12 @@ func (h *Handler) patchPlan(ctx context.Context, httpReq *events.LambdaFunctionU
 		return nil, err
 	}
 
-	plan.UpdatedAt = time.Now()
-
 	if err := plan.Validate(); err != nil {
 		return nil, NewClientError(400, fmt.Sprintf("validation error: %s", err))
 	}
 
 	if err := h.config.UpdatePurchasePlan(ctx, plan); err != nil {
-		return nil, err
+		return nil, mapCreatePlanStorageError(err, "plan not found", "failed to update plan", "patchPlan: UpdatePurchasePlan failed")
 	}
 
 	return plan, nil

internal/api/handler_plans_edit_race_integration_test.go

@@ -0,0 +1,96 @@
+//go:build integration
+
+package api
+
+import (
+	"context"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/LeanerCloud/cloud-commitments-platform/internal/config"
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+type planEditReadBarrier struct {
+	*config.PostgresStore
+	afterRead func()
+}
+
+func (s *planEditReadBarrier) GetPurchasePlan(ctx context.Context, id string) (*config.PurchasePlan, error) {
+	plan, err := s.PostgresStore.GetPurchasePlan(ctx, id)
+	if err == nil && s.afterRead != nil {
+		afterRead := s.afterRead
+		s.afterRead = nil
+		afterRead()
+	}
+	return plan, err
+}
+
+func TestPlanEditsPreserveConcurrentRampCompletion(t *testing.T) {
+	for _, method := range []string{"PUT", "PATCH", "disable"} {
+		for _, race := range []bool{false, true} {
+			name := method + "/fresh"
+			if race {
+				name = method + "/stale"
+			}
+			t.Run(name, func(t *testing.T) {
+				ctx := context.Background()
+				f := newCreateConcurrencyFixture(ctx, t)
+				plan, err := f.store.GetPurchasePlan(ctx, f.planID)
+				require.NoError(t, err)
+				plan.Enabled = true
+				plan.Services = map[string]config.ServiceConfig{"ec2": {Provider: "aws", Service: "ec2"}}
+				plan.RampSchedule.Type = "custom"
+				next := time.Now().Add(24 * time.Hour).Truncate(time.Microsecond)
+				plan.NextExecutionDate = &next
+				require.NoError(t, f.store.UpdatePurchasePlan(ctx, plan))
+				require.NoError(t, f.store.SavePurchaseExecution(ctx, &config.PurchaseExecution{
+					ExecutionID: uuid.NewString(), PlanID: f.planID, Status: "completed", StepNumber: 3, ScheduledDate: next,
+				}))
+				barrier := &planEditReadBarrier{PostgresStore: f.store}
+				var completed *config.PurchasePlan
+				if race {
+					barrier.afterRead = func() {
+						require.NoError(t, f.store.CompletePlanStep(ctx, f.planID, 3))
+						completed, err = f.store.GetPurchasePlan(ctx, f.planID)
+						require.NoError(t, err)
+						require.Equal(t, 3, completed.RampSchedule.CurrentStep)
+					}
+				}
+				f.handler.config = barrier
+				req := &events.LambdaFunctionURLRequest{
+					Headers: map[string]string{"Authorization": "Bearer admin-token", "Content-Type": "application/json"},
+					Body:    `{"name":"renamed","enabled":true,"ramp_schedule":"custom","custom_step_percent":25,"custom_interval_days":7}`,
+				}
+				switch method {
+				case "PUT":
+					_, err = f.handler.updatePlan(ctx, req, f.planID)
+				case "PATCH":
+					req.Body = `{"name":"renamed"}`
+					_, err = f.handler.patchPlan(ctx, req, f.planID)
+				case "disable":
+					err = f.handler.disablePlan(ctx, f.planID)
+				}
+				after, getErr := f.store.GetPurchasePlan(ctx, f.planID)
+				require.NoError(t, getErr)
+				if race {
+					clientErr, ok := IsClientError(err)
+					require.True(t, ok, "stale edit must return a client error, got %v", err)
+					require.Equal(t, http.StatusConflict, clientErr.code)
+					require.Equal(t, completed, after, "the winning completion must remain unchanged")
+					return
+				}
+				require.NoError(t, err)
+				if method == "disable" {
+					require.False(t, after.Enabled)
+				} else {
+					require.Equal(t, "renamed", after.Name)
+				}
+				require.Equal(t, 2, after.RampSchedule.CurrentStep)
+			})
+		}
+	}
+}

internal/config/errors.go

@@ -5,6 +5,9 @@ import "errors"
 // ErrNotFound is returned when a requested config-store row does not exist.
 var ErrNotFound = errors.New("not found")
 
+// ErrPurchasePlanConflict means a plan changed after the caller read it.
+var ErrPurchasePlanConflict = errors.New("purchase plan changed")
+
 // ErrExecutionNotInExpectedStatus is returned by TransitionExecutionStatus
 // when the target execution exists but its current status is not in the
 // allowed `fromStatuses` set -- i.e. the atomic CAS rejected because some

internal/config/interfaces.go

@@ -29,7 +29,9 @@ type StoreInterface interface {
 	// Purchase plans
 	CreatePurchasePlan(ctx context.Context, plan *PurchasePlan) error
 	GetPurchasePlan(ctx context.Context, planID string) (*PurchasePlan, error)
+	// UpdatedAt must be the read version; a stale version returns ErrPurchasePlanConflict.
 	UpdatePurchasePlan(ctx context.Context, plan *PurchasePlan) error
+	StampPlanNotificationSent(ctx context.Context, planID string, at time.Time) error
 	// CompletePlanStep records that ramp step stepNumber finished and advances
 	// the plan's schedule to it, inside a SELECT FOR UPDATE transaction that
 	// prevents the concurrent-write lost-update race of issue #1071.
@@ -77,11 +79,8 @@ type StoreInterface interface {
 	// transaction; without it the answer is advisory and two concurrent callers
 	// both see the step free.
 	OccupiedRampStepsInRangeTx(ctx context.Context, tx pgx.Tx, planID string, from, to int) ([]int, error)
-	// UpdatePurchasePlanTx is the tx-accepting variant of UpdatePurchasePlan.
-	// Used from createPlannedPurchases' WithTx block so the per-row
-	// SavePurchaseExecutionTx writes and the plan's next_execution_date
-	// bump commit atomically — a partial failure leaves no orphaned
-	// rows and no stale plan pointer.
+	// UpdatePurchasePlanTx checks UpdatedAt and replaces it with the returned version.
+	// The caller must discard the transaction-local plan if the transaction rolls back.
 	UpdatePurchasePlanTx(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) error
 	DeletePurchasePlan(ctx context.Context, planID string) error
 	ListPurchasePlans(ctx context.Context, filter PurchasePlanFilter) ([]PurchasePlan, error)

internal/config/store_postgres.go

@@ -516,9 +516,11 @@ func (s *PostgresStore) CreatePurchasePlan(ctx context.Context, plan *PurchasePl
 			services, ramp_schedule, created_at, updated_at,
 			next_execution_date, last_execution_date, last_notification_sent
 		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+		RETURNING updated_at
 	`
 
-	_, err = s.db.Exec(ctx, query,
+	var updatedAt time.Time
+	err = s.db.QueryRow(ctx, query,
 		plan.ID,
 		plan.Name,
 		plan.Enabled,
@@ -531,11 +533,12 @@ func (s *PostgresStore) CreatePurchasePlan(ctx context.Context, plan *PurchasePl
 		plan.NextExecutionDate,
 		plan.LastExecutionDate,
 		plan.LastNotificationSent,
-	)
+	).Scan(&updatedAt)
 
 	if err != nil {
 		return fmt.Errorf("failed to create purchase plan: %w", err)
 	}
+	plan.UpdatedAt = updatedAt
 
 	return nil
 }
@@ -715,39 +718,24 @@ func (s *PostgresStore) persistRampPosition(ctx context.Context, tx pgx.Tx, plan
 
 	now := time.Now()
 	plan.LastExecutionDate = &now
-	// Refresh updated_at on every advance. The plan was read from the DB
-	// with its previous UpdatedAt, so UpdatePurchasePlanTx's zero-value
-	// guard would otherwise persist a stale updated_at timestamp.
-	plan.UpdatedAt = now
 
 	return s.UpdatePurchasePlanTx(ctx, tx, plan)
 }
 
-// UpdatePurchasePlan updates an existing purchase plan. Delegates to
-// UpdatePurchasePlanTx inside a single-call WithTx — keeps the public
-// surface unchanged for callers that don't need to bundle this with
-// other writes, while sharing the SQL with the Tx variant. UpdatedAt
-// is stamped here (before the WithTx call) so existing tests that
-// inspect plan.UpdatedAt without exercising the DB still see it set —
-// see TestPostgresStore_UpdatePurchasePlan_NilDB.
+// UpdatePurchasePlan publishes the new version only after the transaction commits.
 func (s *PostgresStore) UpdatePurchasePlan(ctx context.Context, plan *PurchasePlan) error {
-	plan.UpdatedAt = time.Now()
-	return s.WithTx(ctx, func(tx pgx.Tx) error {
-		return s.UpdatePurchasePlanTx(ctx, tx, plan)
-	})
+	pending := *plan
+	if err := s.WithTx(ctx, func(tx pgx.Tx) error {
+		return s.UpdatePurchasePlanTx(ctx, tx, &pending)
+	}); err != nil {
+		return err
+	}
+	plan.UpdatedAt = pending.UpdatedAt
+	return nil
 }
 
-// UpdatePurchasePlanTx is the tx-accepting variant of UpdatePurchasePlan.
-// Used by createPlannedPurchases so the per-row execution inserts and
-// the plan's next_execution_date bump commit atomically — see the
-// interface doc for the partial-failure rationale. Callers that need
-// the auto-stamp of UpdatedAt either use UpdatePurchasePlan (which
-// stamps before WithTx) or stamp it themselves before calling Tx.
+// UpdatePurchasePlanTx rejects snapshots that lost a race with another plan writer.
 func (s *PostgresStore) UpdatePurchasePlanTx(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) error {
-	if plan.UpdatedAt.IsZero() {
-		plan.UpdatedAt = time.Now()
-	}
-
 	// Marshal services and ramp_schedule to JSONB
 	servicesJSON, err := json.Marshal(plan.Services)
 	if err != nil {
@@ -767,14 +755,15 @@ func (s *PostgresStore) UpdatePurchasePlanTx(ctx context.Context, tx pgx.Tx, pla
 			notification_days_before = $5,
 			services = $6,
 			ramp_schedule = $7,
-			updated_at = $8,
 			next_execution_date = $9,
 			last_execution_date = $10,
 			last_notification_sent = $11
-		WHERE id = $1
+		WHERE id = $1 AND updated_at = $8
+		RETURNING updated_at
 	`
 
-	result, err := tx.Exec(ctx, query,
+	var updatedAt time.Time
+	err = tx.QueryRow(ctx, query,
 		plan.ID,
 		plan.Name,
 		plan.Enabled,
@@ -786,16 +775,34 @@ func (s *PostgresStore) UpdatePurchasePlanTx(ctx context.Context, tx pgx.Tx, pla
 		plan.NextExecutionDate,
 		plan.LastExecutionDate,
 		plan.LastNotificationSent,
-	)
+	).Scan(&updatedAt)
 
+	if errors.Is(err, pgx.ErrNoRows) {
+		var exists bool
+		if probeErr := tx.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM purchase_plans WHERE id = $1)`, plan.ID).Scan(&exists); probeErr != nil {
+			return fmt.Errorf("failed to check purchase plan after conflict: %w", probeErr)
+		}
+		if !exists {
+			return fmt.Errorf("%w: purchase plan %s", ErrNotFound, plan.ID)
+		}
+		return ErrPurchasePlanConflict
+	}
 	if err != nil {
 		return fmt.Errorf("failed to update purchase plan: %w", err)
 	}
+	plan.UpdatedAt = updatedAt
+	return nil
+}
 
+// StampPlanNotificationSent never writes a snapshot of the plan's ramp state.
+func (s *PostgresStore) StampPlanNotificationSent(ctx context.Context, planID string, at time.Time) error {
+	result, err := s.db.Exec(ctx, `UPDATE purchase_plans SET last_notification_sent = $2 WHERE id = $1`, planID, at)
+	if err != nil {
+		return fmt.Errorf("failed to stamp plan notification: %w", err)
+	}
 	if result.RowsAffected() == 0 {
-		return fmt.Errorf("purchase plan not found: %s", plan.ID)
+		return fmt.Errorf("%w: purchase plan %s", ErrNotFound, planID)
 	}
-
 	return nil
 }
 

internal/config/store_postgres_complete_step_test.go

@@ -87,29 +87,14 @@ func (a nullTimeArg) Match(v interface{}) bool {
 	return (tp == nil) == a.wantNil
 }
 
-// afterArg matches a time.Time UPDATE argument that is strictly after the given
-// instant, used to assert updated_at is refreshed rather than persisted stale.
-type afterArg struct{ notBefore time.Time }
-
-func (a afterArg) Match(v interface{}) bool {
-	ts, ok := v.(time.Time)
-	if !ok {
-		return false
-	}
-	return ts.After(a.notBefore)
-}
-
 const purchasePlanUpdateArgs = 11
 
-// completeStepUpdateArgs builds the WithArgs matcher list for the UPDATE issued
-// by CompletePlanStep, asserting the persisted CurrentStep at $7, a refreshed
-// updated_at at $8 (> staleUpdatedAt), and the next_execution_date presence at
-// $9 while leaving the rest as AnyArg.
+// completeStepUpdateArgs pins the completed step, expected read version, and next-date presence.
 func completeStepUpdateArgs(wantStep int, staleUpdatedAt time.Time, wantNextNil bool) []interface{} {
 	args := anyArgsCfg(purchasePlanUpdateArgs)
-	args[6] = rampStepArg{want: wantStep}         // ramp_schedule = $7
-	args[7] = afterArg{notBefore: staleUpdatedAt} // updated_at = $8
-	args[8] = nullTimeArg{wantNil: wantNextNil}   // next_execution_date = $9
+	args[6] = rampStepArg{want: wantStep}       // ramp_schedule = $7
+	args[7] = staleUpdatedAt                    // expected updated_at = $8
+	args[8] = nullTimeArg{wantNil: wantNextNil} // next_execution_date = $9
 	return args
 }
 
@@ -177,9 +162,9 @@ func TestPGXMock_CompletePlanStep_LocksAndAdvances(t *testing.T) {
 		WithArgs("plan-123").
 		WillReturnRows(rampPlanRows(t, "plan-123", ramp, now, stale, sql.NullTime{Valid: false}))
 	expectFanOut(mock, "plan-123", 2, 3, 0)
-	mock.ExpectExec(`UPDATE purchase_plans`).
+	mock.ExpectQuery(`UPDATE purchase_plans`).
 		WithArgs(completeStepUpdateArgs(2, stale, false)...).
-		WillReturnResult(pgxmock.NewResult("UPDATE", 1))
+		WillReturnRows(pgxmock.NewRows([]string{"updated_at"}).AddRow(now))
 	mock.ExpectCommit()
 
 	require.NoError(t, store.CompletePlanStep(ctx, "plan-123", 2))
@@ -372,9 +357,9 @@ func TestPGXMock_CompletePlanStep_CompletedRampClearsNextDate(t *testing.T) {
 		WillReturnRows(rampPlanRows(t, "plan-done", ramp, now, stale, sql.NullTime{Valid: true, Time: now}))
 	// Step stays at 4 (already complete), updated_at refreshed, and
 	// next_execution_date is cleared.
-	mock.ExpectExec(`UPDATE purchase_plans`).
+	mock.ExpectQuery(`UPDATE purchase_plans`).
 		WithArgs(completeStepUpdateArgs(4, stale, true)...).
-		WillReturnResult(pgxmock.NewResult("UPDATE", 1))
+		WillReturnRows(pgxmock.NewRows([]string{"updated_at"}).AddRow(now))
 	mock.ExpectCommit()
 
 	require.NoError(t, store.CompletePlanStep(ctx, "plan-done", 5))