From f388ba3d37e88e21a290e6405434f577b1742ae7 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 12:44:44 +0200 Subject: [PATCH 01/11] feat(db): add last_collection_owner_id to recommendations_state Adds a nullable UUID column that pairs with last_collection_started_at as a compare-and-clear guard for issue #261: the scheduler's unconditional defer-clear wipes another run's in-flight marker when a cron run overlaps a user-triggered collection. This column lets the clear be scoped to the caller that actually owns the marker. --- ...93_recommendations_state_owner_id.down.sql | 2 + ...0093_recommendations_state_owner_id.up.sql | 17 +++++ ...093_recommendations_state_owner_id_test.go | 69 +++++++++++++++++++ 3 files changed, 88 insertions(+) create mode 100644 internal/database/postgres/migrations/000093_recommendations_state_owner_id.down.sql create mode 100644 internal/database/postgres/migrations/000093_recommendations_state_owner_id.up.sql create mode 100644 internal/database/postgres/migrations/000093_recommendations_state_owner_id_test.go 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") +} From 19e2371531c0e589c94b5b7bade97c0a36921a77 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 13:31:44 +0200 Subject: [PATCH 02/11] feat(config,scheduler,api): owner-token compare-and-clear for collection marker (#261) MarkCollectionStarted now stamps a fresh owner token alongside last_collection_started_at, and ClearCollectionStarted only clears when the caller's token still matches last_collection_owner_id. This closes the race where a cron run's unconditional clear could wipe a concurrent user-triggered run's in-flight marker: previously, any caller reaching CollectRecommendations cleared the marker on exit regardless of whether it owned it. The token is generated by MarkCollectionStarted and threaded through the whole async chain: the refresh handler passes it into the Lambda self-invoke payload (owner_token), ParseScheduledEvent extracts it into ScheduledTaskParams, and HandleScheduledTask carries it to CollectRecommendations. Callers that never won a marker (cron, cold-start, background refresh) pass an empty token, and the deferred clear becomes an explicit, logged skip rather than an unconditional wipe. An empty token passed to ClearCollectionStarted directly is a boundary error; a stale/mismatched token is a documented silent no-op. Adds a store-level regression test replaying the exact overlap from the issue (run A wins, goes stale, run B wins with a new token, run A's late clear must not affect run B's marker), plus scheduler and handler unit tests covering the empty-token skip, the owner-token clear (with a deadline-bearing context), the async-invoke rollback path, and the payload's owner_token field. --- cmd/server/main.go | 2 +- internal/analytics/collector_test.go | 6 +- .../api/handler_recommendations_refresh.go | 30 +++++--- internal/api/handler_test.go | 73 +++++++++++++++++- internal/api/mocks_test.go | 4 +- internal/api/types.go | 2 +- internal/config/interfaces.go | 23 ++++-- .../config/store_postgres_recommendations.go | 59 ++++++++++----- .../store_postgres_recommendations_test.go | 66 ++++++++++++++++ internal/mocks/stores.go | 17 +++-- internal/scheduler/scheduler.go | 47 +++++++----- .../scheduler/scheduler_overrides_test.go | 6 +- .../scheduler/scheduler_suppressions_test.go | 6 +- internal/scheduler/scheduler_test.go | 75 ++++++++++++++++--- internal/server/app_test.go | 8 +- internal/server/handler.go | 72 ++++++++++++------ internal/server/handler_ladder_test.go | 2 +- internal/server/handler_ri_exchange_test.go | 2 +- internal/server/handler_test.go | 37 +++++---- internal/server/http.go | 6 +- internal/server/http_test.go | 4 +- internal/server/integration_test.go | 4 +- internal/server/interfaces.go | 2 +- internal/server/lambda.go | 4 +- internal/server/lambda_test.go | 6 +- internal/server/test_helpers_test.go | 6 +- internal/testutil/mocks.go | 6 +- 27 files changed, 430 insertions(+), 145 deletions(-) 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..b7482f627 100644 --- a/internal/api/handler_recommendations_refresh.go +++ b/internal/api/handler_recommendations_refresh.go @@ -66,7 +66,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 +77,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,11 +106,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 { + if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN, token); invokeErr != nil { + if clearErr := h.config.ClearCollectionStarted(ctx, token); clearErr != nil { logging.Warnf("runMarkedCollection: failed to clear collection started marker: %v", clearErr) } return nil, fmt.Errorf("failed to trigger async collection: %w", invokeErr) @@ -119,8 +127,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) @@ -135,7 +144,7 @@ func (h *Handler) runMarkedCollection(ctx context.Context) (*config.Recommendati // 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 +168,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..1a9942d04 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,73 @@ 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_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) + + 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..5705c2069 100644 --- a/internal/config/store_postgres_recommendations.go +++ b/internal/config/store_postgres_recommendations.go @@ -448,39 +448,62 @@ 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) } - return tag.RowsAffected() == 1, nil + if tag.RowsAffected() != 1 { + return "", false, 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 { +// 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") + } if _, 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); err != nil { return fmt.Errorf("failed to clear collection started: %w", err) } return nil diff --git a/internal/config/store_postgres_recommendations_test.go b/internal/config/store_postgres_recommendations_test.go index dfd3eb8ac..cb1a5dca2 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,69 @@ 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 by replaying the exact race the issue +// describes: a cron run (run B) overlaps a user-triggered async run +// (run A). 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. +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 scheduled cron run (run B) is + // allowed to start. + _, 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) +} + func float64Ptr(f float64) *float64 { return &f } 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..8cc9aea63 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -141,26 +141,31 @@ 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. +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 +364,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 +942,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 +1215,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..10ccc719f 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,70 @@ 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) + + 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 +345,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 +413,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 +448,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 +474,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 +613,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 +652,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 +1710,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..4c64f9fa7 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -81,7 +81,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 { @@ -102,32 +102,39 @@ func (app *Application) HandleScheduledTask(ctx context.Context, taskType Schedu defer locker.ReleaseAdvisoryLock(ctx, lockID) } - return app.dispatchTask(ctx, taskType) + return app.dispatchTask(ctx, taskType, params) } // 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 +156,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 +342,33 @@ 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"` } -// ParseScheduledEvent parses a scheduled event and returns the task type. -func ParseScheduledEvent(rawEvent json.RawMessage) (ScheduledTaskType, error) { +// 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 +// 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) } // 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..46f62cc1d 100644 --- a/internal/server/handler_test.go +++ b/internal/server/handler_test.go @@ -39,7 +39,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 +49,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 +181,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 +230,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 +240,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 +258,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 +274,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,7 +297,7 @@ 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") }) @@ -356,16 +356,26 @@ func TestHandleSQSMessage(t *testing.T) { 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": "tok-1"}`, + expectedTask: TaskCollectRecommendations, + expectedToken: "tok-1", + }, { name: "process_scheduled_purchases event", rawEvent: `{"action": "process_scheduled_purchases"}`, @@ -420,12 +430,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..b918edd06 100644 --- a/internal/server/lambda_test.go +++ b/internal/server/lambda_test.go @@ -223,7 +223,7 @@ func TestHandleLambdaScheduledEvent(t *testing.T) { 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 +236,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, @@ -328,7 +328,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 } From 8feac225434ce265d9b6e30f722e0b54197154b3 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 14:00:13 +0200 Subject: [PATCH 03/11] fix(config): stop SetRecommendationsCollectionError from clearing owner marker (#261) SetRecommendationsCollectionError unconditionally cleared last_collection_started_at with no owner check. It is called mid-run from persistCollection on every CollectRecommendations invocation that hits a provider error, including tokenless cron/cold-start/background runs, so a tokenless run's routine provider error could wipe a concurrent owner run's in-flight marker. This reopened the exact race the compare-and-clear guard in the previous commit exists to close. The only place that may clear last_collection_started_at now is the deferred, token-guarded clearCollectionStartedBestEffort in CollectRecommendations, which already runs on both the success and failure exit paths. SetRecommendationsCollectionError now only writes last_collection_error. Adds a regression test replaying the scenario: MarkCollectionStarted stamps an owner token, a tokenless caller invokes SetRecommendationsCollectionError, and the owner's marker must survive. --- .../config/store_postgres_recommendations.go | 18 +++++--- .../store_postgres_recommendations_test.go | 41 +++++++++++++++++++ internal/scheduler/scheduler.go | 6 ++- 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/internal/config/store_postgres_recommendations.go b/internal/config/store_postgres_recommendations.go index 5705c2069..14ee3755f 100644 --- a/internal/config/store_postgres_recommendations.go +++ b/internal/config/store_postgres_recommendations.go @@ -427,16 +427,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) diff --git a/internal/config/store_postgres_recommendations_test.go b/internal/config/store_postgres_recommendations_test.go index cb1a5dca2..ddd408ef2 100644 --- a/internal/config/store_postgres_recommendations_test.go +++ b/internal/config/store_postgres_recommendations_test.go @@ -459,4 +459,45 @@ func TestPostgresStore_ClearCollectionStarted_CompareAndClear(t *testing.T) { 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/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 8cc9aea63..abb427cb9 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -152,7 +152,11 @@ func cacheTTLFromEnv() time.Duration { // 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. +// 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...") From d5e13156663823abaee1457bb56abe352d6124f0 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 14:27:10 +0200 Subject: [PATCH 04/11] fix(server): release owner marker when a collect run is lock-skipped A collect run that won MarkCollectionStarted (so it owns the in-flight marker) but is then skipped by the advisory lock never reaches CollectRecommendations, whose deferred token-scoped clear would normally release the marker. It therefore sat stranded for the full 5-minute auto-recovery window, and every refresh the user attempted in that window was rejected with 409 "collection already in progress" while no run backed by that marker existed. Before the 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 is returned to the caller, which lets the Lambda async invoke retry the same event with the same owner token and still run the collect, so the marker must survive there. The release is scoped by the owner token, so it can only ever clear the marker this run owns, never a concurrent run's, and tokenless callers (cron, the /api/scheduled/ HTTP path, the --task CLI) own nothing and are skipped. Regression test asserts the clear fires with the owning token on a skip, and that tokenless runs, lock errors, and other task types never touch the marker. Verified failing before the fix and passing after. Refs #261 --- internal/server/handler.go | 49 +++++++++++++++ internal/server/handler_test.go | 103 ++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/internal/server/handler.go b/internal/server/handler.go index 4c64f9fa7..fec4df97c 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -6,11 +6,17 @@ import ( "fmt" "hash/fnv" "log" + "time" "github.com/LeanerCloud/CUDly/internal/purchase" "github.com/LeanerCloud/CUDly/internal/scheduler" ) +// 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) @@ -97,6 +103,7 @@ 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) @@ -105,6 +112,48 @@ func (app *Application) HandleScheduledTask(ctx context.Context, taskType Schedu 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. +// +// 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, params ScheduledTaskParams) (any, error) { // Map-based dispatch (rather than a switch) keeps this function under the diff --git a/internal/server/handler_test.go b/internal/server/handler_test.go index 46f62cc1d..0e77c08ec 100644 --- a/internal/server/handler_test.go +++ b/internal/server/handler_test.go @@ -6,9 +6,11 @@ import ( "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. @@ -303,6 +305,107 @@ func TestHandleScheduledTaskAdvisoryLock(t *testing.T) { }) } +// 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) From 3943e9e87fc6082465fdbf843f3411de373cd8e8 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 14:28:04 +0200 Subject: [PATCH 05/11] docs(config): correct run-B framing in compare-and-clear test The store-level compare-and-clear test described run B as "a cron run", but both runs in it go through MarkCollectionStarted, which only the POST /api/recommendations/refresh handler calls. Cron reaches CollectRecommendations directly and never marks, so run B actually models a second user-triggered refresh taking the marker over after run A's 5-minute window lapses. Points the cron half of the race at the scheduler test that does pin it (TestScheduler_CollectRecommendations_EmptyTokenSkipsClear), so both halves of the guard are traceable from either side. Comments only, no behavior change. Refs #261 --- .../store_postgres_recommendations_test.go | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/internal/config/store_postgres_recommendations_test.go b/internal/config/store_postgres_recommendations_test.go index ddd408ef2..445137eda 100644 --- a/internal/config/store_postgres_recommendations_test.go +++ b/internal/config/store_postgres_recommendations_test.go @@ -395,11 +395,19 @@ func TestPostgresStore_UpsertRecommendations_AmbientAndRegisteredCoexist(t *test } // TestPostgresStore_ClearCollectionStarted_CompareAndClear pins the issue -// #261 compare-and-clear guard by replaying the exact race the issue -// describes: a cron run (run B) overlaps a user-triggered async run -// (run A). Pre-fix, ClearCollectionStarted took no token and cleared +// #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) @@ -423,8 +431,8 @@ func TestPostgresStore_ClearCollectionStarted_CompareAndClear(t *testing.T) { assert.Empty(t, tokenNone) // Force run A's marker stale (simulates its Lambda crashing or running - // long past the 5-minute window), so the scheduled cron run (run B) is - // allowed to start. + // 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' From f53bf3736bba2ab3443dcff00f154db5e3acd0b3 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 14:37:03 +0200 Subject: [PATCH 06/11] fix(server): reject malformed owner tokens at the event boundary owner_token arrives as raw event input and is forwarded to a UUID-typed column, so validate it in ParseScheduledEvent instead of letting a corrupt value reach the persistence layer. Empty stays legitimate: EventBridge cron, the /api/scheduled/ HTTP path and the --task CLI never win MarkCollectionStarted and own no marker. A non-empty value can only have come from asyncInvokeSelf, which always sends a uuid.New(), so anything else means the payload is corrupt. Failing loud is deliberate: a token that cannot match any owner would strand the marker it belongs to for the full 5-minute recovery window regardless, so a buried error log on the eventual clear is strictly worse than refusing the event. uuid.Parse's error text describes the shape only and never echoes the value, so wrapping it does not leak the token into logs. Adds end-to-end coverage that a well-formed token survives parsing and reaches CollectRecommendations, which is what scopes the deferred clear to that run: dropping it anywhere in between now fails the test rather than silently stranding the marker. The captured token is asserted from the subtest goroutine rather than inside the mock callback so a failed assertion never calls FailNow off the owning goroutine. Addresses CodeRabbit review findings on PR #1516. Refs #261 --- internal/server/handler.go | 19 +++++++++++ internal/server/handler_test.go | 23 +++++++++++-- internal/server/lambda_test.go | 60 +++++++++++++++++++++++++++++---- 3 files changed, 93 insertions(+), 9 deletions(-) diff --git a/internal/server/handler.go b/internal/server/handler.go index fec4df97c..7b10f8a5b 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -10,6 +10,7 @@ import ( "github.com/LeanerCloud/CUDly/internal/purchase" "github.com/LeanerCloud/CUDly/internal/scheduler" + "github.com/google/uuid" ) // skippedCollectionMarkerClearTimeout bounds the detached best-effort clear @@ -415,6 +416,24 @@ func ParseScheduledEvent(rawEvent json.RawMessage) (ScheduledTaskType, Scheduled 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 text describes the shape only and never echoes the + // token value, so wrapping it does not leak the token into logs. + if event.OwnerToken != "" { + if _, err := uuid.Parse(event.OwnerToken); err != nil { + return "", ScheduledTaskParams{}, fmt.Errorf("invalid owner_token in scheduled event: %w", err) + } + } + // Map action to task type if taskType, ok := scheduledEventActions[event.Action]; ok { return taskType, ScheduledTaskParams{OwnerToken: event.OwnerToken}, nil diff --git a/internal/server/handler_test.go b/internal/server/handler_test.go index 0e77c08ec..e7320a935 100644 --- a/internal/server/handler_test.go +++ b/internal/server/handler_test.go @@ -475,9 +475,28 @@ func TestParseScheduledEvent(t *testing.T) { // 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": "tok-1"}`, + rawEvent: `{"source": "aws.events", "action": "collect_recommendations", "owner_token": "` + testOwnerToken + `"}`, expectedTask: TaskCollectRecommendations, - expectedToken: "tok-1", + 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", diff --git a/internal/server/lambda_test.go b/internal/server/lambda_test.go index b918edd06..d0dc11f50 100644 --- a/internal/server/lambda_test.go +++ b/internal/server/lambda_test.go @@ -212,12 +212,18 @@ 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", @@ -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) } }) } From b16d59fb3583e2a8a64539008b89b5edf5e77fb7 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 18:26:21 +0200 Subject: [PATCH 07/11] fix(api): detach async-invoke rollback clear from request context When asyncInvokeSelf fails, runMarkedCollection rolls back this caller's own in-flight collection marker. It reused the request context to do so, but the dominant cause of an invoke failure is that context itself expiring or being canceled (API Gateway deadline, client disconnect, a slow SDK credential refresh), so the rollback failed in exactly the case that produced it. The marker then sat stranded for the full 5-minute auto-recovery window in MarkCollectionStarted, rejecting every refresh the user attempted with 409 while no collection was running. Extract rollbackMarkedCollection, which issues the token-scoped clear on a detached 5-second context. This is the same pattern the scheduler's deferred clear and Application.releaseSkippedCollectionMarker already use for the same reason; this sibling path was the one that had not adopted it. The regression test drives a request context that dies mid-invoke and pins the clear context to one that is deadline-bounded and still live. The matcher sits on the mock expectation rather than a trailing AssertCalled, since testify's AssertCalled compares arguments by value and does not honor MatchedBy: pre-fix the clear still happened, just with the canceled context, so asserting only that Clear was called would pass with the bug present. --- .../api/handler_recommendations_refresh.go | 34 ++++++++++++++-- internal/api/handler_test.go | 40 +++++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/internal/api/handler_recommendations_refresh.go b/internal/api/handler_recommendations_refresh.go index b7482f627..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. @@ -115,9 +120,7 @@ func (h *Handler) runMarkedCollection(ctx context.Context, token string) (*confi schedulerARN := os.Getenv("SCHEDULER_LAMBDA_ARN") if schedulerARN != "" { if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN, token); invokeErr != nil { - if clearErr := h.config.ClearCollectionStarted(ctx, token); clearErr != nil { - logging.Warnf("runMarkedCollection: failed to clear collection started marker: %v", clearErr) - } + h.rollbackMarkedCollection(token) return nil, fmt.Errorf("failed to trigger async collection: %w", invokeErr) } freshness, err := h.config.GetRecommendationsFreshness(ctx) @@ -139,6 +142,31 @@ func (h *Handler) runMarkedCollection(ctx context.Context, token string) (*confi 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; diff --git a/internal/api/handler_test.go b/internal/api/handler_test.go index 1a9942d04..e7c2300b3 100644 --- a/internal/api/handler_test.go +++ b/internal/api/handler_test.go @@ -574,6 +574,46 @@ func TestRunMarkedCollection_AsyncInvokeFailure_RollsBackOwnerToken(t *testing.T 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; From be1bc975b0c960ddaa7c3d8338d93b0199d7b97c Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 18:27:04 +0200 Subject: [PATCH 08/11] fix(server): stop echoing rejected owner-token bytes in parse error ParseScheduledEvent wrapped uuid.Parse's error with %w, and the comment asserted that error "describes the shape only and never echoes the token value". That is false for one input shape: at 45 characters uuid.Parse takes its urn-prefix branch and formats the error as `invalid urn prefix: %q` over the value's first nine bytes, which then lands in the Lambda logs. No token a legitimate producer issues can reach that branch (asyncInvokeSelf always sends a 36-character uuid.New()), so this is defense in depth rather than a live leak, but the repo forbids logging token material and the comment claimed a guarantee the code did not provide. Return a fixed message instead of wrapping. The parse failure's shape carries no diagnostic value the fixed message does not already give, so the value is dropped rather than masked. The regression test uses the 45-character input specifically: a shorter malformed token passes with the bug present, because uuid.Parse reports those as a bare length/format error containing nothing sensitive. --- internal/server/handler.go | 12 +++++++++--- internal/server/handler_test.go | 24 ++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/internal/server/handler.go b/internal/server/handler.go index 7b10f8a5b..7b4716b1f 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -3,6 +3,7 @@ package server import ( "context" "encoding/json" + "errors" "fmt" "hash/fnv" "log" @@ -426,11 +427,16 @@ func ParseScheduledEvent(rawEvent json.RawMessage) (ScheduledTaskType, Scheduled // 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 text describes the shape only and never echoes the - // token value, so wrapping it does not leak the token into logs. + // + // 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{}, fmt.Errorf("invalid owner_token in scheduled event: %w", err) + return "", ScheduledTaskParams{}, errors.New("invalid owner_token in scheduled event: not a UUID") } } diff --git a/internal/server/handler_test.go b/internal/server/handler_test.go index e7320a935..176d9abc5 100644 --- a/internal/server/handler_test.go +++ b/internal/server/handler_test.go @@ -3,6 +3,7 @@ package server import ( "context" "errors" + "strings" "testing" "time" @@ -457,6 +458,29 @@ 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 From 581414bad98341d758a5929d822c9d67a8359b12 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 18:27:53 +0200 Subject: [PATCH 09/11] test(scheduler,api): make marker AssertNotCalled guards non-vacuous Both tests pinning "this path must not clear the collection marker" passed with the bug present. MockConfigStore short-circuits any method with no registered expectation before reaching mock.Called, so the call never lands in m.Calls, and testify's AssertNotCalled walks m.Calls. Demonstrated by deleting the empty-token guard in clearCollectionStartedBestEffort: TestScheduler_CollectRecommendations_ EmptyTokenSkipsClear stayed green, leaving the central invariant of the compare-and-clear guard unenforced. The same held for the happy-path assertion in TestRunMarkedCollection_AsyncInvokeSuccess_ PayloadCarriesOwnerToken. Register the expectation as .Maybe() in both, so an unwanted call is recorded and the assertion can actually fail. This is the treatment the sibling cases in internal/server/handler_test.go already had; these two were missed. Verified both now fail when the guarded behavior is removed. --- internal/api/handler_test.go | 5 +++++ internal/scheduler/scheduler_test.go | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/internal/api/handler_test.go b/internal/api/handler_test.go index e7c2300b3..52ef69675 100644 --- a/internal/api/handler_test.go +++ b/internal/api/handler_test.go @@ -627,6 +627,11 @@ func TestRunMarkedCollection_AsyncInvokeSuccess_PayloadCarriesOwnerToken(t *test 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{ diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 10ccc719f..6ee6249f3 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -277,6 +277,13 @@ func TestScheduler_CollectRecommendations_EmptyTokenSkipsClear(t *testing.T) { 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, From c47179567935d257b1b7a095dad896a3e6956e60 Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 18:28:39 +0200 Subject: [PATCH 10/11] docs(server): record duplicate-delivery gap in skipped-marker release releaseSkippedCollectionMarker scopes its clear by owner token, but the token identifies the MARKER, not the invocation. It therefore cannot distinguish "the advisory 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: that sibling carries the same token and is still collecting). Lambda async invocation is at-least-once, so the second case is reachable. There the marker is cleared 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. Bounded and self-healing, and it touches no purchase or money path. Releasing remains strictly better than not releasing, since cron overlap is routine while duplicate delivery is rare. Closing the gap properly needs an invocation-scoped identity distinct from the marker token (for example, the lock winner re-stamping last_collection_owner_id with a fresh token), which is a design change kept out of this PR. Documented here so the trade-off is explicit rather than an undiscovered trap. --- internal/server/handler.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/internal/server/handler.go b/internal/server/handler.go index 7b4716b1f..a6a74499e 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -133,6 +133,23 @@ func (app *Application) HandleScheduledTask(ctx context.Context, taskType Schedu // 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 From 036cb13662ca218da6914222636f8c9e57fd4cbd Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Mon, 27 Jul 2026 18:59:34 +0200 Subject: [PATCH 11/11] feat(config): log when compare-and-clear declines to clear ClearCollectionStarted discarded RowsAffected, so nothing in production distinguished "cleared the marker" from "declined to clear someone else's marker". That no-op branch IS the safety mechanism issue #261 adds, and a guard that never fires looks identical to a guard that is broken. Capture the command tag and emit a debug-level line when zero rows matched. Behavior is unchanged. Addresses a CodeRabbit nitpick, with one deviation: the suggested diff logged the token value (`token %q`). The token is the capability that controls the marker, and this PR already stopped rejected token bytes from reaching the logs via ParseScheduledEvent, so the message names the condition without echoing the token. --- internal/config/store_postgres_recommendations.go | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/internal/config/store_postgres_recommendations.go b/internal/config/store_postgres_recommendations.go index 14ee3755f..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" ) @@ -505,14 +506,24 @@ func (s *PostgresStore) ClearCollectionStarted(ctx context.Context, token string if token == "" { return fmt.Errorf("owner token must not be empty") } - if _, err := s.db.Exec(ctx, ` + tag, err := s.db.Exec(ctx, ` UPDATE recommendations_state SET last_collection_started_at = NULL, last_collection_owner_id = NULL WHERE id = 1 AND last_collection_owner_id = $1 - `, token); err != nil { + `, 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 }