feat(api): owner-token compare-and-clear for collection in-flight marker (closes #261)

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

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) {

internal/api/handler_recommendations_refresh.go

@@ -16,6 +16,11 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/lambda/types"
 )
 
+// markedCollectionRollbackTimeout bounds the detached best-effort clear issued
+// when the async self-invoke fails. Matches the scheduler's deferred clear and
+// Application.releaseSkippedCollectionMarker budgets.
+const markedCollectionRollbackTimeout = 5 * time.Second
+
 // LambdaInvokerInterface is the narrow subset of lambda.Client used by the
 // async refresh handler. Extracted so tests can inject a stub without
 // standing up a real Lambda client.
@@ -66,15 +71,18 @@ 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)
 	}
 	if !ok {
 		return nil, NewClientError(409, "recommendation collection already in progress; try again in a few minutes")
 	}
 
-	freshness, err := h.runMarkedCollection(ctx)
+	freshness, err := h.runMarkedCollection(ctx, token)
 	if err != nil {
 		return nil, err
 	}
@@ -103,13 +111,16 @@ func (h *Handler) postRefreshRecommendations(ctx context.Context, req *events.La
 //
 // Extracted from postRefreshRecommendations to keep that function under the
 // project's cyclomatic-complexity gate after the post-async re-read was added.
-func (h *Handler) runMarkedCollection(ctx context.Context) (*config.RecommendationsFreshness, error) {
+//
+// token is the owner token MarkCollectionStarted returned to the caller; it
+// is passed to the async invoke payload (so the scheduler can thread it
+// back into ClearCollectionStarted) and to the rollback clear on this
+// handler's own failure path.
+func (h *Handler) runMarkedCollection(ctx context.Context, token string) (*config.RecommendationsFreshness, error) {
 	schedulerARN := os.Getenv("SCHEDULER_LAMBDA_ARN")
 	if schedulerARN != "" {
-		if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN); invokeErr != nil {
-			if clearErr := h.config.ClearCollectionStarted(ctx); clearErr != nil {
-				logging.Warnf("runMarkedCollection: failed to clear collection started marker: %v", clearErr)
-			}
+		if invokeErr := h.asyncInvokeSelf(ctx, schedulerARN, token); invokeErr != nil {
+			h.rollbackMarkedCollection(token)
 			return nil, fmt.Errorf("failed to trigger async collection: %w", invokeErr)
 		}
 		freshness, err := h.config.GetRecommendationsFreshness(ctx)
@@ -119,8 +130,9 @@ func (h *Handler) runMarkedCollection(ctx context.Context) (*config.Recommendati
 		return freshness, nil
 	}
 	// Non-Lambda (HTTP) mode: collect synchronously. ClearCollectionStarted
-	// is called by the scheduler's deferred clearCollectionStartedBestEffort.
-	if _, collectErr := h.scheduler.CollectRecommendations(ctx); collectErr != nil {
+	// is called by the scheduler's deferred clearCollectionStartedBestEffort,
+	// passed this same token so the clear is scoped to this run.
+	if _, collectErr := h.scheduler.CollectRecommendations(ctx, token); collectErr != nil {
 		return nil, fmt.Errorf("collection failed: %w", collectErr)
 	}
 	freshness, err := h.config.GetRecommendationsFreshness(ctx)
@@ -130,12 +142,37 @@ func (h *Handler) runMarkedCollection(ctx context.Context) (*config.Recommendati
 	return freshness, nil
 }
 
+// rollbackMarkedCollection releases this caller's own in-flight marker after
+// the async self-invoke failed and no collection will ever run to release it.
+//
+// Deliberately detached from the request context. The dominant cause of an
+// asyncInvokeSelf failure is the request context itself expiring or being
+// canceled (API Gateway deadline, client disconnect, a slow SDK credential
+// refresh), so reusing that context here would fail the rollback in exactly
+// the case that produced it. The marker would then sit stranded for the full
+// 5-minute auto-recovery window in MarkCollectionStarted, rejecting every
+// refresh the user attempts with 409 while no collection is running. Matches
+// the detached-clear pattern in the scheduler's deferred clear and in
+// Application.releaseSkippedCollectionMarker.
+//
+// Scoped by token, so it can only ever release the marker this caller owns.
+// Best effort: a failure only defers cleanup to the 5-minute window, and the
+// caller already returns the underlying invoke error, so it is logged rather
+// than masking that error.
+func (h *Handler) rollbackMarkedCollection(token string) {
+	clearCtx, cancel := context.WithTimeout(context.Background(), markedCollectionRollbackTimeout)
+	defer cancel()
+	if err := h.config.ClearCollectionStarted(clearCtx, token); err != nil {
+		logging.Warnf("rollbackMarkedCollection: failed to clear collection started marker: %v", err)
+	}
+}
+
 // asyncInvokeSelf fires an InvocationType=Event invoke of the given Lambda
 // function ARN with the EventBridge-style payload that handleLambdaScheduledEvent
 // recognizes as a "collect recommendations" job. The call returns immediately;
 // the Lambda runtime delivers the event to the next available container
 // (which may be this same container's next invocation).
-func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN string) error {
+func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN, ownerToken string) error {
 	invoker, err := h.getLambdaInvoker(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to build Lambda client: %w", err)
@@ -159,8 +196,9 @@ func (h *Handler) asyncInvokeSelf(ctx context.Context, functionARN string) error
 	// Action != "") classifies this consistently with EventBridge cron
 	// deliveries that already exercise this code path.
 	payload, marshalErr := json.Marshal(map[string]string{
-		"source": "aws.events",
-		"action": "collect_recommendations",
+		"source":      "aws.events",
+		"action":      "collect_recommendations",
+		"owner_token": ownerToken,
 	})
 	if marshalErr != nil {
 		return fmt.Errorf("asyncInvokeSelf: failed to marshal payload: %w", marshalErr)

internal/api/handler_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/LeanerCloud/CUDly/internal/credentials"
 	"github.com/LeanerCloud/CUDly/internal/scheduler"
 	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-sdk-go-v2/service/lambda"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -509,10 +510,10 @@ func TestHandler_HandleRequest_RefreshRecommendations(t *testing.T) {
 	mockAuth.grantAdmin()
 	mockAuth.On("ValidateCSRFToken", ctx, mock.Anything, mock.Anything).Return(nil)
 
-	mockScheduler.On("CollectRecommendations", mock.Anything).Return(&scheduler.CollectResult{Recommendations: 0, TotalSavings: 0}, nil)
+	mockScheduler.On("CollectRecommendations", mock.Anything, mock.Anything).Return(&scheduler.CollectResult{Recommendations: 0, TotalSavings: 0}, nil)
 	mockStore.On("GetRecommendationsFreshness", mock.Anything).
 		Return(&config.RecommendationsFreshness{}, nil)
-	mockStore.On("MarkCollectionStarted", mock.Anything).Return(true, nil)
+	mockStore.On("MarkCollectionStarted", mock.Anything).Return("tok-xyz", true, nil)
 
 	handler := &Handler{config: mockStore, scheduler: mockScheduler, auth: mockAuth, apiKey: "test-key"}
 
@@ -535,7 +536,118 @@ func TestHandler_HandleRequest_RefreshRecommendations(t *testing.T) {
 	require.NoError(t, err)
 	// Sync-fallback path returns the response body fully populated, so 200.
 	assert.Equal(t, 200, resp.StatusCode)
-	mockScheduler.AssertCalled(t, "CollectRecommendations", mock.Anything)
+	mockScheduler.AssertCalled(t, "CollectRecommendations", mock.Anything, "tok-xyz")
+}
+
+// stubLambdaInvoker is a minimal LambdaInvokerInterface for exercising the
+// async-invoke path in runMarkedCollection without a real Lambda client.
+type stubLambdaInvoker struct {
+	invokeFn func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error)
+}
+
+func (s *stubLambdaInvoker) Invoke(ctx context.Context, params *lambda.InvokeInput, _ ...func(*lambda.Options)) (*lambda.InvokeOutput, error) {
+	return s.invokeFn(ctx, params)
+}
+
+// TestRunMarkedCollection_AsyncInvokeFailure_RollsBackOwnerToken pins issue
+// #261: when the async self-invoke fails, the rollback clear must use the
+// same token MarkCollectionStarted returned to this caller (not an
+// unconditional clear that could wipe an unrelated in-flight run).
+func TestRunMarkedCollection_AsyncInvokeFailure_RollsBackOwnerToken(t *testing.T) {
+	ctx := context.Background()
+	t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly")
+
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(t) })
+	mockStore.On("ClearCollectionStarted", mock.Anything, "tok-xyz").Return(nil)
+
+	invoker := &stubLambdaInvoker{
+		invokeFn: func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error) {
+			return nil, assert.AnError
+		},
+	}
+
+	handler := &Handler{config: mockStore, lambdaInvoker: invoker}
+
+	_, err := handler.runMarkedCollection(ctx, "tok-xyz")
+	require.Error(t, err)
+	mockStore.AssertCalled(t, "ClearCollectionStarted", mock.Anything, "tok-xyz")
+}
+
+// TestRunMarkedCollection_AsyncInvokeFailure_RollbackSurvivesCanceledCtx pins
+// the rollback clear to a DETACHED context. The dominant cause of an
+// asyncInvokeSelf failure is the request context itself expiring or being
+// canceled, so a rollback reusing that context would fail in exactly the case
+// that produced it, stranding this caller's marker for the full 5-minute
+// auto-recovery window and 409-ing every refresh in between.
+//
+// The context matcher lives on the expectation itself rather than in a
+// trailing AssertCalled: testify's AssertCalled compares arguments by value
+// and does not honor MatchedBy, whereas an unmatched expectation makes the
+// mock fail the call outright. Pre-fix the clear still happened, just with
+// the already-canceled request context, so asserting only that Clear was
+// called would pass with the bug present.
+func TestRunMarkedCollection_AsyncInvokeFailure_RollbackSurvivesCanceledCtx(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly")
+
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(t) })
+	liveDeadlineCtx := mock.MatchedBy(func(clearCtx context.Context) bool {
+		_, hasDeadline := clearCtx.Deadline()
+		return hasDeadline && clearCtx.Err() == nil
+	})
+	mockStore.On("ClearCollectionStarted", liveDeadlineCtx, "tok-xyz").Return(nil)
+
+	invoker := &stubLambdaInvoker{
+		invokeFn: func(invokeCtx context.Context, _ *lambda.InvokeInput) (*lambda.InvokeOutput, error) {
+			// Model the real failure mode: the request context dies mid-invoke
+			// and the SDK surfaces that as the invoke error.
+			cancel()
+			return nil, invokeCtx.Err()
+		},
+	}
+
+	handler := &Handler{config: mockStore, lambdaInvoker: invoker}
+
+	_, err := handler.runMarkedCollection(ctx, "tok-xyz")
+	require.Error(t, err)
+}
+
+// TestRunMarkedCollection_AsyncInvokeSuccess_PayloadCarriesOwnerToken pins
+// issue #261: the async-invoke payload must carry this run's owner token so
+// the scheduler receiving the event can scope its own clear to this run;
+// on the happy path this handler must NOT call Clear itself (the
+// scheduler's own deferred clear owns that once collection finishes).
+func TestRunMarkedCollection_AsyncInvokeSuccess_PayloadCarriesOwnerToken(t *testing.T) {
+	ctx := context.Background()
+	t.Setenv("SCHEDULER_LAMBDA_ARN", "arn:aws:lambda:us-east-1:123456789012:function:cudly")
+
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(t) })
+	mockStore.On("GetRecommendationsFreshness", mock.Anything).
+		Return(&config.RecommendationsFreshness{}, nil)
+	// Registered (as Maybe) so an unwanted clear is RECORDED: MockConfigStore
+	// short-circuits methods with no registered expectation before reaching
+	// mock.Called, so without this the AssertNotCalled below would pass even
+	// if the happy path did clear the marker.
+	mockStore.On("ClearCollectionStarted", mock.Anything, mock.Anything).Return(nil).Maybe()
+
+	var capturedPayload map[string]string
+	invoker := &stubLambdaInvoker{
+		invokeFn: func(ctx context.Context, params *lambda.InvokeInput) (*lambda.InvokeOutput, error) {
+			require.NoError(t, json.Unmarshal(params.Payload, &capturedPayload))
+			return &lambda.InvokeOutput{}, nil
+		},
+	}
+
+	handler := &Handler{config: mockStore, lambdaInvoker: invoker}
+
+	_, err := handler.runMarkedCollection(ctx, "tok-xyz")
+	require.NoError(t, err)
+
+	assert.Equal(t, "tok-xyz", capturedPayload["owner_token"])
+	mockStore.AssertNotCalled(t, "ClearCollectionStarted", mock.Anything, mock.Anything)
 }
 
 func TestHandler_HandleRequest_ListPlans(t *testing.T) {

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)
 	}

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-

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