diff --git a/internal/config/store_postgres.go b/internal/config/store_postgres.go index 5dc57909..e19fdacd 100644 --- a/internal/config/store_postgres.go +++ b/internal/config/store_postgres.go @@ -728,6 +728,9 @@ func advanceRampStep(ctx context.Context, tx pgx.Tx, plan *PurchasePlan, stepNum // next_execution_date at the step after it (or clearing it on a finished ramp) // and stamping the execution timestamps. func (s *PostgresStore) persistRampPosition(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) error { + if plan.RampSchedule.DatesManaged { + return s.persistManagedRampPosition(ctx, tx, plan) + } if plan.RampSchedule.IsComplete() { plan.NextExecutionDate = nil } else { diff --git a/internal/config/store_postgres_plan_schedule_clock.go b/internal/config/store_postgres_plan_schedule_clock.go new file mode 100644 index 00000000..820ee1d6 --- /dev/null +++ b/internal/config/store_postgres_plan_schedule_clock.go @@ -0,0 +1,127 @@ +package config + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "github.com/jackc/pgx/v5" +) + +func (s *PostgresStore) persistManagedRampPosition(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) error { + next, err := s.planScheduleNextDateTx(ctx, tx, plan) + if err != nil { + return err + } + now := time.Now() + var updated time.Time + err = tx.QueryRow(ctx, `UPDATE purchase_plans SET + ramp_schedule = jsonb_set(ramp_schedule, '{current_step}', to_jsonb($2::int)), + next_execution_date = $3, last_execution_date = $4 + WHERE id = $1 AND updated_at = $5 RETURNING updated_at`, + plan.ID, plan.RampSchedule.CurrentStep, next, now, plan.UpdatedAt).Scan(&updated) + if errors.Is(err, pgx.ErrNoRows) { + return ErrPurchasePlanConflict + } + if err != nil { + return fmt.Errorf("persist managed ramp position: %w", err) + } + plan.NextExecutionDate, plan.LastExecutionDate, plan.UpdatedAt = next, &now, updated + return nil +} + +func (s *PostgresStore) planScheduleNextDateTx(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) (*time.Time, error) { + if plan.RampSchedule.IsComplete() { + return nil, nil + } + if !plan.RampSchedule.DatesManaged { + next := plan.RampSchedule.GetNextPurchaseDate() + return &next, nil + } + root, err := s.managedPlanNextRootTx(ctx, tx, plan) + if err != nil || root == nil { + return nil, err + } + return &root.ScheduledDate, nil +} + +func (s *PostgresStore) managedPlanNextRootTx(ctx context.Context, tx pgx.Tx, plan *PurchasePlan) (*PurchaseExecution, error) { + rows, err := s.ListExecutionsForPlanStepTx(ctx, tx, plan.ID, plan.RampSchedule.CurrentStep+1) + if err != nil { + return nil, err + } + root, err := managedPlanNextRoot(plan, rows) + if err != nil || root == nil { + return nil, err + } + var incoming bool + if err := tx.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM purchase_executions WHERE retry_execution_id = $1)`, root.ExecutionID).Scan(&incoming); err != nil { + return nil, fmt.Errorf("check managed root retry references: %w", err) + } + if incoming { + return nil, nil + } + return root, nil +} + +func managedPlanNextRoot(plan *PurchasePlan, rows []PurchaseExecution) (*PurchaseExecution, error) { + var root *PurchaseExecution + blocked := false + for i := range rows { + row := &rows[i] + if row.StepNumber != plan.RampSchedule.CurrentStep+1 { + continue + } + if row.Status == StatusCanceled || row.Status == LegacyStatusCanceled { + blocked = blocked || managedRootPriorAction(row) + continue + } + if !managedRootIdentity(row) { + blocked = true + continue + } + if root != nil { + return nil, fmt.Errorf("%w: duplicate managed roots for step %d", ErrPurchasePlanConflict, row.StepNumber) + } + root = row + } + if blocked || !managedRootClockAvailable(root) { + return nil, nil + } + return root, nil +} + +func managedRootIdentity(row *PurchaseExecution) bool { + lineage := row.IdempotencyKey + if lineage == "" { + lineage = row.ExecutionID + } + return row.CloudAccountID == nil && row.RetryAttemptN == 0 && row.RetryExecutionID == nil && !strings.Contains(lineage, ":") +} + +func managedRootPriorAction(row *PurchaseExecution) bool { + if row.ExecutedAt != nil || row.CompletedAt != nil || row.PreApprovalSkipReason != nil && *row.PreApprovalSkipReason != "" { + return true + } + for i := range row.Recommendations { + if row.Recommendations[i].Purchased { + return true + } + } + return false +} + +func managedRootClockAvailable(row *PurchaseExecution) bool { + if row == nil || row.ScheduledDate.IsZero() || managedRootPriorAction(row) { + return false + } + switch row.Status { + case "pending", "notified": + return row.ScheduledExecutionAt == nil + case "scheduled": + return row.ScheduledExecutionAt != nil && !row.ScheduledExecutionAt.IsZero() + } + return false +} diff --git a/internal/config/store_postgres_plan_schedule_clock_integration_test.go b/internal/config/store_postgres_plan_schedule_clock_integration_test.go new file mode 100644 index 00000000..28e5c557 --- /dev/null +++ b/internal/config/store_postgres_plan_schedule_clock_integration_test.go @@ -0,0 +1,188 @@ +//go:build integration + +package config + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" +) + +func managedClockFixture(ctx context.Context, t *testing.T, current int) *rampFixture { + t.Helper() + f := newRampFixture(ctx, t, 0, current, 3) + _, err := f.store.db.Exec(ctx, `UPDATE purchase_plans SET + ramp_schedule = ramp_schedule || '{"dates_managed":true,"future_metadata":{"nested":9007199254740993}}', + services = '{"future_service":{"enabled":false,"unknown_filter":["retained"]}}', + ramp_baseline = '{"frozen":{"count":97}}' WHERE id = $1`, f.planID) + require.NoError(t, err) + return f +} + +func nextClockRoot(ctx context.Context, t *testing.T, f *rampFixture, status string) *PurchaseExecution { + t.Helper() + root := f.row(ctx, t, 3, nil, status) + date := time.Now().UTC().Truncate(24*time.Hour).AddDate(0, 0, 47) + _, err := f.store.db.Exec(ctx, `UPDATE purchase_executions SET scheduled_date = $2 WHERE execution_id = $1`, root.ExecutionID, date) + require.NoError(t, err) + root.ScheduledDate = date + return root +} + +func TestManagedCompletionKeepsNextRootAndRawPlan(t *testing.T) { + for _, targetMet := range []bool{false, true} { + t.Run(map[bool]string{false: "bought", true: "target met"}[targetMet], func(t *testing.T) { + ctx := context.Background() + f := managedClockFixture(ctx, t, 1) + root := nextClockRoot(ctx, t, f, "pending") + heldRec := baselineRec(uuid.NewString(), "m5.large", 17) + f.rowWithRecs(ctx, t, 1, nil, StatusCompleted, false, heldRec) + var before []byte + heldBefore := f.committed(ctx, t) + require.Equal(t, 17, heldBefore[baselineKey(t, heldRec)].Count) + require.NoError(t, f.store.db.QueryRow(ctx, `SELECT jsonb_build_object('services',services,'baseline',ramp_baseline, + 'metadata',ramp_schedule->'future_metadata') FROM purchase_plans WHERE id=$1`, f.planID).Scan(&before)) + if targetMet { + current := f.row(ctx, t, 2, nil, "pending") + require.NoError(t, f.store.WithTx(ctx, func(tx pgx.Tx) error { + return f.store.CompleteTargetMetStepTx(ctx, tx, f.planID, 2, current.ExecutionID, "target already held") + })) + } else { + f.row(ctx, t, 2, nil, StatusCompleted) + require.NoError(t, f.store.CompletePlanStep(ctx, f.planID, 2)) + } + plan, err := f.store.GetPurchasePlan(ctx, f.planID) + require.NoError(t, err) + require.Equal(t, 2, plan.RampSchedule.CurrentStep) + require.NotNil(t, plan.NextExecutionDate) + require.True(t, root.ScheduledDate.Equal(*plan.NextExecutionDate), "completion must retain the exact next managed root date") + var after []byte + require.NoError(t, f.store.db.QueryRow(ctx, `SELECT jsonb_build_object('services',services,'baseline',ramp_baseline, + 'metadata',ramp_schedule->'future_metadata') FROM purchase_plans WHERE id=$1`, f.planID).Scan(&after)) + require.Equal(t, before, after) + require.Equal(t, heldBefore, f.committed(ctx, t)) + var ramp map[string]json.RawMessage + require.NoError(t, f.store.db.QueryRow(ctx, `SELECT ramp_schedule FROM purchase_plans WHERE id=$1`, f.planID).Scan(&after)) + require.NoError(t, json.Unmarshal(after, &ramp)) + require.JSONEq(t, "true", string(ramp["dates_managed"])) + }) + } +} + +func TestManagedCompletionStopsAtBlockedNextRoot(t *testing.T) { + cases := []struct { + name, status, sql string + extra string + conflict, healthy bool + }{ + {name: "missing"}, + {name: "canceled", status: StatusCanceled}, + {name: "failed", status: StatusFailed}, + {name: "paused", status: "paused"}, + {name: "legacy uncertain", status: "approved"}, + {name: "duplicate", status: "pending", extra: "duplicate", conflict: true}, + {name: "zero date", status: "pending", sql: "scheduled_date='0001-01-01T00:00:00Z'"}, + {name: "scheduled without delayed fire", status: "scheduled"}, + {name: "pending delayed fire", status: "pending", sql: "scheduled_execution_at=NOW()"}, + {name: "executed", status: "pending", sql: "executed_at=NOW()"}, + {name: "completed provenance", status: "pending", sql: "completed_at=NOW()"}, + {name: "skip provenance", status: "pending", sql: "pre_approval_skip_reason='unknown'"}, + {name: "purchased provenance", status: "pending", sql: `recommendations='[{"purchased":true}]'`}, + {name: "canceled purchased predecessor", status: "pending", extra: "canceled purchased"}, + {name: "canceled executed predecessor", status: "pending", extra: "canceled executed"}, + {name: "live child", status: "pending", extra: "child"}, + {name: "incoming other step", status: "pending", extra: "incoming step"}, + {name: "incoming other plan", status: "pending", extra: "incoming plan"}, + {name: "canceled unacted replacement", status: "pending", extra: "canceled", healthy: true}, + {name: "scheduled trustworthy", status: "scheduled", sql: "scheduled_execution_at=NOW()+INTERVAL '48 days'", healthy: true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + f := managedClockFixture(ctx, t, 1) + f.row(ctx, t, 2, nil, StatusCompleted) + var root *PurchaseExecution + if tc.status != "" { + root = nextClockRoot(ctx, t, f, tc.status) + if tc.sql != "" { + _, err := f.store.db.Exec(ctx, "UPDATE purchase_executions SET "+tc.sql+" WHERE execution_id=$1", root.ExecutionID) + require.NoError(t, err) + } + } + managedClockSibling(ctx, t, f, root, tc.extra) + err := f.store.CompletePlanStep(ctx, f.planID, 2) + if tc.conflict { + require.ErrorIs(t, err, ErrPurchasePlanConflict) + require.Equal(t, 1, f.currentStep(ctx, t)) + return + } + require.NoError(t, err) + plan, err := f.store.GetPurchasePlan(ctx, f.planID) + require.NoError(t, err) + require.Equal(t, 2, plan.RampSchedule.CurrentStep) + if tc.healthy { + require.NotNil(t, plan.NextExecutionDate) + require.True(t, root.ScheduledDate.Equal(*plan.NextExecutionDate)) + } else { + require.Nil(t, plan.NextExecutionDate, "blocked next root must stop the managed clock") + } + }) + } +} + +func managedClockSibling(ctx context.Context, t *testing.T, f *rampFixture, root *PurchaseExecution, extra string) { + t.Helper() + if extra == "" { + return + } + status := StatusCanceled + if extra == "duplicate" || extra == "child" || extra == "incoming step" || extra == "incoming plan" { + status = "pending" + } + owner, step := f, 3 + if extra == "incoming step" { + step = 1 + } + if extra == "incoming plan" { + other, err := f.store.GetPurchasePlan(ctx, f.planID) + require.NoError(t, err) + other.ID, other.Name = "", "Other "+uuid.NewString()[:8] + require.NoError(t, f.store.CreatePurchasePlan(ctx, other)) + owner = &rampFixture{store: f.store, planID: other.ID} + } + sibling := owner.row(ctx, t, step, nil, status) + var err error + switch extra { + case "child": + _, err = f.store.db.Exec(ctx, `UPDATE purchase_executions SET retry_attempt_n=1 WHERE execution_id=$1`, sibling.ExecutionID) + case "canceled purchased": + _, err = f.store.db.Exec(ctx, `UPDATE purchase_executions SET recommendations='[{"purchased":true}]' WHERE execution_id=$1`, sibling.ExecutionID) + case "canceled executed": + _, err = f.store.db.Exec(ctx, `UPDATE purchase_executions SET executed_at=NOW() WHERE execution_id=$1`, sibling.ExecutionID) + case "incoming step", "incoming plan": + _, err = f.store.db.Exec(ctx, `UPDATE purchase_executions SET retry_execution_id=$2 WHERE execution_id=$1`, sibling.ExecutionID, root.ExecutionID) + } + require.NoError(t, err) +} + +func TestManagedCompletionFinishedAndLegacyClock(t *testing.T) { + ctx := context.Background() + f := managedClockFixture(ctx, t, 2) + f.row(ctx, t, 3, nil, StatusCompleted) + require.NoError(t, f.store.CompletePlanStep(ctx, f.planID, 3)) + plan, err := f.store.GetPurchasePlan(ctx, f.planID) + require.NoError(t, err) + require.Nil(t, plan.NextExecutionDate) + legacy := newRampFixture(ctx, t, 0, 1, 3) + legacy.row(ctx, t, 2, nil, StatusCompleted) + require.NoError(t, legacy.store.CompletePlanStep(ctx, legacy.planID, 2)) + plan, err = legacy.store.GetPurchasePlan(ctx, legacy.planID) + require.NoError(t, err) + require.NotNil(t, plan.NextExecutionDate) + require.WithinDuration(t, plan.RampSchedule.GetNextPurchaseDate(), *plan.NextExecutionDate, time.Millisecond) // Postgres stores microseconds +} diff --git a/internal/config/types.go b/internal/config/types.go index 31e0cd5c..80a650e8 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -254,6 +254,7 @@ type RampSchedule struct { CurrentStep int `json:"current_step" dynamodbav:"current_step"` TotalSteps int `json:"total_steps" dynamodbav:"total_steps"` StartDate time.Time `json:"start_date" dynamodbav:"start_date"` + DatesManaged bool `json:"dates_managed,omitempty" dynamodbav:"dates_managed,omitempty"` } // PresetRampSchedules provides common ramp-up configurations.