diff --git a/internal/analytics/collector_test.go b/internal/analytics/collector_test.go index 1493d924..6870feb9 100644 --- a/internal/analytics/collector_test.go +++ b/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 } diff --git a/internal/api/handler_plans.go b/internal/api/handler_plans.go index c0139bbc..7cf5c09b 100644 --- a/internal/api/handler_plans.go +++ b/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 diff --git a/internal/api/handler_plans_edit_race_integration_test.go b/internal/api/handler_plans_edit_race_integration_test.go new file mode 100644 index 00000000..beb6bfeb --- /dev/null +++ b/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) + }) + } + } +} diff --git a/internal/config/errors.go b/internal/config/errors.go index f87ae214..f8ecf87e 100644 --- a/internal/config/errors.go +++ b/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 diff --git a/internal/config/interfaces.go b/internal/config/interfaces.go index a68e76ab..b24e9ac2 100644 --- a/internal/config/interfaces.go +++ b/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) diff --git a/internal/config/store_postgres.go b/internal/config/store_postgres.go index 6aa651cf..b4f018cd 100644 --- a/internal/config/store_postgres.go +++ b/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 } diff --git a/internal/config/store_postgres_complete_step_test.go b/internal/config/store_postgres_complete_step_test.go index ad87c318..d456ab59 100644 --- a/internal/config/store_postgres_complete_step_test.go +++ b/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)) diff --git a/internal/config/store_postgres_coverage_test.go b/internal/config/store_postgres_coverage_test.go index 95c9b494..b4a7a34e 100644 --- a/internal/config/store_postgres_coverage_test.go +++ b/internal/config/store_postgres_coverage_test.go @@ -284,8 +284,7 @@ func TestPostgresStore_UpdatePurchasePlan_NilDB(t *testing.T) { }) assert.True(t, panicked, "expected panic with nil db connection") - // UpdatedAt should be set before JSON marshaling and DB call - assert.False(t, plan.UpdatedAt.IsZero(), "UpdatedAt should have been set") + assert.True(t, plan.UpdatedAt.IsZero(), "failed writes must preserve the read version") } func TestPostgresStore_UpdatePurchasePlan_NilDB_WithServices(t *testing.T) { @@ -317,7 +316,7 @@ func TestPostgresStore_UpdatePurchasePlan_NilDB_WithServices(t *testing.T) { }) assert.True(t, panicked, "expected panic with nil db connection") - assert.False(t, plan.UpdatedAt.IsZero()) + assert.True(t, plan.UpdatedAt.IsZero()) } // ========================================== diff --git a/internal/config/store_postgres_pgxmock_test.go b/internal/config/store_postgres_pgxmock_test.go index 71fe5b48..f7f7fc2a 100644 --- a/internal/config/store_postgres_pgxmock_test.go +++ b/internal/config/store_postgres_pgxmock_test.go @@ -522,20 +522,16 @@ func TestPGXMock_UpdatePurchasePlan_NotFound(t *testing.T) { ctx := context.Background() plan := &PurchasePlan{ID: "missing-id", Services: map[string]ServiceConfig{}, RampSchedule: RampSchedule{}} - // UpdatePurchasePlan now wraps the UPDATE in a WithTx so it can be - // composed with other writes (e.g. createPlannedPurchases bundling - // per-row execution inserts with the plan's next_execution_date - // bump). The pgxmock script needs the matching Begin / Commit - // frame; the inner Exec returns 0 rows-affected, which the store - // surfaces as a "not found" error after the rollback. mock.ExpectBegin() - mock.ExpectExec("UPDATE").WithArgs(anyArgsCfg(11)...). - WillReturnResult(pgxmock.NewResult("UPDATE", 0)) + mock.ExpectQuery("UPDATE").WithArgs(anyArgsCfg(11)...). + WillReturnRows(pgxmock.NewRows([]string{"updated_at"})) + mock.ExpectQuery("SELECT EXISTS").WithArgs("missing-id"). + WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(false)) mock.ExpectRollback() err := store.UpdatePurchasePlan(ctx, plan) require.Error(t, err) - assert.Contains(t, err.Error(), "not found") + assert.ErrorIs(t, err, ErrNotFound) } // ─── DeletePurchasePlan ─────────────────────────────────────────────────────── diff --git a/internal/config/store_postgres_plan_update_integration_test.go b/internal/config/store_postgres_plan_update_integration_test.go new file mode 100644 index 00000000..d14a9512 --- /dev/null +++ b/internal/config/store_postgres_plan_update_integration_test.go @@ -0,0 +1,84 @@ +//go:build integration + +package config + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" +) + +func TestPurchasePlanUpdatesRejectStaleSnapshots(t *testing.T) { + ctx := context.Background() + store := setupRampStepStore(ctx, t) + plan := &PurchasePlan{Name: "original", Services: map[string]ServiceConfig{}} + require.NoError(t, store.CreatePurchasePlan(ctx, plan)) + createdVersion := plan.UpdatedAt + plan.Name = "first edit" + require.NoError(t, store.UpdatePurchasePlan(ctx, plan)) + require.True(t, plan.UpdatedAt.After(createdVersion)) + stale := *plan + plan.Name = "second edit" + require.NoError(t, store.UpdatePurchasePlan(ctx, plan)) + winner, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.Equal(t, plan.UpdatedAt, winner.UpdatedAt) + require.Equal(t, "second edit", winner.Name) + + staleVersion := stale.UpdatedAt + stale.Name = "stale edit" + require.ErrorIs(t, store.UpdatePurchasePlan(ctx, &stale), ErrPurchasePlanConflict) + require.Equal(t, staleVersion, stale.UpdatedAt) + after, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.Equal(t, winner, after) + stale.UpdatedAt = time.Time{} + require.ErrorIs(t, store.UpdatePurchasePlan(ctx, &stale), ErrPurchasePlanConflict) + stale.ID = uuid.NewString() + require.ErrorIs(t, store.UpdatePurchasePlan(ctx, &stale), ErrNotFound) + + stamp := time.Now().Truncate(time.Microsecond) + require.NoError(t, store.StampPlanNotificationSent(ctx, plan.ID, stamp)) + require.ErrorIs(t, store.UpdatePurchasePlan(ctx, plan), ErrPurchasePlanConflict) + after, err = store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.Equal(t, &stamp, after.LastNotificationSent) + require.Equal(t, winner.Name, after.Name) + require.True(t, after.UpdatedAt.After(winner.UpdatedAt)) + require.ErrorIs(t, store.StampPlanNotificationSent(ctx, uuid.NewString(), stamp), ErrNotFound) +} + +func TestPurchasePlanUpdateTxRollsBackWithExecutions(t *testing.T) { + ctx := context.Background() + store := setupRampStepStore(ctx, t) + plan := &PurchasePlan{Name: "original", Services: map[string]ServiceConfig{}} + require.NoError(t, store.CreatePurchasePlan(ctx, plan)) + before, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + executionID := uuid.NewString() + abort := errors.New("abort fixture transaction") + err = store.WithTx(ctx, func(tx pgx.Tx) error { + locked, lockErr := store.LockPurchasePlanTx(ctx, tx, plan.ID) + require.NoError(t, lockErr) + locked.Name = "rolled back" + require.NoError(t, store.UpdatePurchasePlanTx(ctx, tx, locked)) + require.True(t, locked.UpdatedAt.After(before.UpdatedAt)) + locked.Name = "second transaction-local write" + require.NoError(t, store.UpdatePurchasePlanTx(ctx, tx, locked)) + require.NoError(t, store.SavePurchaseExecutionTx(ctx, tx, &PurchaseExecution{ + ExecutionID: executionID, PlanID: plan.ID, Status: "pending", ScheduledDate: time.Now(), + })) + return abort + }) + require.ErrorIs(t, err, abort) + after, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.Equal(t, before, after) + _, err = store.GetExecutionByID(ctx, executionID) + require.ErrorIs(t, err, ErrNotFound) +} diff --git a/internal/config/store_postgres_plan_update_test.go b/internal/config/store_postgres_plan_update_test.go new file mode 100644 index 00000000..69d51d50 --- /dev/null +++ b/internal/config/store_postgres_plan_update_test.go @@ -0,0 +1,48 @@ +package config + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/pashagolub/pgxmock/v4" + "github.com/stretchr/testify/require" +) + +func TestUpdatePurchasePlanPreservesTokenOnCommitFailure(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + readVersion := time.Now().Truncate(time.Microsecond) + plan := &PurchasePlan{ID: "plan", UpdatedAt: readVersion} + commitErr := errors.New("commit failed") + mock.ExpectBegin() + mock.ExpectQuery(`UPDATE purchase_plans[\s\S]*WHERE id = \$1 AND updated_at = \$8[\s\S]*RETURNING updated_at`). + WithArgs(anyArgsCfg(11)...). + WillReturnRows(pgxmock.NewRows([]string{"updated_at"}).AddRow(readVersion.Add(time.Second))) + mock.ExpectCommit().WillReturnError(commitErr) + mock.ExpectRollback().WillReturnError(pgx.ErrTxClosed) + require.ErrorIs(t, store.UpdatePurchasePlan(context.Background(), plan), commitErr) + require.Equal(t, readVersion, plan.UpdatedAt) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestUpdatePurchasePlanPropagatesConflictProbeFailure(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + plan := &PurchasePlan{ID: "plan", UpdatedAt: time.Now()} + readVersion := plan.UpdatedAt + probeErr := errors.New("probe failed") + mock.ExpectBegin() + mock.ExpectQuery("UPDATE purchase_plans").WithArgs(anyArgsCfg(11)...). + WillReturnRows(pgxmock.NewRows([]string{"updated_at"})) + mock.ExpectQuery("SELECT EXISTS").WithArgs(plan.ID).WillReturnError(probeErr) + mock.ExpectRollback() + err := store.UpdatePurchasePlan(context.Background(), plan) + require.ErrorIs(t, err, probeErr) + require.NotErrorIs(t, err, ErrPurchasePlanConflict) + require.NotErrorIs(t, err, ErrNotFound) + require.Equal(t, readVersion, plan.UpdatedAt) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/internal/database/postgres/migrations/000102_purchase_plan_update_token.down.sql b/internal/database/postgres/migrations/000102_purchase_plan_update_token.down.sql new file mode 100644 index 00000000..6836fb3a --- /dev/null +++ b/internal/database/postgres/migrations/000102_purchase_plan_update_token.down.sql @@ -0,0 +1,5 @@ +CREATE OR REPLACE TRIGGER update_purchase_plans_updated_at + BEFORE UPDATE ON purchase_plans + FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +DROP FUNCTION update_purchase_plan_timestamp(); diff --git a/internal/database/postgres/migrations/000102_purchase_plan_update_token.up.sql b/internal/database/postgres/migrations/000102_purchase_plan_update_token.up.sql new file mode 100644 index 00000000..86264a78 --- /dev/null +++ b/internal/database/postgres/migrations/000102_purchase_plan_update_token.up.sql @@ -0,0 +1,11 @@ +CREATE OR REPLACE FUNCTION update_purchase_plan_timestamp() +RETURNS TRIGGER AS $$ +BEGIN + NEW.updated_at = GREATEST(clock_timestamp(), OLD.updated_at + INTERVAL '1 microsecond'); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE TRIGGER update_purchase_plans_updated_at + BEFORE UPDATE ON purchase_plans + FOR EACH ROW EXECUTE FUNCTION update_purchase_plan_timestamp(); diff --git a/internal/database/postgres/migrations/000102_purchase_plan_update_token_test.go b/internal/database/postgres/migrations/000102_purchase_plan_update_token_test.go new file mode 100644 index 00000000..772037da --- /dev/null +++ b/internal/database/postgres/migrations/000102_purchase_plan_update_token_test.go @@ -0,0 +1,52 @@ +//go:build integration + +package migrations_test + +import ( + "context" + "testing" + "time" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/migrations" + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/testhelpers" + "github.com/stretchr/testify/require" +) + +func TestMigration_PurchasePlanUpdateToken(t *testing.T) { + ctx := context.Background() + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, container.Cleanup(ctx)) }) + pool := container.DB.Pool() + path := getMigrationsPath() + require.NoError(t, migrations.RunMigrations(ctx, pool, path, "", "")) + + for _, future := range []bool{false, true} { + t.Run(map[bool]string{false: "current", true: "future"}[future], func(t *testing.T) { + tx, beginErr := pool.Begin(ctx) + require.NoError(t, beginErr) + defer func() { _ = tx.Rollback(ctx) }() + initial := time.Now().Truncate(time.Microsecond) + if future { + initial = initial.Add(time.Hour) + } + var id string + require.NoError(t, tx.QueryRow(ctx, `INSERT INTO purchase_plans (name, updated_at) VALUES ('token fixture', $1) RETURNING id`, initial).Scan(&id)) + var first, second time.Time + require.NoError(t, tx.QueryRow(ctx, `UPDATE purchase_plans SET name = 'first' WHERE id = $1 RETURNING updated_at`, id).Scan(&first)) + require.NoError(t, tx.QueryRow(ctx, `UPDATE purchase_plans SET name = 'second' WHERE id = $1 RETURNING updated_at`, id).Scan(&second)) + require.True(t, first.After(initial), "first token must advance, including beyond a future old token") + require.True(t, second.After(first), "two writes in the same transaction must have distinct tokens") + }) + } + + require.NoError(t, migrations.MigrateToVersion(ctx, pool, path, previousMigrationVersion(t, 102))) + var function string + require.NoError(t, pool.QueryRow(ctx, `SELECT p.proname FROM pg_trigger t JOIN pg_proc p ON p.oid = t.tgfoid WHERE t.tgname = 'update_purchase_plans_updated_at'`).Scan(&function)) + require.Equal(t, "update_updated_at_column", function) + require.NoError(t, migrations.MigrateToVersion(ctx, pool, path, 102)) + require.NoError(t, pool.QueryRow(ctx, `SELECT p.proname FROM pg_trigger t JOIN pg_proc p ON p.oid = t.tgfoid WHERE t.tgname = 'update_purchase_plans_updated_at'`).Scan(&function)) + require.Equal(t, "update_purchase_plan_timestamp", function) + require.NoError(t, pool.QueryRow(ctx, `SELECT p.proname FROM pg_trigger t JOIN pg_proc p ON p.oid = t.tgfoid WHERE t.tgname = 'update_purchase_executions_updated_at'`).Scan(&function)) + require.Equal(t, "update_updated_at_column", function, "other tables retain their existing trigger") +} diff --git a/internal/mocks/stores.go b/internal/mocks/stores.go index e7de8e0f..6c9e6d35 100644 --- a/internal/mocks/stores.go +++ b/internal/mocks/stores.go @@ -191,6 +191,11 @@ func (m *MockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return args.Error(0) } +func (m *MockConfigStore) StampPlanNotificationSent(ctx context.Context, planID string, at time.Time) error { + m.record("StampPlanNotificationSent", ctx, planID, at) + return m.Called(ctx, planID, at).Error(0) +} + // CompletePlanStep mocks the idempotent ramp-step advance operation. func (m *MockConfigStore) CompletePlanStep(ctx context.Context, planID string, stepNumber int) error { m.record("CompletePlanStep", ctx, planID, stepNumber) diff --git a/internal/purchase/notifications.go b/internal/purchase/notifications.go index 3905c885..811dcfab 100644 --- a/internal/purchase/notifications.go +++ b/internal/purchase/notifications.go @@ -115,8 +115,7 @@ func (m *Manager) sendPlanNotification(ctx context.Context, plan *config.Purchas // not persisted, so the next run would resend — return false to avoid // double-counting and let the caller decide whether to retry. now := time.Now() - plan.LastNotificationSent = &now - if err := m.config.UpdatePurchasePlan(ctx, plan); err != nil { + if err := m.config.StampPlanNotificationSent(ctx, plan.ID, now); err != nil { logging.Errorf("Failed to update plan notification timestamp: %v", err) return false } diff --git a/internal/purchase/notifications_ramp_race_integration_test.go b/internal/purchase/notifications_ramp_race_integration_test.go new file mode 100644 index 00000000..34da18ea --- /dev/null +++ b/internal/purchase/notifications_ramp_race_integration_test.go @@ -0,0 +1,74 @@ +//go:build integration + +package purchase + +import ( + "context" + "testing" + "time" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/LeanerCloud/cloud-commitments-platform/internal/email" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +func TestNotificationStampPreservesCompletedRampAndNextExecution(t *testing.T) { + ctx := context.Background() + store, _ := rampStepStore(ctx, t) + global, err := store.GetGlobalConfig(ctx) + require.NoError(t, err) + recipient := "fixture@example.invalid" + global.NotificationEmail = &recipient + require.NoError(t, store.SaveGlobalConfig(ctx, global)) + next := time.Now().Add(24 * time.Hour).Truncate(time.Microsecond) + plan := &config.PurchasePlan{ + Name: "notification race", Enabled: true, AutoPurchase: true, NotificationDaysBefore: 30, + Services: map[string]config.ServiceConfig{"ec2": {Provider: "aws", Service: "ec2"}}, + RampSchedule: config.RampSchedule{Type: "weekly", PercentPerStep: 25, StepIntervalDays: 7, + CurrentStep: 1, TotalSteps: 4, StartDate: next.AddDate(0, 0, -7)}, + NextExecutionDate: &next, + } + require.NoError(t, plan.Validate()) + require.NoError(t, store.CreatePurchasePlan(ctx, plan)) + sender := new(MockEmailSender) + manager := &Manager{config: store, email: sender} + var completed *config.PurchasePlan + sender.On("SendScheduledPurchaseNotification", mock.Anything, mock.Anything).Run(func(args mock.Arguments) { + data := args.Get(1).(email.NotificationData) + execution, getErr := store.GetExecutionByID(ctx, data.ExecutionID) + require.NoError(t, getErr) + require.Equal(t, 2, execution.StepNumber) + execution.Status = "completed" + require.NoError(t, store.SavePurchaseExecution(ctx, execution)) + require.NoError(t, store.CompletePlanStep(ctx, plan.ID, 2)) + completed, getErr = store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, getErr) + require.Equal(t, 2, completed.RampSchedule.CurrentStep) + require.NotNil(t, completed.LastExecutionDate) + require.True(t, completed.NextExecutionDate.After(next)) + }).Return(nil).Once() + result, err := manager.SendUpcomingPurchaseNotifications(ctx) + require.NoError(t, err) + require.Equal(t, 1, result.Notified) + sender.AssertExpectations(t) + after, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.NotNil(t, after.LastNotificationSent) + require.Equal(t, completed.RampSchedule.CurrentStep, after.RampSchedule.CurrentStep) + require.Equal(t, completed.NextExecutionDate, after.NextExecutionDate) + require.Equal(t, completed.LastExecutionDate, after.LastExecutionDate) + + execution, _, _, err := manager.getOrCreateExecution(ctx, after) + require.NoError(t, err) + require.Equal(t, 3, execution.StepNumber) + require.Equal(t, *completed.NextExecutionDate, execution.ScheduledDate) + execution.Status = "completed" + require.NoError(t, store.SavePurchaseExecution(ctx, execution)) + require.NoError(t, store.CompletePlanStep(ctx, plan.ID, 3)) + continued, err := store.GetPurchasePlan(ctx, plan.ID) + require.NoError(t, err) + require.Equal(t, 3, continued.RampSchedule.CurrentStep) + require.True(t, continued.NextExecutionDate.After(*after.NextExecutionDate)) + require.Equal(t, after.LastNotificationSent, continued.LastNotificationSent) +} diff --git a/internal/purchase/notifications_test.go b/internal/purchase/notifications_test.go index 2025e8f5..d591f291 100644 --- a/internal/purchase/notifications_test.go +++ b/internal/purchase/notifications_test.go @@ -442,7 +442,7 @@ func TestManager_SendUpcomingPurchaseNotifications_WithNotification(t *testing.T mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")).Return(nil) mockStore.On("GetGlobalConfig", ctx).Return(globalCfg, nil) mockEmail.On("SendScheduledPurchaseNotification", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("StampPlanNotificationSent", ctx, "plan-123", mock.AnythingOfType("time.Time")).Return(nil) manager := &Manager{ config: mockStore, @@ -656,6 +656,7 @@ func TestManager_SendPlanNotification_SendFailureLeavesLiveTokenAlone(t *testing manager := &Manager{config: mockStore, email: mockEmail, dashboardURL: "https://dashboard.example.com"} assert.False(t, manager.sendPlanNotification(ctx, plan)) mockStore.AssertNotCalled(t, "RotatePendingApprovalToken", mock.Anything, mock.Anything, mock.Anything, mock.Anything) + mockStore.AssertNotCalled(t, "StampPlanNotificationSent", mock.Anything, mock.Anything, mock.Anything) mockStore.AssertExpectations(t) mockEmail.AssertExpectations(t) } @@ -683,7 +684,7 @@ func TestManager_SendPlanNotification_PersistsEmailedTokenHashAfterSend(t *testi require.NotEmpty(t, emailedToken, "the token must be persisted only after the email went out") persistedHash = args.String(2) }).Return(true, nil) - mockStore.On("UpdatePurchasePlan", ctx, plan).Return(nil) + mockStore.On("StampPlanNotificationSent", ctx, plan.ID, mock.AnythingOfType("time.Time")).Return(nil) manager := &Manager{config: mockStore, email: mockEmail, dashboardURL: "https://dashboard.example.com"} assert.True(t, manager.sendPlanNotification(ctx, plan)) @@ -693,3 +694,20 @@ func TestManager_SendPlanNotification_PersistsEmailedTokenHashAfterSend(t *testi mockStore.AssertExpectations(t) mockEmail.AssertExpectations(t) } + +func TestManager_SendPlanNotification_StampFailure(t *testing.T) { + ctx := context.Background() + store := new(MockConfigStore) + sender := new(MockEmailSender) + plan, _ := pendingExecForNotification(ctx, store) + recipient := "fixture@example.invalid" + store.On("GetGlobalConfig", ctx).Return(&config.GlobalConfig{NotificationEmail: &recipient}, nil) + sender.On("SendScheduledPurchaseNotification", ctx, mock.Anything).Return(nil) + store.On("RotatePendingApprovalToken", ctx, "exec-rot", mock.Anything, mock.Anything).Return(true, nil) + store.On("StampPlanNotificationSent", ctx, plan.ID, mock.Anything).Return(errors.New("stamp failed")) + manager := &Manager{config: store, email: sender} + require.False(t, manager.sendPlanNotification(ctx, plan)) + require.Nil(t, plan.LastNotificationSent) + store.AssertExpectations(t) + sender.AssertExpectations(t) +} diff --git a/internal/server/test_helpers_test.go b/internal/server/test_helpers_test.go index 8fb94df9..0afc4e08 100644 --- a/internal/server/test_helpers_test.go +++ b/internal/server/test_helpers_test.go @@ -49,6 +49,10 @@ func (m *mockConfigStoreForHealth) UpdatePurchasePlan(ctx context.Context, plan return nil } +func (m *mockConfigStoreForHealth) StampPlanNotificationSent(context.Context, string, time.Time) error { + return nil +} + func (m *mockConfigStoreForHealth) CompletePlanStep(_ context.Context, _ string, _ int) error { return nil }