From 52f404495db86b2b7416d38dedfef03806139ead Mon Sep 17 00:00:00 2001 From: Astrid Gealer Date: Mon, 19 Oct 2026 12:07:06 +0100 Subject: [PATCH 1/2] fix(scheduler): scope the collection marker release to the invocation, not the marker The owner token MarkCollectionStarted hands the refresh handler is shared by every Lambda delivery of the async collect event, and async delivery is at-least-once. When a duplicate delivery was skipped by the task advisory lock, releaseSkippedCollectionMarker cleared the marker with that shared token while the lock winner was still collecting, so the in-progress banner dropped early and a refresh in the remaining window returned 202 without collecting. The lock winner now calls ClaimCollectionInvocation right after taking the lock. It re-stamps last_collection_owner_id with a fresh invocation token when the marker is still owned by the marker token, or free or stale (the latter covers a duplicate that cleared in the window between the winner's lock and its re-stamp, and a late redelivery). The winner collects and clears under the new token, so the duplicate's skip-path clear no longer matches. A run skipped behind a tokenless cron run still releases its own marker. No schema change: the owner column already exists. Closes #84 Co-Authored-By: Claude Opus 5.5 --- internal/analytics/collector_test.go | 3 + internal/config/interfaces.go | 8 + ...ore_postgres_collection_invocation_test.go | 269 +++++++++++++ .../config/store_postgres_recommendations.go | 55 +++ internal/mocks/stores.go | 18 + internal/scheduler/scheduler.go | 4 +- internal/server/collection_invocation.go | 35 ++ internal/server/collection_invocation_test.go | 361 ++++++++++++++++++ internal/server/handler.go | 21 +- internal/server/test_helpers_test.go | 3 + 10 files changed, 760 insertions(+), 17 deletions(-) create mode 100644 internal/config/store_postgres_collection_invocation_test.go create mode 100644 internal/server/collection_invocation.go create mode 100644 internal/server/collection_invocation_test.go diff --git a/internal/analytics/collector_test.go b/internal/analytics/collector_test.go index f5422c072..c6d58d470 100644 --- a/internal/analytics/collector_test.go +++ b/internal/analytics/collector_test.go @@ -496,6 +496,9 @@ func (m *mockConfigStore) MarkCollectionStarted(_ context.Context) (string, bool func (m *mockConfigStore) ClearCollectionStarted(_ context.Context, _ string) error { return nil } +func (m *mockConfigStore) ClaimCollectionInvocation(_ context.Context, _ string) (string, bool, error) { + return "mock-invocation-token", true, nil +} func (m *mockConfigStore) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { return nil, nil } diff --git a/internal/config/interfaces.go b/internal/config/interfaces.go index 6483836dc..ffd0703a2 100644 --- a/internal/config/interfaces.go +++ b/internal/config/interfaces.go @@ -561,6 +561,14 @@ type StoreInterface interface { // 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 + // ClaimCollectionInvocation is called by a token-carrying collect run + // once it holds the task advisory lock (issue #84). It re-stamps the + // marker with a fresh invocation-scoped token when the marker is still + // owned by markerToken, or free/stale, so a lock-skipped duplicate + // delivery of the same event (same markerToken) can no longer clear + // it. Returns the new token and true when the caller now owns the + // marker; ("", false, nil) when another run holds a fresh marker. + ClaimCollectionInvocation(ctx context.Context, markerToken string) (token string, ok bool, err error) // RI utilization cache. Postgres-backed TTL cache for Cost Explorer // GetReservationUtilization; shared across Lambda containers so diff --git a/internal/config/store_postgres_collection_invocation_test.go b/internal/config/store_postgres_collection_invocation_test.go new file mode 100644 index 000000000..db295f0ba --- /dev/null +++ b/internal/config/store_postgres_collection_invocation_test.go @@ -0,0 +1,269 @@ +//go:build integration +// +build integration + +package config_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/migrations" + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/testhelpers" + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// collectionMarker is the raw recommendations_state marker, read straight +// from the row because GetRecommendationsFreshness does not expose the owner. +type collectionMarker struct { + startedAt *time.Time + owner *string +} + +func readCollectionMarker(ctx context.Context, t *testing.T, pool *pgxpool.Pool) collectionMarker { + t.Helper() + var m collectionMarker + require.NoError(t, pool.QueryRow(ctx, ` + SELECT last_collection_started_at, last_collection_owner_id::text + FROM recommendations_state + WHERE id = 1 + `).Scan(&m.startedAt, &m.owner)) + return m +} + +func resetCollectionMarker(ctx context.Context, t *testing.T, pool *pgxpool.Pool) { + t.Helper() + _, err := pool.Exec(ctx, ` + UPDATE recommendations_state + SET last_collection_started_at = NULL, + last_collection_owner_id = NULL + WHERE id = 1 + `) + require.NoError(t, err) +} + +func markCollection(ctx context.Context, t *testing.T, store *config.PostgresStore) string { + t.Helper() + token, ok, err := store.MarkCollectionStarted(ctx) + require.NoError(t, err) + require.True(t, ok) + return token +} + +// TestPostgresStore_ClaimCollectionInvocation pins issue #84 at the store +// level. The marker token MarkCollectionStarted hands out is shared by every +// Lambda delivery of the refresh's async event; the delivery that wins the +// task advisory lock re-stamps the marker with a fresh invocation token so a +// lock-skipped duplicate's ClearCollectionStarted(markerToken) no longer +// matches it. +func TestPostgresStore_ClaimCollectionInvocation(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) + + t.Run("owned marker is re-stamped and the duplicate's clear no longer matches", func(t *testing.T) { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + before := readCollectionMarker(ctx, t, pool) + + invocationToken, ok, err := store.ClaimCollectionInvocation(ctx, markerToken) + require.NoError(t, err) + require.True(t, ok) + _, parseErr := uuid.Parse(invocationToken) + require.NoError(t, parseErr, "invocation token must be a UUID") + require.NotEqual(t, markerToken, invocationToken) + + after := readCollectionMarker(ctx, t, pool) + require.NotNil(t, after.owner) + assert.Equal(t, invocationToken, *after.owner) + require.NotNil(t, after.startedAt) + assert.True(t, before.startedAt.Equal(*after.startedAt), + "re-stamping an owned marker must keep the started_at the refresh recorded") + + // The duplicate delivery, lock-skipped, releases with the marker token. + require.NoError(t, store.ClearCollectionStarted(ctx, markerToken)) + assert.NotNil(t, readCollectionMarker(ctx, t, pool).startedAt, + "the duplicate's marker-token clear must not release the winner's marker") + + // The winner's deferred clear with its invocation token releases it. + require.NoError(t, store.ClearCollectionStarted(ctx, invocationToken)) + m := readCollectionMarker(ctx, t, pool) + assert.Nil(t, m.startedAt) + assert.Nil(t, m.owner) + }) + + t.Run("marker already cleared by the duplicate is re-established", func(t *testing.T) { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + require.NoError(t, store.ClearCollectionStarted(ctx, markerToken)) + + invocationToken, ok, err := store.ClaimCollectionInvocation(ctx, markerToken) + require.NoError(t, err) + require.True(t, ok) + m := readCollectionMarker(ctx, t, pool) + require.NotNil(t, m.startedAt, "the lock winner must re-establish a cleared marker") + require.NotNil(t, m.owner) + assert.Equal(t, invocationToken, *m.owner) + assert.WithinDuration(t, time.Now(), *m.startedAt, time.Minute) + }) + + t.Run("fresh marker owned by another run is left alone", func(t *testing.T) { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + require.NoError(t, store.ClearCollectionStarted(ctx, markerToken)) + otherToken := markCollection(ctx, t, store) + + token, ok, err := store.ClaimCollectionInvocation(ctx, markerToken) + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, token) + m := readCollectionMarker(ctx, t, pool) + require.NotNil(t, m.owner) + assert.Equal(t, otherToken, *m.owner, "another run's fresh marker must keep its owner") + }) + + t.Run("stale marker owned by another run is taken over", func(t *testing.T) { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + require.NoError(t, store.ClearCollectionStarted(ctx, markerToken)) + markCollection(ctx, t, store) + _, err := pool.Exec(ctx, ` + UPDATE recommendations_state + SET last_collection_started_at = NOW() - INTERVAL '6 minutes' + WHERE id = 1 + `) + require.NoError(t, err) + + invocationToken, ok, err := store.ClaimCollectionInvocation(ctx, markerToken) + require.NoError(t, err) + require.True(t, ok) + m := readCollectionMarker(ctx, t, pool) + require.NotNil(t, m.owner) + assert.Equal(t, invocationToken, *m.owner) + assert.WithinDuration(t, time.Now(), *m.startedAt, time.Minute, + "taking over a stale marker must restart its 5-minute window") + }) + + t.Run("empty marker token is a boundary error", func(t *testing.T) { + _, ok, err := store.ClaimCollectionInvocation(ctx, "") + require.Error(t, err) + assert.False(t, ok) + }) +} + +// TestPostgresStore_ClaimCollectionInvocation_Races runs the issue #84 +// interleavings concurrently against real Postgres, where the two UPDATEs +// on recommendations_state row 1 serialize on its row lock and the loser +// re-evaluates its WHERE clause against the winner's committed row. +func TestPostgresStore_ClaimCollectionInvocation_Races(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) + + // The lock winner's claim races the lock-skipped duplicate's clear, both + // with the shared marker token. Whichever commits first, the winner must + // end up owning a set marker: claim-first makes the clear a no-op, and + // clear-first leaves a free marker the claim re-establishes. + t.Run("winner's claim against the duplicate's clear", func(t *testing.T) { + claimFirst, clearFirst := 0, 0 + for i := 0; i < 50; i++ { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + markedAt := readCollectionMarker(ctx, t, pool).startedAt + + var ( + wg sync.WaitGroup + invocationToken string + claimed bool + claimErr, clearErr error + start = make(chan struct{}) + ) + wg.Add(2) + go func() { + defer wg.Done() + <-start + invocationToken, claimed, claimErr = store.ClaimCollectionInvocation(ctx, markerToken) + }() + go func() { + defer wg.Done() + <-start + clearErr = store.ClearCollectionStarted(ctx, markerToken) + }() + close(start) + wg.Wait() + + require.NoError(t, claimErr) + require.NoError(t, clearErr) + require.True(t, claimed, "iteration %d: the lock winner must own the marker", i) + m := readCollectionMarker(ctx, t, pool) + require.NotNil(t, m.startedAt, "iteration %d: marker must be set while the winner collects", i) + require.NotNil(t, m.owner) + require.Equal(t, invocationToken, *m.owner, "iteration %d", i) + // An unchanged started_at means the claim re-stamped the owned + // marker; a new one means the clear landed first and the claim + // re-established it. + if m.startedAt.Equal(*markedAt) { + claimFirst++ + } else { + clearFirst++ + } + } + t.Logf("orderings observed: claim-first=%d clear-first=%d", claimFirst, clearFirst) + }) + + // Concurrent claims with the same marker token: exactly one becomes the + // owner. Production only claims under the advisory lock, so this is a + // defense-in-depth check that the UPDATE itself is single-winner. + t.Run("concurrent claims elect one owner", func(t *testing.T) { + const claimers = 8 + for i := 0; i < 20; i++ { + resetCollectionMarker(ctx, t, pool) + markerToken := markCollection(ctx, t, store) + + var ( + wg sync.WaitGroup + mu sync.Mutex + owners []string + errs []error + start = make(chan struct{}) + ) + for c := 0; c < claimers; c++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + token, ok, err := store.ClaimCollectionInvocation(ctx, markerToken) + mu.Lock() + defer mu.Unlock() + if err != nil { + errs = append(errs, err) + } + if ok { + owners = append(owners, token) + } + }() + } + close(start) + wg.Wait() + + require.Empty(t, errs) + require.Len(t, owners, 1, "iteration %d: exactly one claim may win", i) + m := readCollectionMarker(ctx, t, pool) + require.NotNil(t, m.owner) + require.Equal(t, owners[0], *m.owner) + } + }) +} diff --git a/internal/config/store_postgres_recommendations.go b/internal/config/store_postgres_recommendations.go index d95877c47..9efd969fa 100644 --- a/internal/config/store_postgres_recommendations.go +++ b/internal/config/store_postgres_recommendations.go @@ -601,3 +601,58 @@ func (s *PostgresStore) ClearCollectionStarted(ctx context.Context, token string } return nil } + +// ClaimCollectionInvocation re-stamps the collection in-flight marker with a +// fresh, invocation-scoped owner token (issue #84). The scheduler calls it +// once a token-carrying collect run has WON the task advisory lock, before it +// starts collecting. +// +// The token MarkCollectionStarted hands out identifies the marker, not the +// run that serves it: Lambda's async delivery is at-least-once, so a +// duplicate delivery of the same event carries the same token. Without the +// re-stamp, a duplicate that is skipped by the advisory lock clears the +// marker with that shared token while the lock winner is still collecting. +// After the re-stamp only the winner's new token matches, so the duplicate's +// skip-path clear is a no-op and the winner's deferred clear (with the new +// token) is the one that releases the marker. +// +// The marker is taken when it is still owned by markerToken (started_at is +// kept as MarkCollectionStarted set it) or when it is free or stale under the +// same 5-minute rule as MarkCollectionStarted (started_at is set to NOW()). +// The second case covers the window between the winner taking the lock and +// running this UPDATE, in which a lock-skipped duplicate may already have +// cleared the marker: the lock winner is the run actually collecting, so it +// re-establishes the marker rather than collecting with none. It also covers +// a late redelivery after the original run already finished and cleared. +// +// Returns the new token and true when the caller now owns the marker. Returns +// ("", false, nil) when a different run holds a fresh marker; the caller then +// owns nothing and must not clear. An empty markerToken is a boundary error, +// as in ClearCollectionStarted. +func (s *PostgresStore) ClaimCollectionInvocation(ctx context.Context, markerToken string) (token string, ok bool, err error) { + if markerToken == "" { + return "", false, fmt.Errorf("marker token must not be empty") + } + token = uuid.New().String() + tag, err := s.db.Exec(ctx, ` + UPDATE recommendations_state + SET last_collection_started_at = CASE + WHEN last_collection_owner_id = $1 THEN last_collection_started_at + ELSE NOW() + END, + last_collection_owner_id = $2 + WHERE id = 1 + AND ( + last_collection_owner_id = $1 + OR last_collection_started_at IS NULL + OR last_collection_started_at < NOW() - INTERVAL '5 minutes' + ) + `, markerToken, token) + if err != nil { + return "", false, fmt.Errorf("failed to claim collection invocation: %w", err) + } + if tag.RowsAffected() != 1 { + return "", false, nil + } + return token, true, nil +} diff --git a/internal/mocks/stores.go b/internal/mocks/stores.go index a4c9546f7..83e295a09 100644 --- a/internal/mocks/stores.go +++ b/internal/mocks/stores.go @@ -20,6 +20,11 @@ import ( // implicitly keep getting a non-empty token to thread through. const MockOwnerToken = "mock-owner-token" +// MockInvocationToken is the default token MockConfigStore.ClaimCollectionInvocation +// returns when no expectation is registered. It differs from MockOwnerToken +// so tests can tell the re-stamped invocation token from the marker token. +const MockInvocationToken = "mock-invocation-token" + // MockConfigStore is a shared testify-based mock for config.StoreInterface. // // Most methods dispatch through m.Called only when an expectation has been @@ -1981,6 +1986,19 @@ func (m *MockConfigStore) ClearCollectionStarted(ctx context.Context, token stri return m.Called(ctx, token).Error(0) } +// ClaimCollectionInvocation mocks the ClaimCollectionInvocation operation. +// Defaults to (MockInvocationToken, true, nil) when no expectation is +// registered, so a token-carrying collect run that wins the lock keeps +// owning the marker under a distinct invocation token. +func (m *MockConfigStore) ClaimCollectionInvocation(ctx context.Context, markerToken string) (token string, ok bool, err error) { + m.record("ClaimCollectionInvocation", ctx, markerToken) + if !isExpected(&m.Mock, "ClaimCollectionInvocation") { + return MockInvocationToken, true, nil + } + args := m.Called(ctx, markerToken) + return args.String(0), args.Bool(1), args.Error(2) +} + // StampRIExchangeApprovedBy mocks the StampRIExchangeApprovedBy operation. func (m *MockConfigStore) StampRIExchangeApprovedBy(ctx context.Context, id, approverEmail string) error { m.record("StampRIExchangeApprovedBy", ctx, id, approverEmail) diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index f6d4b4051..e1dee9daa 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -158,7 +158,9 @@ func cacheTTLFromEnv() time.Duration { // 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; kickColdStartCollection's +// returned the token that ownerToken carries here (in Lambda, re-stamped to +// an invocation-scoped token once the run won the task advisory lock, see +// ClaimCollectionInvocation, issue #84); kickColdStartCollection's // background collect does too (issue #106), for the same reason. The cron // case calls with an empty ownerToken since it never wins a marker of its // own (it coordinates via the advisory lock instead, see diff --git a/internal/server/collection_invocation.go b/internal/server/collection_invocation.go new file mode 100644 index 000000000..60f121219 --- /dev/null +++ b/internal/server/collection_invocation.go @@ -0,0 +1,35 @@ +package server + +import ( + "context" + "log" +) + +// claimCollectionInvocation runs once a token-carrying collect run holds the +// advisory lock (see HandleScheduledTask), and swaps the marker token it was +// invoked with for a fresh invocation-scoped one (issue #84). The marker +// token MarkCollectionStarted handed the refresh handler is shared by every +// delivery of the same async event, and Lambda async delivery is +// at-least-once. After the swap a lock-skipped duplicate delivery's +// releaseSkippedCollectionMarker no longer matches the marker, so it cannot +// clear it while this run is still collecting. +// +// When another run holds a fresh marker this run owns nothing, so it collects +// tokenless and never clears. A store error keeps the marker token: that is +// the pre-#84 behavior, and failing a collection the lock already admitted +// over a bookkeeping write would be worse. +func (app *Application) claimCollectionInvocation(ctx context.Context, taskType ScheduledTaskType, params ScheduledTaskParams) ScheduledTaskParams { + if taskType != TaskCollectRecommendations || params.OwnerToken == "" || app.Config == nil { + return params + } + token, ok, err := app.Config.ClaimCollectionInvocation(ctx, params.OwnerToken) + if err != nil { + log.Printf("Failed to claim collection marker for this invocation, keeping the marker token: %v", err) + return params + } + if !ok { + log.Println("Collection marker is held by another run; collecting without an owner token") + return ScheduledTaskParams{} + } + return ScheduledTaskParams{OwnerToken: token} +} diff --git a/internal/server/collection_invocation_test.go b/internal/server/collection_invocation_test.go new file mode 100644 index 000000000..c9815425b --- /dev/null +++ b/internal/server/collection_invocation_test.go @@ -0,0 +1,361 @@ +package server + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/mocks" + "github.com/LeanerCloud/cloud-commitments-platform/internal/scheduler" + "github.com/LeanerCloud/cloud-commitments-platform/internal/testutil" + "github.com/google/uuid" + "github.com/stretchr/testify/mock" +) + +// markerStore is an in-memory stand-in for the recommendations_state +// collection marker. It mirrors the PostgresStore semantics of +// MarkCollectionStarted, ClearCollectionStarted (compare-and-clear) and +// ClaimCollectionInvocation (re-stamp when owned, free or, in Postgres, +// stale; this fake never goes stale), so the server-level tests can observe +// whether the marker is still set while a collect run is in progress. The +// real SQL is pinned against Postgres in +// internal/config/store_postgres_recommendations_test.go. +type markerStore struct { + *mocks.MockConfigStore + mu sync.Mutex + owner string + set bool +} + +func newMarkerStore() *markerStore { + return &markerStore{MockConfigStore: new(mocks.MockConfigStore)} +} + +func (s *markerStore) MarkCollectionStarted(_ context.Context) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.set { + return "", false, nil + } + s.owner, s.set = uuid.NewString(), true + return s.owner, true, nil +} + +func (s *markerStore) ClearCollectionStarted(_ context.Context, token string) error { + if token == "" { + return errors.New("owner token must not be empty") + } + s.mu.Lock() + defer s.mu.Unlock() + if s.set && s.owner == token { + s.owner, s.set = "", false + } + return nil +} + +func (s *markerStore) ClaimCollectionInvocation(_ context.Context, markerToken string) (string, bool, error) { + if markerToken == "" { + return "", false, errors.New("marker token must not be empty") + } + s.mu.Lock() + defer s.mu.Unlock() + if s.set && s.owner != markerToken { + return "", false, nil + } + s.owner, s.set = uuid.NewString(), true + return s.owner, true, nil +} + +func (s *markerStore) markerSet() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.set +} + +// sharedTaskLocker models pg_try_advisory_lock across concurrent +// invocations: the first TryAdvisoryLock for a lock ID wins until it is +// released, and every other attempt in the meantime is refused. +// +// afterAcquire, when set, runs once right after a successful acquisition and +// before TryAdvisoryLock returns, so a test can interleave another +// invocation between the winner taking the lock and its next step. +type sharedTaskLocker struct { + afterAcquire func() + held map[int64]bool + mu sync.Mutex +} + +func (l *sharedTaskLocker) TryAdvisoryLock(_ context.Context, lockID int64) (bool, error) { + l.mu.Lock() + if l.held == nil { + l.held = map[int64]bool{} + } + if l.held[lockID] { + l.mu.Unlock() + return false, nil + } + l.held[lockID] = true + hook := l.afterAcquire + l.afterAcquire = nil + l.mu.Unlock() + if hook != nil { + hook() + } + return true, nil +} + +func (l *sharedTaskLocker) ReleaseAdvisoryLock(_ context.Context, lockID int64) { + l.mu.Lock() + defer l.mu.Unlock() + delete(l.held, lockID) +} + +// blockingCollector stands in for scheduler.CollectRecommendations: it +// reports that collection started, blocks until released, then performs the +// scheduler's deferred token-scoped clear with the token it was handed. +type blockingCollector struct { + store *markerStore + started chan struct{} + release chan struct{} + gotToken string +} + +func newBlockingCollector(store *markerStore) *blockingCollector { + return &blockingCollector{store: store, started: make(chan struct{}), release: make(chan struct{})} +} + +func (c *blockingCollector) collect(ctx context.Context, ownerToken string) (*scheduler.CollectResult, error) { + c.gotToken = ownerToken + close(c.started) + <-c.release + if ownerToken != "" { + _ = c.store.ClearCollectionStarted(ctx, ownerToken) + } + return &scheduler.CollectResult{}, nil +} + +// TestHandleScheduledTask_DuplicateDeliveryKeepsWinnersMarker pins issue +// #84. Lambda's async invoke is at-least-once, so the refresh handler's +// self-invoke can be delivered twice with the same owner token. The first +// delivery wins the advisory lock and collects; the duplicate is +// lock-skipped. Before the fix the duplicate's skip-path release cleared the +// marker with that shared token while the winner was still collecting, so +// the in-progress banner dropped early and a refresh in the remaining window +// returned 202 without collecting. The winner now re-stamps the marker with +// its own invocation token once it holds the lock, so the duplicate's clear +// no longer matches. +func TestHandleScheduledTask_DuplicateDeliveryKeepsWinnersMarker(t *testing.T) { + ctx := testutil.TestContext(t) + store := newMarkerStore() + collector := newBlockingCollector(store) + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{CollectRecommendationsFunc: collector.collect}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &sharedTaskLocker{}, + } + + markerToken, ok, err := store.MarkCollectionStarted(ctx) + testutil.AssertNoError(t, err) + testutil.AssertTrue(t, ok, "refresh must win the marker") + params := ScheduledTaskParams{OwnerToken: markerToken} + + winnerDone := make(chan error, 1) + go func() { + _, runErr := app.HandleScheduledTask(ctx, TaskCollectRecommendations, params) + winnerDone <- runErr + }() + <-collector.started + + // The duplicate delivery: same event, same token, lock held by the winner. + result, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, params) + testutil.AssertNoError(t, err) + m, isMap := result.(map[string]string) + testutil.AssertTrue(t, isMap && m["status"] == "skipped", "duplicate delivery must be lock-skipped") + + markerDuringCollect := store.markerSet() + + close(collector.release) + testutil.AssertNoError(t, <-winnerDone) + + testutil.AssertTrue(t, markerDuringCollect, + "the lock-skipped duplicate must not clear the marker while the winner is still collecting") + testutil.AssertTrue(t, collector.gotToken != "" && collector.gotToken != markerToken, + "the winner must collect under its own invocation token, not the shared marker token") + testutil.AssertFalse(t, store.markerSet(), + "the winner's deferred clear with its invocation token must release the marker") +} + +// TestHandleScheduledTask_DuplicateClearsBeforeWinnerClaims covers the +// window between the winner taking the advisory lock and re-stamping the +// marker: a duplicate delivery that is lock-skipped inside that window still +// matches the shared marker token and clears the marker. The winner must +// then re-establish the marker under its own invocation token rather than +// collect with none, or the issue #84 symptom survives in that window. +func TestHandleScheduledTask_DuplicateClearsBeforeWinnerClaims(t *testing.T) { + ctx := testutil.TestContext(t) + store := newMarkerStore() + collector := newBlockingCollector(store) + locker := &sharedTaskLocker{} + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{CollectRecommendationsFunc: collector.collect}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: locker, + } + + markerToken, ok, err := store.MarkCollectionStarted(ctx) + testutil.AssertNoError(t, err) + testutil.AssertTrue(t, ok, "refresh must win the marker") + params := ScheduledTaskParams{OwnerToken: markerToken} + + var duplicateErr error + markerClearedInWindow := false + locker.afterAcquire = func() { + _, duplicateErr = app.HandleScheduledTask(ctx, TaskCollectRecommendations, params) + markerClearedInWindow = !store.markerSet() + } + + winnerDone := make(chan error, 1) + go func() { + _, runErr := app.HandleScheduledTask(ctx, TaskCollectRecommendations, params) + winnerDone <- runErr + }() + <-collector.started + markerDuringCollect := store.markerSet() + close(collector.release) + testutil.AssertNoError(t, <-winnerDone) + testutil.AssertNoError(t, duplicateErr) + + testutil.AssertTrue(t, markerClearedInWindow, "precondition: the duplicate cleared the marker inside the window") + testutil.AssertTrue(t, markerDuringCollect, "the lock winner must re-establish the marker before collecting") + testutil.AssertTrue(t, collector.gotToken != "" && collector.gotToken != markerToken, + "the winner must collect under its own invocation token") + testutil.AssertFalse(t, store.markerSet(), "the winner's deferred clear must release the marker it re-established") +} + +// TestHandleScheduledTask_SkippedBehindCronStillReleases pins the case the +// lock-skip release exists for (issue #261 follow-up), which the issue #84 +// fix must keep: when the advisory lock is held by a tokenless cron run, a +// token-carrying refresh run that is skipped still owns the marker and must +// release it, or every refresh is rejected with 409 for five minutes. +func TestHandleScheduledTask_SkippedBehindCronStillReleases(t *testing.T) { + ctx := testutil.TestContext(t) + store := newMarkerStore() + collector := newBlockingCollector(store) + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{CollectRecommendationsFunc: collector.collect}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &sharedTaskLocker{}, + } + + cronDone := make(chan error, 1) + go func() { + _, runErr := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{}) + cronDone <- runErr + }() + <-collector.started + + markerToken, ok, err := store.MarkCollectionStarted(ctx) + testutil.AssertNoError(t, err) + testutil.AssertTrue(t, ok, "refresh must win the marker") + _, err = app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{OwnerToken: markerToken}) + testutil.AssertNoError(t, err) + + markerAfterSkip := store.markerSet() + close(collector.release) + testutil.AssertNoError(t, <-cronDone) + + testutil.AssertFalse(t, markerAfterSkip, "a refresh run skipped behind cron must release its own marker") + testutil.AssertEqual(t, "", collector.gotToken) +} + +// TestHandleScheduledTask_ClaimInvocationOutcomes pins how the lock winner +// threads the ClaimCollectionInvocation result into the collect run. +func TestHandleScheduledTask_ClaimInvocationOutcomes(t *testing.T) { + tests := []struct { + claimErr error + name string + claimToken string + wantToken string + claimOK bool + }{ + { + name: "claimed: collect under the invocation token", + claimToken: "tok-invocation", claimOK: true, + wantToken: "tok-invocation", + }, + { + // Another run holds a fresh marker: this run owns nothing, so it + // collects tokenless and its deferred clear is skipped rather + // than wiping the other run's marker. + name: "marker held by another run: collect tokenless", + wantToken: "", + }, + { + // A store error degrades to the pre-#84 marker-scoped behavior + // rather than failing a collection the lock already admitted. + name: "claim error: keep the marker token", + claimErr: errors.New("db connection lost"), + wantToken: "tok-marker", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + t.Cleanup(func() { store.AssertExpectations(t) }) + store.On("ClaimCollectionInvocation", mock.Anything, "tok-marker"). + Return(tt.claimToken, tt.claimOK, tt.claimErr).Once() + + gotToken := "unset" + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{ + CollectRecommendationsFunc: func(_ context.Context, ownerToken string) (*scheduler.CollectResult, error) { + gotToken = ownerToken + return &scheduler.CollectResult{}, nil + }, + }, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{acquired: true}, + } + + _, err := app.HandleScheduledTask(ctx, TaskCollectRecommendations, ScheduledTaskParams{OwnerToken: "tok-marker"}) + testutil.AssertNoError(t, err) + testutil.AssertEqual(t, tt.wantToken, gotToken) + }) + } +} + +// TestHandleScheduledTask_ClaimOnlyForTokenCarryingCollect pins that the +// re-stamp is confined to a token-carrying collect_recommendations run that +// won the lock: a tokenless (cron) collect owns no marker and another task +// type never touches it. +func TestHandleScheduledTask_ClaimOnlyForTokenCarryingCollect(t *testing.T) { + cases := []struct { + name string + taskType ScheduledTaskType + params ScheduledTaskParams + }{ + {name: "tokenless collect", taskType: TaskCollectRecommendations}, + {name: "other task type with a stray token", taskType: TaskProcessScheduledPurchases, params: ScheduledTaskParams{OwnerToken: "tok-marker"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ctx := testutil.TestContext(t) + store := new(mocks.MockConfigStore) + app := &Application{ + Config: store, + Scheduler: &testutil.MockScheduler{}, + Purchase: &testutil.MockPurchaseManager{}, + TaskLocker: &mockTaskLocker{acquired: true}, + } + _, err := app.HandleScheduledTask(ctx, tc.taskType, tc.params) + testutil.AssertNoError(t, err) + store.AssertNumberOfCalls(t, "ClaimCollectionInvocation", 0) + }) + } +} diff --git a/internal/server/handler.go b/internal/server/handler.go index 71c1659ad..e81df4d5a 100644 --- a/internal/server/handler.go +++ b/internal/server/handler.go @@ -112,6 +112,7 @@ func (app *Application) HandleScheduledTask(ctx context.Context, taskType Schedu return map[string]string{"status": "skipped", "reason": "already_running"}, nil } defer locker.ReleaseAdvisoryLock(ctx, lockID) + params = app.claimCollectionInvocation(ctx, taskType, params) } return app.dispatchTask(ctx, taskType, params) @@ -136,22 +137,10 @@ 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. +// A lock held by a duplicate delivery of this same event (Lambda async invoke +// is at-least-once, so it carries the same token) is safe too: that winner +// re-stamped the marker with its own invocation token on taking the lock +// (claimCollectionInvocation, issue #84), so this clear no longer matches it. // // Scoped by ownerToken, so this can only ever release the marker this run // owns, never a concurrent run's. Callers that never won MarkCollectionStarted diff --git a/internal/server/test_helpers_test.go b/internal/server/test_helpers_test.go index bf5d1ee90..0965fc55a 100644 --- a/internal/server/test_helpers_test.go +++ b/internal/server/test_helpers_test.go @@ -347,6 +347,9 @@ func (m *mockConfigStoreForHealth) MarkCollectionStarted(_ context.Context) (str func (m *mockConfigStoreForHealth) ClearCollectionStarted(_ context.Context, _ string) error { return nil } +func (m *mockConfigStoreForHealth) ClaimCollectionInvocation(_ context.Context, _ string) (string, bool, error) { + return "mock-invocation-token", true, nil +} func (m *mockConfigStoreForHealth) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { return nil, nil } From d0f151eceef8107118176f2ea3e84684f2578cf3 Mon Sep 17 00:00:00 2001 From: Astrid Gealer Date: Fri, 23 Oct 2026 03:15:57 +0000 Subject: [PATCH 2/2] test(config): move the #84 store tests onto the shared Postgres container After merging main, TestIntegrationTestsShareOnePostgresContainer (#1165) flagged the two new ClaimCollectionInvocation integration tests for starting a container each. They now take a migrated database from testhelpers.RequireMigratedDatabase, like the rest of internal/config. Also from review and the merge: - give sharedTaskLocker AdvisoryLockHeldFor, which main added to TaskLocker - claimCollectionInvocation clears OwnerToken on a copy of params instead of returning a fresh struct, so a field added later is not dropped - point the markerStore doc at store_postgres_collection_invocation_test.go and the claim doc at runScheduledTask - use mocks.MockInvocationToken in the analytics and server test fakes Co-Authored-By: Claude Opus 5.5 --- internal/analytics/collector_test.go | 3 ++- ...ore_postgres_collection_invocation_test.go | 19 ++++++------------- internal/server/collection_invocation.go | 8 +++++--- internal/server/collection_invocation_test.go | 12 +++++++++++- internal/server/test_helpers_test.go | 3 ++- 5 files changed, 26 insertions(+), 19 deletions(-) diff --git a/internal/analytics/collector_test.go b/internal/analytics/collector_test.go index dd13ddaf9..14d284433 100644 --- a/internal/analytics/collector_test.go +++ b/internal/analytics/collector_test.go @@ -10,6 +10,7 @@ import ( "github.com/LeanerCloud/cloud-commitments-go/pkg/common" "github.com/LeanerCloud/cloud-commitments-go/pkg/ladder" "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/LeanerCloud/cloud-commitments-platform/internal/mocks" "github.com/jackc/pgx/v5" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -512,7 +513,7 @@ func (m *mockConfigStore) ClearCollectionStarted(_ context.Context, _ string) er return nil } func (m *mockConfigStore) ClaimCollectionInvocation(_ context.Context, _ string) (string, bool, error) { - return "mock-invocation-token", true, nil + return mocks.MockInvocationToken, true, nil } func (m *mockConfigStore) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { return nil, nil diff --git a/internal/config/store_postgres_collection_invocation_test.go b/internal/config/store_postgres_collection_invocation_test.go index db295f0ba..2db99ab89 100644 --- a/internal/config/store_postgres_collection_invocation_test.go +++ b/internal/config/store_postgres_collection_invocation_test.go @@ -10,7 +10,6 @@ import ( "time" "github.com/LeanerCloud/cloud-commitments-platform/internal/config" - "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/migrations" "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/testhelpers" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" @@ -63,12 +62,9 @@ func markCollection(ctx context.Context, t *testing.T, store *config.PostgresSto // matches it. func TestPostgresStore_ClaimCollectionInvocation(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) + db := testhelpers.RequireMigratedDatabase(ctx, t, getMigrationsPath()).DB + pool := db.Pool() + store := config.NewPostgresStore(db) t.Run("owned marker is re-stamped and the duplicate's clear no longer matches", func(t *testing.T) { resetCollectionMarker(ctx, t, pool) @@ -166,12 +162,9 @@ func TestPostgresStore_ClaimCollectionInvocation(t *testing.T) { // re-evaluates its WHERE clause against the winner's committed row. func TestPostgresStore_ClaimCollectionInvocation_Races(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) + db := testhelpers.RequireMigratedDatabase(ctx, t, getMigrationsPath()).DB + pool := db.Pool() + store := config.NewPostgresStore(db) // The lock winner's claim races the lock-skipped duplicate's clear, both // with the shared marker token. Whichever commits first, the winner must diff --git a/internal/server/collection_invocation.go b/internal/server/collection_invocation.go index 60f121219..c3b5d00d9 100644 --- a/internal/server/collection_invocation.go +++ b/internal/server/collection_invocation.go @@ -6,7 +6,7 @@ import ( ) // claimCollectionInvocation runs once a token-carrying collect run holds the -// advisory lock (see HandleScheduledTask), and swaps the marker token it was +// advisory lock (see runScheduledTask), and swaps the marker token it was // invoked with for a fresh invocation-scoped one (issue #84). The marker // token MarkCollectionStarted handed the refresh handler is shared by every // delivery of the same async event, and Lambda async delivery is @@ -29,7 +29,9 @@ func (app *Application) claimCollectionInvocation(ctx context.Context, taskType } if !ok { log.Println("Collection marker is held by another run; collecting without an owner token") - return ScheduledTaskParams{} + params.OwnerToken = "" + return params } - return ScheduledTaskParams{OwnerToken: token} + params.OwnerToken = token + return params } diff --git a/internal/server/collection_invocation_test.go b/internal/server/collection_invocation_test.go index c9815425b..42dd9c0ab 100644 --- a/internal/server/collection_invocation_test.go +++ b/internal/server/collection_invocation_test.go @@ -5,6 +5,7 @@ import ( "errors" "sync" "testing" + "time" "github.com/LeanerCloud/cloud-commitments-platform/internal/mocks" "github.com/LeanerCloud/cloud-commitments-platform/internal/scheduler" @@ -20,7 +21,7 @@ import ( // stale; this fake never goes stale), so the server-level tests can observe // whether the marker is still set while a collect run is in progress. The // real SQL is pinned against Postgres in -// internal/config/store_postgres_recommendations_test.go. +// internal/config/store_postgres_collection_invocation_test.go. type markerStore struct { *mocks.MockConfigStore mu sync.Mutex @@ -111,6 +112,15 @@ func (l *sharedTaskLocker) ReleaseAdvisoryLock(_ context.Context, lockID int64) delete(l.held, lockID) } +// AdvisoryLockHeldFor reports a held lock as freshly taken, so a lock-skipped +// invocation takes the normal "already_running" path rather than the +// stuck-lock error (issue #394), which these tests do not exercise. +func (l *sharedTaskLocker) AdvisoryLockHeldFor(_ context.Context, lockID int64) (time.Duration, bool, error) { + l.mu.Lock() + defer l.mu.Unlock() + return 0, l.held[lockID], nil +} + // blockingCollector stands in for scheduler.CollectRecommendations: it // reports that collection started, blocks until released, then performs the // scheduler's deferred token-scoped clear with the token it was handed. diff --git a/internal/server/test_helpers_test.go b/internal/server/test_helpers_test.go index 652b5b902..a01c6c968 100644 --- a/internal/server/test_helpers_test.go +++ b/internal/server/test_helpers_test.go @@ -9,6 +9,7 @@ import ( "github.com/LeanerCloud/cloud-commitments-go/pkg/ladder" "github.com/LeanerCloud/cloud-commitments-platform/internal/config" "github.com/LeanerCloud/cloud-commitments-platform/internal/database" + "github.com/LeanerCloud/cloud-commitments-platform/internal/mocks" "github.com/jackc/pgx/v5" ) @@ -363,7 +364,7 @@ func (m *mockConfigStoreForHealth) ClearCollectionStarted(_ context.Context, _ s return nil } func (m *mockConfigStoreForHealth) ClaimCollectionInvocation(_ context.Context, _ string) (string, bool, error) { - return "mock-invocation-token", true, nil + return mocks.MockInvocationToken, true, nil } func (m *mockConfigStoreForHealth) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) { return nil, nil