diff --git a/cmd/server/main.go b/cmd/server/main.go index 5cfd3ae74..266abd46f 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -58,7 +58,7 @@ func main() { log.Printf("Running scheduled task: %s (timeout: %v)", *task, timeout) taskType := server.ScheduledTaskType(*task) - result, err := app.HandleScheduledTask(taskCtx, taskType) + result, err := app.HandleScheduledTask(taskCtx, taskType, server.ScheduledTaskParams{}) cancel() if err != nil { log.Fatalf("Scheduled task %q failed: %v", *task, err) //nolint:gocritic // exitAfterDefer: intentional fatal; app.Close() not needed on task failure diff --git a/internal/analytics/collector_test.go b/internal/analytics/collector_test.go index d65dabb3c..a07bb0607 100644 --- a/internal/analytics/collector_test.go +++ b/internal/analytics/collector_test.go @@ -424,10 +424,10 @@ func (m *mockConfigStore) GetRecommendationsFreshness(_ context.Context) (*confi func (m *mockConfigStore) SetRecommendationsCollectionError(_ context.Context, _ string) error { return nil } -func (m *mockConfigStore) MarkCollectionStarted(_ context.Context) (bool, error) { - return true, nil +func (m *mockConfigStore) MarkCollectionStarted(_ context.Context) (string, bool, error) { + return "mock-owner-token", true, nil } -func (m *mockConfigStore) ClearCollectionStarted(_ context.Context) error { +func (m *mockConfigStore) ClearCollectionStarted(_ context.Context, _ string) error { return nil } func (m *mockConfigStore) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { diff --git a/internal/api/handler_recommendations_refresh.go b/internal/api/handler_recommendations_refresh.go index 20e8cce37..45a08d30f 100644 --- a/internal/api/handler_recommendations_refresh.go +++ b/internal/api/handler_recommendations_refresh.go @@ -16,6 +16,11 @@ import ( "github.com/aws/aws-sdk-go-v2/service/lambda/types" ) +// markedCollectionRollbackTimeout bounds the detached best-effort clear issued +// when the async self-invoke fails. Matches the scheduler's deferred clear and +// Application.releaseSkippedCollectionMarker budgets. +const markedCollectionRollbackTimeout = 5 * time.Second + // LambdaInvokerInterface is the narrow subset of lambda.Client used by the // async refresh handler. Extracted so tests can inject a stub without // standing up a real Lambda client. @@ -66,7 +71,10 @@ func (h *Handler) postRefreshRecommendations(ctx context.Context, req *events.La // Atomically mark collection as started. Returns false (409) if another // collection is already in flight (started_at set within the last 5 minutes). - ok, err := h.config.MarkCollectionStarted(ctx) + // The returned token identifies this caller as the marker's owner; it is + // threaded through the async invoke (or the sync collect call) so only + // this run can later clear the marker (issue #261 compare-and-clear guard). + token, ok, err := h.config.MarkCollectionStarted(ctx) if err != nil { return nil, fmt.Errorf("failed to mark collection started: %w", err) } @@ -74,7 +82,7 @@ func (h *Handler) postRefreshRecommendations(ctx context.Context, req *events.La return nil, NewClientError(409, "recommendation collection already in progress; try again in a few minutes") } - freshness, err := h.runMarkedCollection(ctx) + freshness, err := h.runMarkedCollection(ctx, token) if err != nil { return nil, err } @@ -103,13 +111,16 @@ func (h *Handler) postRefreshRecommendations(ctx context.Context, req *events.La // // Extracted from postRefreshRecommendations to keep that function under the // project's cyclomatic-complexity gate after the post-async re-read was added. -func (h *Handler) runMarkedCollection(ctx context.Context) (*config.RecommendationsFreshness, error) { +// +// token is the owner token MarkCollectionStarted returned to the caller; it +// is passed to the async invoke payload (so the scheduler can thread it +// back into ClearCollectionStarted) and to the rollback clear on this +// handler's own failure path. +func (h *Handler) runMarkedCollection(ctx context.Context, token string) (*config.RecommendationsFreshness, error) { schedulerARN := os.Getenv("SCHEDULER_LAMBDA_ARN") if schedulerARN != "" { - if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN); invokeErr != nil { - if clearErr := h.config.ClearCollectionStarted(ctx); clearErr != nil { - logging.Warnf("runMarkedCollection: failed to clear collection started marker: %v", clearErr) - } + if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN, token); invokeErr != nil { + h.rollbackMarkedCollection(token) return nil, fmt.Errorf("failed to trigger async collection: %w", invokeErr) } freshness, err := h.config.GetRecommendationsFreshness(ctx) @@ -119,8 +130,9 @@ func (h *Handler) runMarkedCollection(ctx context.Context) (*config.Recommendati return freshness, nil } // Non-Lambda (HTTP) mode: collect synchronously. ClearCollectionStarted - // is called by the scheduler's deferred clearCollectionStartedBestEffort. - if _, collectErr := h.scheduler.CollectRecommendations(ctx); collectErr != nil { + // is called by the scheduler's deferred clearCollectionStartedBestEffort, + // passed this same token so the clear is scoped to this run. + if _, collectErr := h.scheduler.CollectRecommendations(ctx, token); collectErr != nil { return nil, fmt.Errorf("collection failed: %w", collectErr) } freshness, err := h.config.GetRecommendationsFreshness(ctx) @@ -130,12 +142,37 @@ func (h *Handler) runMarkedCollection(ctx context.Context) (*config.Recommendati return freshness, nil } +// rollbackMarkedCollection releases this caller's own in-flight marker after +// the async self-invoke failed and no collection will ever run to release it. +// +// Deliberately detached from the request context. The dominant cause of an +// asyncInvokeSelf failure is the request context itself expiring or being +// canceled (API Gateway deadline, client disconnect, a slow SDK credential +// refresh), so reusing that context here would fail the rollback in exactly +// the case that produced it. The marker would then sit stranded for the full +// 5-minute auto-recovery window in MarkCollectionStarted, rejecting every +// refresh the user attempts with 409 while no collection is running. Matches +// the detached-clear pattern in the scheduler's deferred clear and in +// Application.releaseSkippedCollectionMarker. +// +// Scoped by token, so it can only ever release the marker this caller owns. +// Best effort: a failure only defers cleanup to the 5-minute window, and the +// caller already returns the underlying invoke error, so it is logged rather +// than masking that error. +func (h *Handler) rollbackMarkedCollection(token string) { + clearCtx, cancel := context.WithTimeout(context.Background(), markedCollectionRollbackTimeout) + defer cancel() + if err := h.config.ClearCollectionStarted(clearCtx, token); err != nil { + logging.Warnf("rollbackMarkedCollection: failed to clear collection started marker: %v", err) + } +} + // asyncInvokeSelf fires an InvocationType=Event invoke of the given Lambda // function ARN with the EventBridge-style payload that handleLambdaScheduledEvent // recognizes as a "collect recommendations" job. The call returns immediately; // the Lambda runtime delivers the event to the next available container // (which may be this same container's next invocation). -func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN string) error { +func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN, ownerToken string) error { invoker, err := h.getLambdaInvoker(ctx) if err != nil { return fmt.Errorf("failed to build Lambda client: %w", err) @@ -159,8 +196,9 @@ func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN string) error // Action != "") classifies this consistently with EventBridge cron // deliveries that already exercise this code path. payload, marshalErr := json.Marshal(map[string]string{ - "source": "aws.events", - "action": "collect_recommendations", + "source": "aws.events", + "action": "collect_recommendations", + "owner_token": ownerToken, }) if marshalErr != nil { return fmt.Errorf("asyncInvokeSelf: failed to marshal payload: %w", marshalErr) diff --git a/internal/api/handler_test.go b/internal/api/handler_test.go index 20c47345f..52ef69675 100644 --- a/internal/api/handler_test.go +++ b/internal/api/handler_test.go @@ -11,6 +11,7 @@ import ( "github.com/LeanerCloud/CUDly/internal/credentials" "github.com/LeanerCloud/CUDly/internal/scheduler" "github.com/aws/aws-lambda-go/events" + "github.com/aws/aws-sdk-go-v2/service/lambda" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" @@ -509,10 +510,10 @@ func TestHandler_HandleRequest_RefreshRecommendations(t *testing.T) { mockAuth.grantAdmin() mockAuth.On("ValidateCSRFToken", ctx, mock.Anything, mock.Anything).Return(nil) - mockScheduler.On("CollectRecommendations", mock.Anything).Return(&scheduler.CollectResult{Recommendations: 0, TotalSavings: 0}, nil) + mockScheduler.On("CollectRecommendations", mock.Anything, mock.Anything).Return(&scheduler.CollectResult{Recommendations: 0, TotalSavings: 0}, nil) mockStore.On("GetRecommendationsFreshness", mock.Anything). Return(&config.RecommendationsFreshness{}, nil) - mockStore.On("MarkCollectionStarted", mock.Anything).Return(true, nil) + mockStore.On("MarkCollectionStarted", mock.Anything).Return("tok-xyz", true, nil) handler := &Handler{config: mockStore, scheduler: mockScheduler, auth: mockAuth, apiKey: "test-key"} @@ -535,7 +536,118 @@ func TestHandler_HandleRequest_RefreshRecommendations(t *testing.T) { require.NoError(t, err) // Sync-fallback path returns the response body fully populated, so 200. assert.Equal(t, 200, resp.StatusCode) - mockScheduler.AssertCalled(t, "CollectRecommendations", mock.Anything) + mockScheduler.AssertCalled(t, "CollectRecommendations", mock.Anything, "tok-xyz") +} + +// stubLambdaInvoker is a minimal LambdaInvokerInterface for exercising the +// async-invoke path in runMarkedCollection without a real Lambda client. +type stubLambdaInvoker struct { + invokeFn func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error) +} + +func (s *stubLambdaInvoker) Invoke(ctx context.Context, params *lambda.InvokeInput, _ ...func(*lambda.Options)) (*lambda.InvokeOutput, error) { + return s.invokeFn(ctx, params) +} + +// TestRunMarkedCollection_AsyncInvokeFailure_RollsBackOwnerToken pins issue +// #261: when the async self-invoke fails, the rollback clear must use the +// same token MarkCollectionStarted returned to this caller (not an +// unconditional clear that could wipe an unrelated in-flight run). +func TestRunMarkedCollection_AsyncInvokeFailure_RollsBackOwnerToken(t *testing.T) { + ctx := context.Background() + t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly") + + mockStore := new(MockConfigStore) + t.Cleanup(func() { mockStore.AssertExpectations(t) }) + mockStore.On("ClearCollectionStarted", mock.Anything, "tok-xyz").Return(nil) + + invoker := &stubLambdaInvoker{ + invokeFn: func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error) { + return nil, assert.AnError + }, + } + + handler := &Handler{config: mockStore, lambdaInvoker: invoker} + + _, err := handler.runMarkedCollection(ctx, "tok-xyz") + require.Error(t, err) + mockStore.AssertCalled(t, "ClearCollectionStarted", mock.Anything, "tok-xyz") +} + +// TestRunMarkedCollection_AsyncInvokeFailure_RollbackSurvivesCanceledCtx pins +// the rollback clear to a DETACHED context. The dominant cause of an +// asyncInvokeSelf failure is the request context itself expiring or being +// canceled, so a rollback reusing that context would fail in exactly the case +// that produced it, stranding this caller's marker for the full 5-minute +// auto-recovery window and 409-ing every refresh in between. +// +// The context matcher lives on the expectation itself rather than in a +// trailing AssertCalled: testify's AssertCalled compares arguments by value +// and does not honor MatchedBy, whereas an unmatched expectation makes the +// mock fail the call outright. Pre-fix the clear still happened, just with +// the already-canceled request context, so asserting only that Clear was +// called would pass with the bug present. +func TestRunMarkedCollection_AsyncInvokeFailure_RollbackSurvivesCanceledCtx(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly") + + mockStore := new(MockConfigStore) + t.Cleanup(func() { mockStore.AssertExpectations(t) }) + liveDeadlineCtx := mock.MatchedBy(func(clearCtx context.Context) bool { + _, hasDeadline := clearCtx.Deadline() + return hasDeadline && clearCtx.Err() == nil + }) + mockStore.On("ClearCollectionStarted", liveDeadlineCtx, "tok-xyz").Return(nil) + + invoker := &stubLambdaInvoker{ + invokeFn: func(invokeCtx context.Context, _ *lambda.InvokeInput) (*lambda.InvokeOutput, error) { + // Model the real failure mode: the request context dies mid-invoke + // and the SDK surfaces that as the invoke error. + cancel() + return nil, invokeCtx.Err() + }, + } + + handler := &Handler{config: mockStore, lambdaInvoker: invoker} + + _, err := handler.runMarkedCollection(ctx, "tok-xyz") + require.Error(t, err) +} + +// TestRunMarkedCollection_AsyncInvokeSuccess_PayloadCarriesOwnerToken pins +// issue #261: the async-invoke payload must carry this run's owner token so +// the scheduler receiving the event can scope its own clear to this run; +// on the happy path this handler must NOT call Clear itself (the +// scheduler's own deferred clear owns that once collection finishes). +func TestRunMarkedCollection_AsyncInvokeSuccess_PayloadCarriesOwnerToken(t *testing.T) { + ctx := context.Background() + t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly") + + mockStore := new(MockConfigStore) + t.Cleanup(func() { mockStore.AssertExpectations(t) }) + mockStore.On("GetRecommendationsFreshness", mock.Anything). + Return(&config.RecommendationsFreshness{}, nil) + // Registered (as Maybe) so an unwanted clear is RECORDED: MockConfigStore + // short-circuits methods with no registered expectation before reaching + // mock.Called, so without this the AssertNotCalled below would pass even + // if the happy path did clear the marker. + mockStore.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe() + + var capturedPayload map[string]string + invoker := &stubLambdaInvoker{ + invokeFn: func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error) { + require.NoError(t, json.Unmarshal(params.Payload, &capturedPayload)) + return &lambda.InvokeOutput{}, nil + }, + } + + handler := &Handler{config: mockStore, lambdaInvoker: invoker} + + _, err := handler.runMarkedCollection(ctx, "tok-xyz") + require.NoError(t, err) + + assert.Equal(t, "tok-xyz", capturedPayload["owner_token"]) + mockStore.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything) } func TestHandler_HandleRequest_ListPlans(t *testing.T) { diff --git a/internal/api/mocks_test.go b/internal/api/mocks_test.go index 4bfc68568..5d41e12b9 100644 --- a/internal/api/mocks_test.go +++ b/internal/api/mocks_test.go @@ -67,8 +67,8 @@ type MockScheduler struct { mock.Mock } -func (m *MockScheduler) CollectRecommendations(ctx context.Context) (*scheduler.CollectResult, error) { - args := m.Called(ctx) +func (m *MockScheduler) CollectRecommendations(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { + args := m.Called(ctx, ownerToken) if args.Get(0) == nil { return nil, args.Error(1) } diff --git a/internal/api/types.go b/internal/api/types.go index bbc3ea92d..293d21fda 100644 --- a/internal/api/types.go +++ b/internal/api/types.go @@ -142,7 +142,7 @@ type PurchaseManagerInterface interface { // SchedulerInterface defines scheduler methods used by handler. type SchedulerInterface interface { - CollectRecommendations(ctx context.Context) (*scheduler.CollectResult, error) + CollectRecommendations(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) ListRecommendations(ctx context.Context, filter config.RecommendationFilter) ([]config.RecommendationRecord, error) // GetRecommendationByID fetches a single rec by its application-level id, // bypassing account-override filtering so deep-linked URLs to override- diff --git a/internal/config/interfaces.go b/internal/config/interfaces.go index f78bd8796..3edbe9e79 100644 --- a/internal/config/interfaces.go +++ b/internal/config/interfaces.go @@ -330,14 +330,21 @@ type StoreInterface interface { SetRecommendationsCollectionError(ctx context.Context, errMsg string) error // MarkCollectionStarted atomically sets last_collection_started_at = now // only when no in-flight collection is running (last_collection_started_at IS NULL - // OR older than 5 minutes). Returns true when this caller won the race and - // should proceed with the async invoke; false when another collection is - // already in flight and the caller should return 409. - MarkCollectionStarted(ctx context.Context) (bool, error) - // ClearCollectionStarted clears last_collection_started_at. Called by the - // scheduler at the end of every CollectRecommendations run, whether it - // succeeded or failed, so the UI knows the collection has finished. - ClearCollectionStarted(ctx context.Context) error + // OR older than 5 minutes), and stamps a freshly generated owner token + // alongside it. Returns the token and true when this caller won the race + // and should proceed with the async invoke, passing the token through so + // only this caller can later clear the marker; ("", false, nil) when + // another collection is already in flight and the caller should return + // 409. + MarkCollectionStarted(ctx context.Context) (token string, ok bool, err error) + // ClearCollectionStarted clears last_collection_started_at, but only if + // last_collection_owner_id still matches token: a compare-and-clear + // guard (issue #261) so a caller that never won MarkCollectionStarted + // (cron, cold-start) cannot wipe another caller's in-flight marker. A + // mismatched token is a documented silent no-op (the marker belongs to + // someone else); an empty token is a boundary error, since only a + // caller that actually owns a marker should ever call Clear. + ClearCollectionStarted(ctx context.Context, token string) error // RI utilization cache. Postgres-backed TTL cache for Cost Explorer // GetReservationUtilization; shared across Lambda containers so diff --git a/internal/config/store_postgres_recommendations.go b/internal/config/store_postgres_recommendations.go index 97ab3c60f..9a8335e9f 100644 --- a/internal/config/store_postgres_recommendations.go +++ b/internal/config/store_postgres_recommendations.go @@ -8,6 +8,7 @@ import ( "strings" "time" + "github.com/LeanerCloud/CUDly/pkg/logging" "github.com/google/uuid" "github.com/jackc/pgx/v5" ) @@ -427,16 +428,24 @@ func (s *PostgresStore) GetRecommendationsFreshness(ctx context.Context) (*Recom } // SetRecommendationsCollectionError records the most recent collection's -// error message without touching last_collected_at. Also clears -// last_collection_started_at so the frontend knows the collection has -// finished (with an error). Used by the scheduler when a collect fails +// error message without touching last_collected_at or +// last_collection_started_at. Used by the scheduler when a collect fails // partially or fully so the frontend banner surfaces the issue while // existing cached rows stay visible. +// +// This method must NOT clear last_collection_started_at (issue #261): it is +// called mid-run from persistCollection on every CollectRecommendations +// invocation that hits a provider error, including tokenless cron/cold-start/ +// background runs. Clearing the marker here, unconditionally and with no +// owner check, previously let a tokenless run's routine provider error wipe +// a concurrent owner run's in-flight marker, reopening the exact race the +// compare-and-clear guard exists to close. Only the deferred, token-guarded +// clearCollectionStartedBestEffort (which runs on both the success and +// failure exit paths of CollectRecommendations) may clear started_at. func (s *PostgresStore) SetRecommendationsCollectionError(ctx context.Context, errMsg string) error { if _, err := s.db.Exec(ctx, ` UPDATE recommendations_state - SET last_collection_error = $1, - last_collection_started_at = NULL + SET last_collection_error = $1 WHERE id = 1 `, errMsg); err != nil { return fmt.Errorf("failed to set collection error: %w", err) @@ -448,40 +457,73 @@ func (s *PostgresStore) SetRecommendationsCollectionError(ctx context.Context, e // only when no in-flight collection is currently running. The WHERE clause // treats a started_at older than 5 minutes as stale (the scheduler Lambda // must have crashed) so a new collection can proceed rather than being -// permanently blocked. +// permanently blocked. A fresh owner token is stamped into +// last_collection_owner_id alongside the timestamp so the caller that wins +// the race is the only one that can later clear the marker (issue #261 +// compare-and-clear guard; see ClearCollectionStarted). // -// Returns true when this caller won the race (rowsAffected == 1) and should -// proceed with the async invoke. Returns false when another collection is -// already in flight (rowsAffected == 0), signaling the handler to return -// 409 Conflict. -func (s *PostgresStore) MarkCollectionStarted(ctx context.Context) (bool, error) { +// Returns the token and true when this caller won the race (rowsAffected == +// 1) and should proceed with the async invoke, threading the token through +// so it can be passed to ClearCollectionStarted later. Returns ("", false, +// nil) when another collection is already in flight (rowsAffected == 0), +// signaling the handler to return 409 Conflict. +func (s *PostgresStore) MarkCollectionStarted(ctx context.Context) (token string, ok bool, err error) { + token = uuid.New().String() tag, err := s.db.Exec(ctx, ` UPDATE recommendations_state - SET last_collection_started_at = NOW() + SET last_collection_started_at = NOW(), + last_collection_owner_id = $1 WHERE id = 1 AND ( last_collection_started_at IS NULL OR last_collection_started_at < NOW() - INTERVAL '5 minutes' ) - `) + `, token) if err != nil { - return false, fmt.Errorf("failed to mark collection started: %w", err) + return "", false, fmt.Errorf("failed to mark collection started: %w", err) + } + if tag.RowsAffected() != 1 { + return "", false, nil } - return tag.RowsAffected() == 1, nil + return token, true, nil } -// ClearCollectionStarted clears last_collection_started_at so the frontend -// knows an async collection has finished. Called by the scheduler on both -// success and failure paths. On the success path, last_collected_at and -// last_collection_error are updated by UpsertRecommendations/ReplaceRecommendations, -// so this method only touches started_at. -func (s *PostgresStore) ClearCollectionStarted(ctx context.Context) error { - if _, err := s.db.Exec(ctx, ` +// ClearCollectionStarted clears last_collection_started_at (and the owner +// token) so the frontend knows an async collection has finished. Called by +// the scheduler on both success and failure paths. On the success path, +// last_collected_at and last_collection_error are updated by +// UpsertRecommendations/ReplaceRecommendations, so this method only touches +// started_at and the owner column. +// +// The clear is scoped to rows where last_collection_owner_id still matches +// token (issue #261): a caller whose token no longer matches (another run +// has since started and won the race) has nothing left to clear and this is +// a documented silent no-op, not an error. An empty token is a boundary +// error: only a caller that actually won MarkCollectionStarted should ever +// call Clear; callers with no marker to own (cron, cold-start) must skip +// the call entirely rather than pass an empty token. +func (s *PostgresStore) ClearCollectionStarted(ctx context.Context, token string) error { + if token == "" { + return fmt.Errorf("owner token must not be empty") + } + tag, err := s.db.Exec(ctx, ` UPDATE recommendations_state - SET last_collection_started_at = NULL + SET last_collection_started_at = NULL, + last_collection_owner_id = NULL WHERE id = 1 - `); err != nil { + AND last_collection_owner_id = $1 + `, token) + if err != nil { return fmt.Errorf("failed to clear collection started: %w", err) } + // The no-op branch IS the safety mechanism issue #261 adds, so surface it + // at debug level: without this there is no signal in production + // distinguishing "cleared" from "declined to clear someone else's marker", + // and a guard that never fires looks identical to a guard that is broken. + // The token is deliberately not logged: it is the capability that controls + // the marker, and the repo forbids putting token material in logs. + if tag.RowsAffected() == 0 { + logging.Debugf("ClearCollectionStarted: no-op, caller's token no longer owns the collection marker") + } return nil } diff --git a/internal/config/store_postgres_recommendations_test.go b/internal/config/store_postgres_recommendations_test.go index dfd3eb8ac..445137eda 100644 --- a/internal/config/store_postgres_recommendations_test.go +++ b/internal/config/store_postgres_recommendations_test.go @@ -12,6 +12,7 @@ import ( "github.com/LeanerCloud/CUDly/internal/config" "github.com/LeanerCloud/CUDly/internal/database/postgres/migrations" "github.com/LeanerCloud/CUDly/internal/database/postgres/testhelpers" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -393,4 +394,118 @@ func TestPostgresStore_UpsertRecommendations_AmbientAndRegisteredCoexist(t *test assert.True(t, registeredFound, "registered row must be upserted") } +// TestPostgresStore_ClearCollectionStarted_CompareAndClear pins the issue +// #261 compare-and-clear guard at the store level: a late clear from an +// abandoned run (run A) must not wipe the marker a later run (run B) now +// owns. Pre-fix, ClearCollectionStarted took no token and cleared +// unconditionally, so run A's deferred clear (arriving after run B had +// already started) would wipe run B's in-flight marker. +// +// Both runs here go through MarkCollectionStarted, which only the +// POST /api/recommendations/refresh handler calls, so run B models a second +// user-triggered refresh taking over after run A's 5-minute window lapses. +// The cron variant of the same race (a tokenless scheduler run clearing a +// user run's marker) cannot be expressed at this level because cron never +// marks; it is pinned in the scheduler by +// TestScheduler_CollectRecommendations_EmptyTokenSkipsClear. +func TestPostgresStore_ClearCollectionStarted_CompareAndClear(t *testing.T) { + ctx := context.Background() + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + defer container.Cleanup(ctx) + pool := container.DB.Pool() + require.NoError(t, migrations.RunMigrations(ctx, pool, getMigrationsPath(), "", "")) + store := config.NewPostgresStore(container.DB) + + // Run A wins the race and is stamped with a fresh owner token. + tokenA, okA, err := store.MarkCollectionStarted(ctx) + require.NoError(t, err) + require.True(t, okA) + _, uuidErr := uuid.Parse(tokenA) + require.NoError(t, uuidErr, "MarkCollectionStarted must return a UUID token") + + // A concurrent Mark while A's marker is fresh is blocked (409 path). + tokenNone, okBlocked, err := store.MarkCollectionStarted(ctx) + require.NoError(t, err) + assert.False(t, okBlocked) + assert.Empty(t, tokenNone) + + // Force run A's marker stale (simulates its Lambda crashing or running + // long past the 5-minute window), so the next refresh (run B) is allowed + // to take the marker over. + _, err = pool.Exec(ctx, ` + UPDATE recommendations_state + SET last_collection_started_at = NOW() - INTERVAL '6 minutes' + WHERE id = 1 + `) + require.NoError(t, err) + + tokenB, okB, err := store.MarkCollectionStarted(ctx) + require.NoError(t, err) + require.True(t, okB) + require.NotEqual(t, tokenA, tokenB, "run B must win its own distinct token") + + // The regression assertion: run A finally finishes (late, after its own + // stale window) and calls Clear with its own token. Because + // last_collection_owner_id now points at run B, this must be a silent + // no-op: run B's in-flight marker must survive. + require.NoError(t, store.ClearCollectionStarted(ctx, tokenA)) + freshness, err := store.GetRecommendationsFreshness(ctx) + require.NoError(t, err) + assert.NotNil(t, freshness.LastCollectionStartedAt, + "run B's marker must survive run A's stale-token clear") + + // Run B's own clear matches the current owner and actually clears. + require.NoError(t, store.ClearCollectionStarted(ctx, tokenB)) + freshness, err = store.GetRecommendationsFreshness(ctx) + require.NoError(t, err) + assert.Nil(t, freshness.LastCollectionStartedAt) + + // Empty token is a boundary error: only a caller that actually won + // MarkCollectionStarted should ever call Clear. + err = store.ClearCollectionStarted(ctx, "") + require.Error(t, err) +} + +// TestPostgresStore_SetRecommendationsCollectionError_DoesNotClearOwnerMarker +// pins a second issue #261 side door found in adversarial review: +// SetRecommendationsCollectionError is called mid-run from persistCollection +// on EVERY CollectRecommendations invocation that hits a provider error, +// including tokenless cron/cold-start/background runs. Pre-fix, it cleared +// last_collection_started_at unconditionally (no owner check), so a +// tokenless run's routine provider error would wipe a concurrent owner run's +// in-flight marker, reopening the exact race the compare-and-clear guard +// exists to close. +func TestPostgresStore_SetRecommendationsCollectionError_DoesNotClearOwnerMarker(t *testing.T) { + ctx := context.Background() + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + defer container.Cleanup(ctx) + pool := container.DB.Pool() + require.NoError(t, migrations.RunMigrations(ctx, pool, getMigrationsPath(), "", "")) + store := config.NewPostgresStore(container.DB) + + // Owner run A wins the race and is stamped with a fresh owner token. + tokenA, okA, err := store.MarkCollectionStarted(ctx) + require.NoError(t, err) + require.True(t, okA) + + // A tokenless caller (cron/cold-start/background run hitting a routine + // provider error) records the error. It must NOT touch run A's marker. + require.NoError(t, store.SetRecommendationsCollectionError(ctx, "aws: transient throttling")) + + freshness, err := store.GetRecommendationsFreshness(ctx) + require.NoError(t, err) + require.NotNil(t, freshness.LastCollectionError) + assert.Equal(t, "aws: transient throttling", *freshness.LastCollectionError) + assert.NotNil(t, freshness.LastCollectionStartedAt, + "run A's marker must survive a tokenless SetRecommendationsCollectionError call") + + // Run A's own token-guarded clear still works afterward. + require.NoError(t, store.ClearCollectionStarted(ctx, tokenA)) + freshness, err = store.GetRecommendationsFreshness(ctx) + require.NoError(t, err) + assert.Nil(t, freshness.LastCollectionStartedAt) +} + func float64Ptr(f float64) *float64 { return &f } diff --git a/internal/database/postgres/migrations/000093_recommendations_state_owner_id.down.sql b/internal/database/postgres/migrations/000093_recommendations_state_owner_id.down.sql new file mode 100644 index 000000000..b381bd563 --- /dev/null +++ b/internal/database/postgres/migrations/000093_recommendations_state_owner_id.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE recommendations_state + DROP COLUMN IF EXISTS last_collection_owner_id; diff --git a/internal/database/postgres/migrations/000093_recommendations_state_owner_id.up.sql b/internal/database/postgres/migrations/000093_recommendations_state_owner_id.up.sql new file mode 100644 index 000000000..50b117cab --- /dev/null +++ b/internal/database/postgres/migrations/000093_recommendations_state_owner_id.up.sql @@ -0,0 +1,17 @@ +-- Add last_collection_owner_id to recommendations_state. +-- +-- Issue #261: the scheduler's defer-clear of last_collection_started_at +-- (added in 000047) unconditionally wipes the marker on every run, +-- including cron runs that never called MarkCollectionStarted. If a cron +-- run overlaps a user-triggered async collection, the cron run's clear +-- silently erases the user run's in-flight marker, letting a second +-- concurrent collection start. +-- +-- This column pairs with last_collection_started_at as a compare-and-clear +-- guard: MarkCollectionStarted stamps a fresh owner token alongside the +-- timestamp, and ClearCollectionStarted only clears when the caller's +-- token still matches (WHERE last_collection_owner_id = $1). A caller that +-- never owned a marker (cron, cold-start) carries no token and skips the +-- clear entirely rather than clearing unconditionally. +ALTER TABLE recommendations_state + ADD COLUMN IF NOT EXISTS last_collection_owner_id UUID; diff --git a/internal/database/postgres/migrations/000093_recommendations_state_owner_id_test.go b/internal/database/postgres/migrations/000093_recommendations_state_owner_id_test.go new file mode 100644 index 000000000..77791fed2 --- /dev/null +++ b/internal/database/postgres/migrations/000093_recommendations_state_owner_id_test.go @@ -0,0 +1,69 @@ +//go:build integration +// +build integration + +package migrations_test + +import ( + "context" + "testing" + + "github.com/LeanerCloud/CUDly/internal/database/postgres/migrations" + "github.com/LeanerCloud/CUDly/internal/database/postgres/testhelpers" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestMigration_RecommendationsStateOwnerID locks down migration 000093: +// recommendations_state gains a nullable last_collection_owner_id UUID +// column (issue #261's compare-and-clear guard), and the down migration +// removes it cleanly. +func TestMigration_RecommendationsStateOwnerID(t *testing.T) { + ctx := context.Background() + migrationsPath := getMigrationsPath() + + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + defer container.Cleanup(ctx) + pool := container.DB.Pool() + + // Pin at 000092 (pre-migration) so the assertions exercise 000093's + // direct effect rather than whatever migration happens to be newest. + require.NoError(t, migrations.MigrateToVersion(ctx, pool, migrationsPath, 92)) + + var columnExists bool + err = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'recommendations_state' + AND column_name = 'last_collection_owner_id' + ) + `).Scan(&columnExists) + require.NoError(t, err) + assert.False(t, columnExists, "last_collection_owner_id must not exist before 000093") + + require.NoError(t, migrations.MigrateToVersion(ctx, pool, migrationsPath, 93)) + + var dataType, isNullable string + err = pool.QueryRow(ctx, ` + SELECT data_type, is_nullable + FROM information_schema.columns + WHERE table_name = 'recommendations_state' + AND column_name = 'last_collection_owner_id' + `).Scan(&dataType, &isNullable) + require.NoError(t, err, "last_collection_owner_id must exist after 000093") + assert.Equal(t, "uuid", dataType) + assert.Equal(t, "YES", isNullable) + + // Rollback restores the pre-migration schema. + require.NoError(t, migrations.MigrateToVersion(ctx, pool, migrationsPath, 92)) + + err = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'recommendations_state' + AND column_name = 'last_collection_owner_id' + ) + `).Scan(&columnExists) + require.NoError(t, err) + assert.False(t, columnExists, "last_collection_owner_id must be dropped after rollback to 92") +} diff --git a/internal/mocks/stores.go b/internal/mocks/stores.go index 9e283cb64..a7d4a53f4 100644 --- a/internal/mocks/stores.go +++ b/internal/mocks/stores.go @@ -13,6 +13,11 @@ import ( "github.com/stretchr/testify/mock" ) +// MockOwnerToken is the default token MockConfigStore.MarkCollectionStarted +// returns when no expectation is registered, so existing tests that call it +// implicitly keep getting a non-empty token to thread through. +const MockOwnerToken = "mock-owner-token" + // MockConfigStore is a shared testify-based mock for config.StoreInterface. // // Most methods dispatch through m.Called only when an expectation has been @@ -1339,22 +1344,22 @@ func (m *MockConfigStore) GetScheduledExecutionsDue(ctx context.Context) ([]conf } // MarkCollectionStarted mocks the MarkCollectionStarted operation. -// Defaults to (true, nil) when no expectation is registered. -func (m *MockConfigStore) MarkCollectionStarted(ctx context.Context) (bool, error) { +// Defaults to (MockOwnerToken, true, nil) when no expectation is registered. +func (m *MockConfigStore) MarkCollectionStarted(ctx context.Context) (token string, ok bool, err error) { if !isExpected(&m.Mock, "MarkCollectionStarted") { - return true, nil + return MockOwnerToken, true, nil } args := m.Called(ctx) - return args.Bool(0), args.Error(1) + return args.String(0), args.Bool(1), args.Error(2) } // ClearCollectionStarted mocks the ClearCollectionStarted operation. // Defaults to nil when no expectation is registered. -func (m *MockConfigStore) ClearCollectionStarted(ctx context.Context) error { +func (m *MockConfigStore) ClearCollectionStarted(ctx context.Context, token string) error { if !isExpected(&m.Mock, "ClearCollectionStarted") { return nil } - return m.Called(ctx).Error(0) + return m.Called(ctx, token).Error(0) } // StampRIExchangeApprovedBy mocks the StampRIExchangeApprovedBy operation. diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index b7c0dfa03..abb427cb9 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -141,26 +141,35 @@ func cacheTTLFromEnv() time.Duration { // Persists results to the recommendations cache so read handlers can serve // from SQL instead of re-fetching live. // -// Bookkeeping: always clears last_collection_started_at on exit (success or -// failure) so the frontend polling loop can detect completion. The scheduler -// is invoked either by the cron EventBridge rule or by an async self-invoke -// from the POST /api/recommendations/refresh handler. In the async case, -// MarkCollectionStarted has already set last_collection_started_at; the -// cron case leaves it NULL (no async-invoke bookkeeping for cron runs, which -// are expected and not user-triggered). -func (s *Scheduler) CollectRecommendations(ctx context.Context) (*CollectResult, error) { +// Bookkeeping: clears last_collection_started_at on exit (success or +// failure) so the frontend polling loop can detect completion, but only if +// ownerToken is non-empty and still matches last_collection_owner_id +// (issue #261 compare-and-clear guard). The scheduler is invoked either by +// the cron EventBridge rule or by an async self-invoke from the POST +// /api/recommendations/refresh handler. In the async case, +// MarkCollectionStarted has already set last_collection_started_at and +// returned the token that ownerToken carries here; the cron case (and +// cold-start) call with an empty ownerToken since they never won a marker, +// and the clear is skipped entirely rather than clearing unconditionally, +// which previously let a cron run wipe a concurrent user-triggered run's +// marker. This deferred clear is the ONLY place that touches started_at: +// persistCollection's SetRecommendationsCollectionError call (below, on a +// provider failure) intentionally leaves started_at alone, so a tokenless +// run hitting a routine provider error cannot wipe another run's marker +// either. +func (s *Scheduler) CollectRecommendations(ctx context.Context, ownerToken string) (*CollectResult, error) { logging.Info("Collecting recommendations from cloud providers...") - // Always clear last_collection_started_at on exit so the frontend knows - // the collection has finished. Use a fresh background context with a - // short timeout: the request ctx may already be canceled by the time - // the defer runs (e.g. caller deadline expired during a slow collect), + // Clear last_collection_started_at on exit so the frontend knows the + // collection has finished. Use a fresh background context with a short + // timeout: the request ctx may already be canceled by the time the + // defer runs (e.g. caller deadline expired during a slow collect), // which would cause ClearCollectionStarted to fail and leave the // "in flight" marker until the 5-min auto-recovery window kicks in. defer func() { clearCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - s.clearCollectionStartedBestEffort(clearCtx) + s.clearCollectionStartedBestEffort(clearCtx, ownerToken) }() // Get global config @@ -359,8 +368,16 @@ func (s *Scheduler) collectAllProviders(ctx context.Context, globalCfg *config.G // scheduler's exit path. Best-effort — a failure here is logged but does not // prevent returning the collection result. Extracted so CollectRecommendations // stays under the cyclomatic-complexity gate. -func (s *Scheduler) clearCollectionStartedBestEffort(ctx context.Context) { - if err := s.config.ClearCollectionStarted(ctx); err != nil { +// +// A caller that never won MarkCollectionStarted (cron, cold-start) passes +// an empty ownerToken and by design owns no marker to clear; this is an +// explicit, logged skip, not a silent fallback. +func (s *Scheduler) clearCollectionStartedBestEffort(ctx context.Context, ownerToken string) { + if ownerToken == "" { + logging.Debugf("skipping collection-started clear: caller holds no owner token (cron/cold-start run)") + return + } + if err := s.config.ClearCollectionStarted(ctx, ownerToken); err != nil { logging.Errorf("failed to clear collection started: %v", err) } } @@ -929,7 +946,7 @@ func (s *Scheduler) ListRecommendations(ctx context.Context, filter config.Recom if freshness.LastCollectedAt == nil { logging.Info("Recommendations cache is empty; performing synchronous cold-start collect") - _, collectErr := s.CollectRecommendations(ctx) + _, collectErr := s.CollectRecommendations(ctx, "") if collectErr != nil { return nil, fmt.Errorf("cold-start collect failed: %w", collectErr) } @@ -1202,7 +1219,7 @@ func (s *Scheduler) maybeKickBackgroundRefresh(freshness *config.Recommendations logging.Errorf("background recommendations refresh panic: %v", r) } }() - if _, err := s.CollectRecommendations(bgCtx); err != nil { + if _, err := s.CollectRecommendations(bgCtx, ""); err != nil { // CollectRecommendations already surfaces errors via // recommendations_state.last_collection_error, so just log // locally here for operator visibility. diff --git a/internal/scheduler/scheduler_overrides_test.go b/internal/scheduler/scheduler_overrides_test.go index 01119cdcb..5ba7e3686 100644 --- a/internal/scheduler/scheduler_overrides_test.go +++ b/internal/scheduler/scheduler_overrides_test.go @@ -42,10 +42,10 @@ func (m *mockOverrideStore) GetRecommendationsFreshness(_ context.Context) (*con now := time.Now() return &config.RecommendationsFreshness{LastCollectedAt: &now}, nil } -func (m *mockOverrideStore) MarkCollectionStarted(_ context.Context) (bool, error) { - return true, nil +func (m *mockOverrideStore) MarkCollectionStarted(_ context.Context) (string, bool, error) { + return "mock-owner-token", true, nil } -func (m *mockOverrideStore) ClearCollectionStarted(_ context.Context) error { +func (m *mockOverrideStore) ClearCollectionStarted(_ context.Context, _ string) error { return nil } func (m *mockOverrideStore) GetServiceConfig(_ context.Context, provider, service string) (*config.ServiceConfig, error) { diff --git a/internal/scheduler/scheduler_suppressions_test.go b/internal/scheduler/scheduler_suppressions_test.go index 70b11c66e..1ec7879c6 100644 --- a/internal/scheduler/scheduler_suppressions_test.go +++ b/internal/scheduler/scheduler_suppressions_test.go @@ -31,10 +31,10 @@ func (m *mockSuppressionStore) GetRecommendationsFreshness(_ context.Context) (* now := time.Now() return &config.RecommendationsFreshness{LastCollectedAt: &now}, nil } -func (m *mockSuppressionStore) MarkCollectionStarted(_ context.Context) (bool, error) { - return true, nil +func (m *mockSuppressionStore) MarkCollectionStarted(_ context.Context) (string, bool, error) { + return "mock-owner-token", true, nil } -func (m *mockSuppressionStore) ClearCollectionStarted(_ context.Context) error { +func (m *mockSuppressionStore) ClearCollectionStarted(_ context.Context, _ string) error { return nil } diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 931096829..6ee6249f3 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -226,7 +226,7 @@ func TestScheduler_CollectRecommendations_NoProviders(t *testing.T) { dashboardURL: "https://dashboard.example.com", } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, 0, result.Recommendations) @@ -257,13 +257,77 @@ func TestScheduler_CollectRecommendations_AWSProvider(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) // Provider returns error, so no recommendations assert.Equal(t, 0, result.Recommendations) } +// TestScheduler_CollectRecommendations_EmptyTokenSkipsClear pins issue #261: +// a caller with no owner token (cron, cold-start) never won +// MarkCollectionStarted and must not call ClearCollectionStarted at all, +// rather than clearing unconditionally and risking wiping another run's +// marker. +func TestScheduler_CollectRecommendations_EmptyTokenSkipsClear(t *testing.T) { + ctx := context.Background() + mockStore := new(MockConfigStore) + t.Cleanup(func() { mockStore.AssertExpectations(t) }) + mockEmail := new(MockEmailSender) + + globalCfg := &config.GlobalConfig{EnabledProviders: []string{}} + mockStore.On("GetGlobalConfig", ctx).Return(globalCfg, nil) + // The expectation is registered (as Maybe) purely so an unwanted call is + // RECORDED. MockConfigStore short-circuits methods with no registered + // expectation before reaching mock.Called, so the call never lands in + // m.Calls and the AssertNotCalled below would pass vacuously: deleting the + // empty-token guard in clearCollectionStartedBestEffort kept this test + // green until this line was added. + mockStore.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe() + + scheduler := &Scheduler{ + config: mockStore, + email: mockEmail, + dashboardURL: "https://dashboard.example.com", + } + + _, err := scheduler.CollectRecommendations(ctx, "") + require.NoError(t, err) + + mockStore.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything) +} + +// TestScheduler_CollectRecommendations_ClearsWithOwnerToken pins issue #261: +// a caller holding an owner token (the async self-invoke path) must clear +// with that exact token, using a background context carrying a deadline +// (not the possibly-canceled request ctx) so the clear survives a caller +// timeout. +func TestScheduler_CollectRecommendations_ClearsWithOwnerToken(t *testing.T) { + ctx := context.Background() + mockStore := new(MockConfigStore) + t.Cleanup(func() { mockStore.AssertExpectations(t) }) + mockEmail := new(MockEmailSender) + + globalCfg := &config.GlobalConfig{EnabledProviders: []string{}} + mockStore.On("GetGlobalConfig", ctx).Return(globalCfg, nil) + mockStore.On("ClearCollectionStarted", + mock.MatchedBy(func(clearCtx context.Context) bool { + _, hasDeadline := clearCtx.Deadline() + return hasDeadline + }), + "tok-1", + ).Return(nil) + + scheduler := &Scheduler{ + config: mockStore, + email: mockEmail, + dashboardURL: "https://dashboard.example.com", + } + + _, err := scheduler.CollectRecommendations(ctx, "tok-1") + require.NoError(t, err) +} + func TestScheduler_CollectRecommendations_AllProviders(t *testing.T) { ctx := context.Background() mockStore := new(MockConfigStore) @@ -288,7 +352,7 @@ func TestScheduler_CollectRecommendations_AllProviders(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, 0, result.Recommendations) @@ -356,7 +420,7 @@ func TestScheduler_CollectRecommendations_ParallelProviders(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, []string{"gcp", "azure"}, result.SuccessfulProviders, @@ -391,7 +455,7 @@ func TestScheduler_CollectRecommendations_ParallelProviders(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() // cancel BEFORE the call - _, err := scheduler.CollectRecommendations(ctx) + _, err := scheduler.CollectRecommendations(ctx, "") require.Error(t, err, "expected context.Canceled to propagate from CollectRecommendations") assert.ErrorIs(t, err, context.Canceled, "CollectRecommendations must propagate the parent ctx error after the provider fan-out's g.Wait()") @@ -417,7 +481,7 @@ func TestScheduler_CollectRecommendations_UnknownProvider(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, 0, result.Recommendations) @@ -556,7 +620,7 @@ func TestScheduler_CollectRecommendations_WithNotification(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, 0, result.Recommendations) @@ -595,7 +659,7 @@ func TestScheduler_CollectRecommendations_ConfigError(t *testing.T) { dashboardURL: "https://dashboard.example.com", } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") assert.Error(t, err) assert.Nil(t, result) } @@ -1653,7 +1717,7 @@ func TestScheduler_CollectRecommendations_WithSuccessfulRecs(t *testing.T) { providerFactory: mockFactory, } - result, err := scheduler.CollectRecommendations(ctx) + result, err := scheduler.CollectRecommendations(ctx, "") require.NoError(t, err) assert.Equal(t, 1, result.Recommendations) diff --git a/internal/server/app_test.go b/internal/server/app_test.go index 700e6f43a..c27788cf6 100644 --- a/internal/server/app_test.go +++ b/internal/server/app_test.go @@ -230,7 +230,7 @@ func TestHandleProcessScheduledPurchases_Error(t *testing.T) { }, } - _, err := app.HandleScheduledTask(ctx, TaskProcessScheduledPurchases) + _, err := app.HandleScheduledTask(ctx, TaskProcessScheduledPurchases, ScheduledTaskParams{}) testutil.AssertError(t, err) } @@ -244,7 +244,7 @@ func TestHandleSendNotifications_Error(t *testing.T) { }, } - _, err := app.HandleScheduledTask(ctx, TaskSendNotifications) + _, err := app.HandleScheduledTask(ctx, TaskSendNotifications, ScheduledTaskParams{}) testutil.AssertError(t, err) } @@ -252,7 +252,7 @@ func TestHandleCollectRecommendations_WithResults(t *testing.T) { ctx := testutil.TestContext(t) app := &Application{ Scheduler: &testutil.MockScheduler{ - CollectRecommendationsFunc: func(ctx context.Context) (*scheduler.CollectResult, error) { + CollectRecommendationsFunc: func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{ Recommendations: 15, TotalSavings: 2500.50, @@ -261,7 +261,7 @@ func TestHandleCollectRecommendations_WithResults(t *testing.T) { }, } - result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations) + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) testutil.AssertNoError(t, err) testutil.AssertTrue(t, result != nil, "Result should not be nil") } diff --git a/internal/server/handler.go b/internal/server/handler.go index 51a662872..a6a74499e 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -3,14 +3,22 @@ package server import ( "context" "encoding/json" + "errors" "fmt" "hash/fnv" "log" + "time" "github.com/LeanerCloud/CUDly/internal/purchase" "github.com/LeanerCloud/CUDly/internal/scheduler" + "github.com/google/uuid" ) +// skippedCollectionMarkerClearTimeout bounds the detached best-effort clear +// issued when a marker-owning collect run is skipped by the advisory lock. +// Matches the scheduler's deferred clear budget. +const skippedCollectionMarkerClearTimeout = 5 * time.Second + // TaskLocker abstracts advisory lock operations for scheduled task concurrency control. type TaskLocker interface { TryAdvisoryLock(ctx context.Context, lockID int64) (bool, error) @@ -81,7 +89,7 @@ var scheduledEventActions = map[string]ScheduledTaskType{ // HandleScheduledTask processes a scheduled task by type. // It acquires a PostgreSQL advisory lock to prevent concurrent execution of the same task. -func (app *Application) HandleScheduledTask(ctx context.Context, taskType ScheduledTaskType) (any, error) { +func (app *Application) HandleScheduledTask(ctx context.Context, taskType ScheduledTaskType, params ScheduledTaskParams) (any, error) { log.Printf("Handling scheduled task: %q", taskType) // #nosec G706 -- taskType is looked up from a known-value map; %q quotes the value to prevent CR/LF log injection if err := app.ensureDB(ctx); err != nil { @@ -97,37 +105,104 @@ func (app *Application) HandleScheduledTask(ctx context.Context, taskType Schedu } if !acquired { log.Printf("Task %q already running (advisory lock held), skipping", taskType) // #nosec G706 -- taskType is looked up from a known-value map; %q quotes the value to prevent CR/LF log injection + app.releaseSkippedCollectionMarker(taskType, params.OwnerToken) return map[string]string{"status": "skipped", "reason": "already_running"}, nil } defer locker.ReleaseAdvisoryLock(ctx, lockID) } - return app.dispatchTask(ctx, taskType) + return app.dispatchTask(ctx, taskType, params) +} + +// releaseSkippedCollectionMarker releases this run's OWN collection in-flight +// marker when a TaskCollectRecommendations invocation that holds an owner +// token is skipped by the advisory lock, and therefore never reaches +// CollectRecommendations, whose deferred token-scoped clear would normally +// release it. +// +// Without this the abandoned run's marker sits stranded until the 5-minute +// auto-recovery window in MarkCollectionStarted expires, and every refresh +// the user attempts during that window is rejected with 409 "collection +// already in progress" while no collection backed by that marker is running. +// Before the issue #261 compare-and-clear guard, the overlapping run holding +// the lock happened to cover this case with its unconditional clear; scoping +// the clear to its owner correctly stopped that cross-run wipe, so the +// abandoning run now has to release its own marker explicitly. +// +// Only the lock-skip path releases. A lock-check error returns an error, which +// lets the Lambda async-invoke retry the same event (same owner token) and +// still run the collect, so the marker must survive that path. +// +// Known residual gap: the token identifies the MARKER, not the invocation, so +// this cannot distinguish "the lock is held by a tokenless cron run" (release +// is required, or the marker strands for the full 5-minute window) from "the +// lock is held by a concurrent duplicate delivery of my own event" (release is +// premature, since that sibling carries the same token and is still +// collecting). Lambda's async invocation is at-least-once, so the second case +// is reachable, and there it clears the marker mid-run: the frontend banner +// drops early and a refresh issued during the remainder wins a fresh marker +// only to be lock-skipped and released again, returning 202 without collecting. +// It is bounded and self-healing (the next refresh after the run completes +// behaves normally) and touches no purchase or money path. Releasing is still +// strictly better than not releasing, because the cron-overlap case is routine +// while duplicate delivery is rare. Closing it properly needs an +// invocation-scoped identity distinct from the marker token (e.g. the lock +// winner re-stamping last_collection_owner_id with a fresh token), which is a +// design change deliberately left out of this PR. +// +// Scoped by ownerToken, so this can only ever release the marker this run +// owns, never a concurrent run's. Callers that never won MarkCollectionStarted +// (EventBridge cron, the /api/scheduled/ HTTP path, the --task CLI) carry no +// token, own no marker, and are skipped. Best effort: a failure only defers +// cleanup to the 5-minute window, so it is logged rather than failing the +// task. Uses a detached short-deadline context for the same reason the +// scheduler's deferred clear does: ctx may already be near its deadline by +// the time cleanup runs. +func (app *Application) releaseSkippedCollectionMarker(taskType ScheduledTaskType, ownerToken string) { + if taskType != TaskCollectRecommendations || ownerToken == "" { + return + } + if app.Config == nil { + log.Println("Cannot release collection marker for skipped run: no config store configured") + return + } + clearCtx, cancel := context.WithTimeout(context.Background(), skippedCollectionMarkerClearTimeout) + defer cancel() + if err := app.Config.ClearCollectionStarted(clearCtx, ownerToken); err != nil { + log.Printf("Failed to release collection marker for skipped run: %v", err) + } } // dispatchTask routes a scheduled task to its handler. -func (app *Application) dispatchTask(ctx context.Context, taskType ScheduledTaskType) (any, error) { +func (app *Application) dispatchTask(ctx context.Context, taskType ScheduledTaskType, params ScheduledTaskParams) (any, error) { // Map-based dispatch (rather than a switch) keeps this function under the // cyclomatic-complexity limit as the task roster grows. Each handler adapts - // its concrete return type to (any, error) at the call site. - handlers := map[ScheduledTaskType]func(context.Context) (any, error){ - TaskCollectRecommendations: func(c context.Context) (any, error) { return app.handleCollectRecommendations(c) }, - TaskProcessScheduledPurchases: func(c context.Context) (any, error) { return app.handleProcessScheduledPurchases(c) }, - TaskSendNotifications: func(c context.Context) (any, error) { return app.handleSendNotifications(c) }, - TaskCleanupExpiredRecords: func(c context.Context) (any, error) { return app.handleCleanupExpiredRecords(c) }, - TaskRefreshAnalytics: func(c context.Context) (any, error) { return app.handleRefreshAnalytics(c) }, - TaskCollectAnalytics: func(c context.Context) (any, error) { return app.handleCollectAnalytics(c) }, - TaskRIExchangeReshape: func(c context.Context) (any, error) { return app.handleRIExchangeReshape(c) }, - TaskReapStuckPurchases: func(c context.Context) (any, error) { return app.handleReapStuckPurchases(c) }, - TaskFireScheduledPurchases: func(c context.Context) (any, error) { return app.handleFireScheduledPurchases(c) }, - TaskFinalizeRevocations: func(c context.Context) (any, error) { return app.handleFinalizeRevocations(c) }, - TaskLadderRun: func(c context.Context) (any, error) { return app.handleLadderRun(c) }, + // its concrete return type to (any, error) at the call site. Only + // TaskCollectRecommendations reads params; every other closure ignores it. + handlers := map[ScheduledTaskType]func(context.Context, ScheduledTaskParams) (any, error){ + TaskCollectRecommendations: func(c context.Context, p ScheduledTaskParams) (any, error) { + return app.handleCollectRecommendations(c, p.OwnerToken) + }, + TaskProcessScheduledPurchases: func(c context.Context, _ ScheduledTaskParams) (any, error) { + return app.handleProcessScheduledPurchases(c) + }, + TaskSendNotifications: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleSendNotifications(c) }, + TaskCleanupExpiredRecords: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleCleanupExpiredRecords(c) }, + TaskRefreshAnalytics: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleRefreshAnalytics(c) }, + TaskCollectAnalytics: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleCollectAnalytics(c) }, + TaskRIExchangeReshape: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleRIExchangeReshape(c) }, + TaskReapStuckPurchases: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleReapStuckPurchases(c) }, + TaskFireScheduledPurchases: func(c context.Context, _ ScheduledTaskParams) (any, error) { + return app.handleFireScheduledPurchases(c) + }, + TaskFinalizeRevocations: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleFinalizeRevocations(c) }, + TaskLadderRun: func(c context.Context, _ ScheduledTaskParams) (any, error) { return app.handleLadderRun(c) }, } handler, ok := handlers[taskType] if !ok { return nil, fmt.Errorf("unknown scheduled task type: %s", taskType) } - return handler(ctx) + return handler(ctx, params) } // taskLocker returns the configured TaskLocker, falling back to DB if set. @@ -149,9 +224,11 @@ func taskLockID(taskType ScheduledTaskType) int64 { } // handleCollectRecommendations collects cost optimization recommendations. -func (app *Application) handleCollectRecommendations(ctx context.Context) (*scheduler.CollectResult, error) { +// ownerToken threads through to CollectRecommendations's compare-and-clear +// guard (issue #261); it is empty for cron-triggered runs. +func (app *Application) handleCollectRecommendations(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { log.Println("Collecting recommendations...") - result, err := app.Scheduler.CollectRecommendations(ctx) + result, err := app.Scheduler.CollectRecommendations(ctx, ownerToken) if err != nil { log.Printf("Failed to collect recommendations: %v", err) return nil, err @@ -333,18 +410,56 @@ type ScheduledEvent struct { DetailType string `json:"detail-type"` Action string `json:"action"` Detail json.RawMessage `json:"detail"` + // OwnerToken carries the collection-marker owner token (issue #261 + // compare-and-clear guard) for TaskCollectRecommendations events fired + // by the async self-invoke in handler_recommendations_refresh.go. + // EventBridge cron deliveries and the --task CLI never set this field, + // so it decodes to "" and the scheduler treats the run as owning no + // marker (see ScheduledTaskParams / clearCollectionStartedBestEffort). + OwnerToken string `json:"owner_token"` +} + +// ScheduledTaskParams carries per-task-type data extracted from a +// ScheduledEvent. Only TaskCollectRecommendations currently reads a field +// (OwnerToken); other task types receive a zero-value struct. +type ScheduledTaskParams struct { + OwnerToken string } -// ParseScheduledEvent parses a scheduled event and returns the task type. -func ParseScheduledEvent(rawEvent json.RawMessage) (ScheduledTaskType, error) { +// ParseScheduledEvent parses a scheduled event and returns the task type +// plus any per-task parameters. +func ParseScheduledEvent(rawEvent json.RawMessage) (ScheduledTaskType, ScheduledTaskParams, error) { var event ScheduledEvent if err := json.Unmarshal(rawEvent, &event); err != nil { - return "", fmt.Errorf("failed to parse scheduled event: %w", err) + return "", ScheduledTaskParams{}, fmt.Errorf("failed to parse scheduled event: %w", err) + } + + // Validate the owner token at the boundary rather than letting a malformed + // value reach the UUID-typed persistence layer. Empty is legitimate and + // expected (EventBridge cron, the /api/scheduled/ HTTP path, the --task + // CLI): those callers never won MarkCollectionStarted and own no marker. + // A non-empty value, though, can only have come from asyncInvokeSelf, + // which always sends a uuid.New(), so anything else means the payload is + // corrupt. Failing loud here is deliberate: a token that cannot match any + // owner would leave the marker it belongs to stranded for the full + // 5-minute recovery window anyway, so a buried error log on the eventual + // clear is strictly worse than refusing the malformed event outright. + // + // uuid.Parse's error is deliberately NOT wrapped: for a 45-character input + // it formats as "invalid urn prefix: %q" over the value's first nine bytes + // (google/uuid uuid.go), so propagating it would echo part of the rejected + // token into the error and from there into logs. The shape of the failure + // carries no diagnostic value the fixed message below does not already + // give, so the value is dropped rather than masked. + if event.OwnerToken != "" { + if _, err := uuid.Parse(event.OwnerToken); err != nil { + return "", ScheduledTaskParams{}, errors.New("invalid owner_token in scheduled event: not a UUID") + } } // Map action to task type if taskType, ok := scheduledEventActions[event.Action]; ok { - return taskType, nil + return taskType, ScheduledTaskParams{OwnerToken: event.OwnerToken}, nil } - return "", fmt.Errorf("unknown scheduled task action: %q", event.Action) + return "", ScheduledTaskParams{}, fmt.Errorf("unknown scheduled task action: %q", event.Action) } diff --git a/internal/server/handler_ladder_test.go b/internal/server/handler_ladder_test.go index 3610365f9..d5d284c5f 100644 --- a/internal/server/handler_ladder_test.go +++ b/internal/server/handler_ladder_test.go @@ -294,7 +294,7 @@ func testBaselineLowWater(lowWater float64) pkgladder.UsageBaseline { // ============================================================ func TestHandleLadderRun_ParseScheduledEvent(t *testing.T) { - taskType, err := ParseScheduledEvent([]byte(`{"action":"ladder_run"}`)) + taskType, _, err := ParseScheduledEvent([]byte(`{"action":"ladder_run"}`)) require.NoError(t, err) assert.Equal(t, TaskLadderRun, taskType) } diff --git a/internal/server/handler_ri_exchange_test.go b/internal/server/handler_ri_exchange_test.go index 7221d4dbd..a11dd93a3 100644 --- a/internal/server/handler_ri_exchange_test.go +++ b/internal/server/handler_ri_exchange_test.go @@ -59,7 +59,7 @@ func TestHandleRIExchangeReshape_ConfigLoadFailure(t *testing.T) { } func TestParseScheduledEvent_RIExchangeReshape(t *testing.T) { - taskType, err := ParseScheduledEvent([]byte(`{"action": "ri_exchange_reshape"}`)) + taskType, _, err := ParseScheduledEvent([]byte(`{"action": "ri_exchange_reshape"}`)) testutil.AssertNoError(t, err) testutil.AssertEqual(t, TaskRIExchangeReshape, taskType) } diff --git a/internal/server/handler_test.go b/internal/server/handler_test.go index 91a144987..176d9abc5 100644 --- a/internal/server/handler_test.go +++ b/internal/server/handler_test.go @@ -3,12 +3,15 @@ package server import ( "context" "errors" + "strings" "testing" "time" + "github.com/LeanerCloud/CUDly/internal/mocks" "github.com/LeanerCloud/CUDly/internal/purchase" "github.com/LeanerCloud/CUDly/internal/scheduler" "github.com/LeanerCloud/CUDly/internal/testutil" + "github.com/stretchr/testify/mock" ) // mockTaskLocker implements TaskLocker for testing. @@ -39,7 +42,7 @@ func TestHandleScheduledTask(t *testing.T) { name: "collect_recommendations success", taskType: TaskCollectRecommendations, setupMocks: func(s *testutil.MockScheduler, p *testutil.MockPurchaseManager) { - s.CollectRecommendationsFunc = func(ctx context.Context) (*scheduler.CollectResult, error) { + s.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{}, nil } }, @@ -49,7 +52,7 @@ func TestHandleScheduledTask(t *testing.T) { name: "collect_recommendations failure", taskType: TaskCollectRecommendations, setupMocks: func(s *testutil.MockScheduler, p *testutil.MockPurchaseManager) { - s.CollectRecommendationsFunc = func(ctx context.Context) (*scheduler.CollectResult, error) { + s.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return nil, errors.New("collection failed") } }, @@ -181,7 +184,7 @@ func TestHandleScheduledTask(t *testing.T) { Purchase: mockPurchase, } - _, err := app.HandleScheduledTask(ctx, tt.taskType) + _, err := app.HandleScheduledTask(ctx, tt.taskType, ScheduledTaskParams{}) if tt.expectError { testutil.AssertError(t, err) @@ -230,7 +233,7 @@ func TestHandleScheduledTaskSkipsWhenDBNil(t *testing.T) { ctx := testutil.TestContext(t) mockScheduler := &testutil.MockScheduler{} - mockScheduler.CollectRecommendationsFunc = func(ctx context.Context) (*scheduler.CollectResult, error) { + mockScheduler.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{Recommendations: 5}, nil } @@ -240,7 +243,7 @@ func TestHandleScheduledTaskSkipsWhenDBNil(t *testing.T) { DB: nil, // No DB — lock path skipped } - result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations) + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) testutil.AssertNoError(t, err) if result == nil { t.Fatal("expected non-nil result") @@ -258,7 +261,7 @@ func TestHandleScheduledTaskAdvisoryLock(t *testing.T) { TaskLocker: locker, } - _, err := app.HandleScheduledTask(ctx, TaskCleanupExpiredRecords) + _, err := app.HandleScheduledTask(ctx, TaskCleanupExpiredRecords, ScheduledTaskParams{}) testutil.AssertNoError(t, err) testutil.AssertEqual(t, 1, locker.lockCalls) testutil.AssertEqual(t, 1, locker.unlockCalls) @@ -274,7 +277,7 @@ func TestHandleScheduledTaskAdvisoryLock(t *testing.T) { TaskLocker: locker, } - result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations) + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) testutil.AssertNoError(t, err) testutil.AssertEqual(t, 1, locker.lockCalls) testutil.AssertEqual(t, 0, locker.unlockCalls) @@ -297,12 +300,113 @@ func TestHandleScheduledTaskAdvisoryLock(t *testing.T) { TaskLocker: locker, } - _, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations) + _, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) testutil.AssertError(t, err) testutil.AssertContains(t, err.Error(), "failed to check task lock") }) } +// TestHandleScheduledTaskReleasesMarkerWhenSkipped pins the abandoned-marker +// leak in the issue #261 compare-and-clear guard: a collect run that WON +// MarkCollectionStarted (so it owns the in-flight marker) but is then skipped +// by the advisory lock never reaches CollectRecommendations, so the deferred +// token-scoped clear never fires. Before this release the marker sat stranded +// for the full 5-minute auto-recovery window, rejecting every refresh the user +// attempted with 409 while nothing backed by that marker was running. +func TestHandleScheduledTaskReleasesMarkerWhenSkipped(t *testing.T) { + t.Run("owner token released when the run is skipped", func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + t.Cleanup(func() { store.AssertExpectations(t) }) + store.On("ClearCollectionStarted", mock.Anything, "tok-owner").Return(nil) + + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{acquired: false}, + } + + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, + ScheduledTaskParams{OwnerToken: "tok-owner"}) + testutil.AssertNoError(t, err) + + m, ok := result.(map[string]string) + if !ok { + t.Fatalf("expected map[string]string, got %T", result) + } + testutil.AssertEqual(t, "skipped", m["status"]) + store.AssertCalled(t, "ClearCollectionStarted", mock.Anything, "tok-owner") + }) + + // A tokenless run (EventBridge cron, the /api/scheduled/ HTTP path, the + // --task CLI) owns no marker, so a skip must not clear anything: that + // would be exactly the cross-run wipe issue #261 closes. The expectation + // is registered (as Maybe) so an unwanted call is still recorded rather + // than falling through the mock's no-expectation default and letting + // AssertNotCalled pass vacuously. + t.Run("tokenless skipped run clears nothing", func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + t.Cleanup(func() { store.AssertExpectations(t) }) + store.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe() + + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{acquired: false}, + } + + _, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) + testutil.AssertNoError(t, err) + store.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything) + }) + + // A lock-check error is returned to the caller, so the Lambda async invoke + // retries the same event with the same owner token and the collect can + // still run. Releasing the marker there would strand the retry. + t.Run("lock error keeps the marker for the retry", func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + t.Cleanup(func() { store.AssertExpectations(t) }) + store.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe() + + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{err: errors.New("db connection lost")}, + } + + _, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, + ScheduledTaskParams{OwnerToken: "tok-owner"}) + testutil.AssertError(t, err) + store.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything) + }) + + // Only collect_recommendations carries an owner token. A stray token on + // another task type must never reach the collection marker. + t.Run("other task types never touch the marker", func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + t.Cleanup(func() { store.AssertExpectations(t) }) + store.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe() + + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{acquired: false}, + } + + _, err := app.HandleScheduledTask(ctx, TaskCleanupExpiredRecords, + ScheduledTaskParams{OwnerToken: "tok-owner"}) + testutil.AssertNoError(t, err) + store.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything) + }) +} + func TestHandleSQSMessage(t *testing.T) { tests := []struct { setupMocks func(*testutil.MockPurchaseManager) @@ -354,18 +458,70 @@ func TestHandleSQSMessage(t *testing.T) { } } +// TestParseScheduledEvent_MalformedTokenNotEchoedInError pins that the +// rejection error for a malformed owner_token never carries any of the +// rejected value. The 45-character input below is the one shape where +// uuid.Parse formats its error as "invalid urn prefix: %q" over the value's +// first nine bytes (google/uuid uuid.go), so wrapping that error with %w +// would echo "OWNERTOK-" into the error string and from there into the +// Lambda logs. Asserting on a shorter malformed token would pass with the +// bug present, because uuid.Parse reports those as a bare length/format +// error that happens to contain nothing sensitive. +func TestParseScheduledEvent_MalformedTokenNotEchoedInError(t *testing.T) { + // 45 characters, so uuid.Parse takes its urn-prefix branch. + const leakyToken = "OWNERTOK-6b1f2c34-5d6e-4a7b-8c9d-0e1f2a3b4c5d" + testutil.AssertEqual(t, 45, len(leakyToken)) + + _, _, err := ParseScheduledEvent([]byte( + `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "` + leakyToken + `"}`)) + + testutil.AssertError(t, err) + if strings.Contains(err.Error(), leakyToken[:9]) { + t.Fatalf("rejection error must not echo the token value, got: %s", err.Error()) + } +} + func TestParseScheduledEvent(t *testing.T) { tests := []struct { - name string - rawEvent string - expectedTask ScheduledTaskType - expectError bool + name string + rawEvent string + expectedTask ScheduledTaskType + expectedToken string + expectError bool }{ { name: "collect_recommendations event", rawEvent: `{"action": "collect_recommendations"}`, expectedTask: TaskCollectRecommendations, }, + { + // Async self-invoke payload carries the owner token so the + // scheduler can scope ClearCollectionStarted to this run + // (issue #261 compare-and-clear guard). + name: "collect_recommendations event with owner_token", + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "` + testOwnerToken + `"}`, + expectedTask: TaskCollectRecommendations, + expectedToken: testOwnerToken, + }, + { + // A non-empty owner_token is validated at the boundary: the only + // legitimate producer is asyncInvokeSelf, which always sends a + // uuid.New(), so a non-UUID value means a corrupt payload. It + // could never match a marker owner, and letting it through would + // strand that marker for the full 5-minute recovery window with + // only a buried error log to show for it. + name: "collect_recommendations event with malformed owner_token", + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "tok-1"}`, + expectError: true, + }, + { + // An absent owner_token stays legitimate: cron, the + // /api/scheduled/ HTTP path and the --task CLI never win + // MarkCollectionStarted and own no marker to clear. + name: "collect_recommendations event with empty owner_token", + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": ""}`, + expectedTask: TaskCollectRecommendations, + }, { name: "process_scheduled_purchases event", rawEvent: `{"action": "process_scheduled_purchases"}`, @@ -420,12 +576,13 @@ func TestParseScheduledEvent(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - taskType, err := ParseScheduledEvent([]byte(tt.rawEvent)) + taskType, params, err := ParseScheduledEvent([]byte(tt.rawEvent)) if tt.expectError { testutil.AssertError(t, err) } else { testutil.AssertNoError(t, err) testutil.AssertEqual(t, tt.expectedTask, taskType) + testutil.AssertEqual(t, tt.expectedToken, params.OwnerToken) } }) } diff --git a/internal/server/http.go b/internal/server/http.go index 119c5a21d..957014789 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -243,8 +243,10 @@ func (app *Application) handleScheduledHTTP(w http.ResponseWriter, r *http.Reque } taskType := ScheduledTaskType(taskTypeStr) - // Execute scheduled task - result, err := app.HandleScheduledTask(ctx, taskType) + // Execute scheduled task. This HTTP-triggered path carries no owner + // token; TaskCollectRecommendations runs as if cron-triggered and skips + // the collection-marker clear (see ScheduledTaskParams). + result, err := app.HandleScheduledTask(ctx, taskType, ScheduledTaskParams{}) if err != nil { log.Printf("Scheduled task %q error: %v", taskTypeStr, err) // #nosec G706 -- taskTypeStr printed with %q which escapes special chars; validated to contain no '/' before this point http.Error(w, "Internal server error", http.StatusInternalServerError) diff --git a/internal/server/http_test.go b/internal/server/http_test.go index 401222fa4..deb654939 100644 --- a/internal/server/http_test.go +++ b/internal/server/http_test.go @@ -245,7 +245,7 @@ func TestHandleScheduledHTTP(t *testing.T) { path: "/api/scheduled/collect_recommendations", setupApp: func(_ *testing.T, app *Application) { app.Scheduler = &testutil.MockScheduler{ - CollectRecommendationsFunc: func(ctx context.Context) (*scheduler.CollectResult, error) { + CollectRecommendationsFunc: func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{ Recommendations: 10, TotalSavings: 1000.0, @@ -289,7 +289,7 @@ func TestHandleScheduledHTTP(t *testing.T) { setupApp: func(t *testing.T, app *Application) { app.scheduledAuth = newBearerValidator(t, "my-secret") app.Scheduler = &testutil.MockScheduler{ - CollectRecommendationsFunc: func(ctx context.Context) (*scheduler.CollectResult, error) { + CollectRecommendationsFunc: func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{}, nil }, } diff --git a/internal/server/integration_test.go b/internal/server/integration_test.go index 02e32eba2..15a78b481 100644 --- a/internal/server/integration_test.go +++ b/internal/server/integration_test.go @@ -83,7 +83,7 @@ func TestScheduledTaskIntegration(t *testing.T) { // Create application with mock dependencies mockScheduler := &testutil.MockScheduler{ - CollectRecommendationsFunc: func(ctx context.Context) (*scheduler.CollectResult, error) { + CollectRecommendationsFunc: func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { // Simulate actual work time.Sleep(100 * time.Millisecond) return &scheduler.CollectResult{}, nil @@ -95,7 +95,7 @@ func TestScheduledTaskIntegration(t *testing.T) { } // Test collect_recommendations task - result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations) + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) testutil.AssertNoError(t, err) testutil.AssertTrue(t, result != nil, "Result should not be nil") diff --git a/internal/server/interfaces.go b/internal/server/interfaces.go index 4a797e6af..5b04c1290 100644 --- a/internal/server/interfaces.go +++ b/internal/server/interfaces.go @@ -11,7 +11,7 @@ import ( // SchedulerInterface defines the methods required for the scheduler component. type SchedulerInterface interface { - CollectRecommendations(ctx context.Context) (*scheduler.CollectResult, error) + CollectRecommendations(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) ListRecommendations(ctx context.Context, filter config.RecommendationFilter) ([]config.RecommendationRecord, error) // GetRecommendationByID fetches a single rec by application-level id, // bypassing account-override filtering. hiddenBy is non-nil when the rec diff --git a/internal/server/lambda.go b/internal/server/lambda.go index a620a17fe..348d349b1 100644 --- a/internal/server/lambda.go +++ b/internal/server/lambda.go @@ -249,10 +249,10 @@ func (app *Application) handleLambdaSQSEvent(ctx context.Context, rawEvent json. // handleLambdaScheduledEvent processes scheduled/cron events. func (app *Application) handleLambdaScheduledEvent(ctx context.Context, rawEvent json.RawMessage) (any, error) { - taskType, err := ParseScheduledEvent(rawEvent) + taskType, params, err := ParseScheduledEvent(rawEvent) if err != nil { return nil, fmt.Errorf("failed to parse scheduled event: %w", err) } - return app.HandleScheduledTask(ctx, taskType) + return app.HandleScheduledTask(ctx, taskType, params) } diff --git a/internal/server/lambda_test.go b/internal/server/lambda_test.go index 1a81cc41d..d0dc11f50 100644 --- a/internal/server/lambda_test.go +++ b/internal/server/lambda_test.go @@ -212,18 +212,24 @@ func TestHandleLambdaSQSEvent(t *testing.T) { } } +// testOwnerToken is a fixed, well-formed collection owner token. It has to be +// a real UUID because ParseScheduledEvent rejects malformed non-empty tokens +// at the boundary rather than letting them reach the UUID-typed column. +const testOwnerToken = "6b1f2c34-5d6e-4a7b-8c9d-0e1f2a3b4c5d" + func TestHandleLambdaScheduledEvent(t *testing.T) { tests := []struct { - setupMocks func(*testutil.MockScheduler) - name string - rawEvent string - expectError bool + setupMocks func(*testutil.MockScheduler) + name string + rawEvent string + expectedToken string + expectError bool }{ { name: "collect_recommendations event", rawEvent: `{"action": "collect_recommendations"}`, setupMocks: func(s *testutil.MockScheduler) { - s.CollectRecommendationsFunc = func(ctx context.Context) (*scheduler.CollectResult, error) { + s.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{ Recommendations: 10, TotalSavings: 500.0, @@ -236,7 +242,7 @@ func TestHandleLambdaScheduledEvent(t *testing.T) { name: "EventBridge format", rawEvent: `{"source": "aws.events", "action": "collect_recommendations"}`, setupMocks: func(s *testutil.MockScheduler) { - s.CollectRecommendationsFunc = func(ctx context.Context) (*scheduler.CollectResult, error) { + s.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{ Recommendations: 5, TotalSavings: 250.0, @@ -245,6 +251,28 @@ func TestHandleLambdaScheduledEvent(t *testing.T) { }, expectError: false, }, + { + // Issue #261: the async self-invoke payload's owner_token must + // survive parsing and reach CollectRecommendations, which is what + // scopes the deferred clear to this run. Asserting the token the + // scheduler actually received (rather than only that the event + // parses) means dropping it anywhere between ParseScheduledEvent + // and the scheduler fails the test instead of silently stranding + // the marker for the full 5-minute recovery window. + name: "async self-invoke carries owner_token through to the scheduler", + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "` + testOwnerToken + `"}`, + expectedToken: testOwnerToken, + expectError: false, + }, + { + // A non-empty owner_token that is not a UUID can only come from a + // corrupt payload (asyncInvokeSelf always sends a uuid.New()), and + // could never match a marker owner, so it is rejected at the + // boundary rather than reaching the UUID-typed persistence layer. + name: "malformed owner_token is rejected at the boundary", + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "not-a-uuid"}`, + expectError: true, + }, } for _, tt := range tests { @@ -252,7 +280,19 @@ func TestHandleLambdaScheduledEvent(t *testing.T) { ctx := testutil.TestContext(t) mockScheduler := &testutil.MockScheduler{} - tt.setupMocks(mockScheduler) + // Capture the token the scheduler was handed and assert on it from + // the subtest goroutine rather than inside the mock callback, so a + // failed assertion never calls FailNow off the owning goroutine. + gotToken := "" + collected := false + mockScheduler.CollectRecommendationsFunc = func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { + gotToken = ownerToken + collected = true + return &scheduler.CollectResult{}, nil + } + if tt.setupMocks != nil { + tt.setupMocks(mockScheduler) + } app := &Application{ Scheduler: mockScheduler, @@ -262,8 +302,14 @@ func TestHandleLambdaScheduledEvent(t *testing.T) { if tt.expectError { testutil.AssertError(t, err) - } else { - testutil.AssertNoError(t, err) + return + } + testutil.AssertNoError(t, err) + if tt.expectedToken != "" { + if !collected { + t.Fatal("expected CollectRecommendations to be called") + } + testutil.AssertEqual(t, tt.expectedToken, gotToken) } }) } @@ -328,7 +374,7 @@ func TestHandleLambdaEvent(t *testing.T) { rawEvent: `{"action": "collect_recommendations"}`, setupApp: func(app *Application) { app.Scheduler = &testutil.MockScheduler{ - CollectRecommendationsFunc: func(ctx context.Context) (*scheduler.CollectResult, error) { + CollectRecommendationsFunc: func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { return &scheduler.CollectResult{}, nil }, } diff --git a/internal/server/test_helpers_test.go b/internal/server/test_helpers_test.go index 64bbe7033..c38871acb 100644 --- a/internal/server/test_helpers_test.go +++ b/internal/server/test_helpers_test.go @@ -274,10 +274,10 @@ func (m *mockConfigStoreForHealth) GetRecommendationsFreshness(_ context.Context func (m *mockConfigStoreForHealth) SetRecommendationsCollectionError(_ context.Context, _ string) error { return nil } -func (m *mockConfigStoreForHealth) MarkCollectionStarted(_ context.Context) (bool, error) { - return true, nil +func (m *mockConfigStoreForHealth) MarkCollectionStarted(_ context.Context) (string, bool, error) { + return "mock-owner-token", true, nil } -func (m *mockConfigStoreForHealth) ClearCollectionStarted(_ context.Context) error { +func (m *mockConfigStoreForHealth) ClearCollectionStarted(_ context.Context, _ string) error { return nil } func (m *mockConfigStoreForHealth) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { diff --git a/internal/testutil/mocks.go b/internal/testutil/mocks.go index 472f1706e..14de628e2 100644 --- a/internal/testutil/mocks.go +++ b/internal/testutil/mocks.go @@ -11,14 +11,14 @@ import ( // MockScheduler is a mock implementation of server.SchedulerInterface. type MockScheduler struct { - CollectRecommendationsFunc func(ctx context.Context) (*scheduler.CollectResult, error) + CollectRecommendationsFunc func(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) ListRecommendationsFunc func(ctx context.Context, filter config.RecommendationFilter) ([]config.RecommendationRecord, error) GetRecommendationByIDFunc func(ctx context.Context, id string) (*config.RecommendationRecord, []string, error) } -func (m *MockScheduler) CollectRecommendations(ctx context.Context) (*scheduler.CollectResult, error) { +func (m *MockScheduler) CollectRecommendations(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { if m.CollectRecommendationsFunc != nil { - return m.CollectRecommendationsFunc(ctx) + return m.CollectRecommendationsFunc(ctx, ownerToken) } return &scheduler.CollectResult{}, nil }