fix(scheduler): scope the collection marker release to the invocation, not the marker

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"
@@ -515,6 +516,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 mocks.MockInvocationToken, true, nil
+}
 func (m *mockConfigStore) GetRIUtilizationCache(_ context.Context, _ string, _ int) (*config.RIUtilizationCacheEntry, error) {
 	return nil, nil
 }

internal/config/interfaces.go

@@ -587,6 +587,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

internal/config/store_postgres_collection_invocation_test.go

@@ -0,0 +1,262 @@
+//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/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()
+	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)
+		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()
+	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
+	// 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)
+		}
+	})
+}

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

internal/mocks/stores.go

@@ -21,6 +21,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
@@ -2064,6 +2069,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)

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