fix(revoke): make the Azure in-flight flip an exclusive claim so only one revoke posts a Return

internal/api/handler_purchases_revoke.go

@@ -200,15 +200,13 @@ func (h *Handler) loadAndRevokePurchaseHistory(ctx context.Context, req *events.
 
 	// Partial-success reconciliation (issue #290 Finding #6): if the
 	// revocation_in_flight flag is set but revoked_at is still NULL, Azure
-	// already issued the refund but our DB write failed. Return 207 so the
-	// frontend does not retry the Azure call (which would fail with "already
-	// returned"). The finalize_revocations sweep will reconcile the row.
+	// already issued the refund but our DB write failed, or another request
+	// holds the claim and is posting the Return now (issue #1518). Return 207
+	// so the frontend does not retry the Azure call (which would fail with
+	// "already returned"). The finalize_revocations sweep will reconcile the
+	// row.
 	if record.RevocationInFlight {
-		return &revokeReconcilePendingResult{
-			Code:          "RECONCILE_PENDING",
-			AzureReturned: true,
-			Message:       "Refund already issued. We will reconcile your audit record shortly. Do not retry.",
-		}, nil
+		return revocationInFlightResult(), nil
 	}
 
 	// Parse the confirmed quote from the request body. Azure revocations
@@ -787,8 +785,8 @@ func (h *Handler) callAzureReturn(
 		return nil, quoteErr
 	}
 
-	if revoked := h.flipRevocationInFlight(ctx, record); revoked != nil {
-		return revoked, nil
+	if stop, claimErr := h.claimRevocationInFlight(ctx, record); stop != nil || claimErr != nil {
+		return stop, claimErr
 	}
 
 	// Step 2: Return (post the actual refund request).
@@ -810,30 +808,58 @@ func (h *Handler) callAzureReturn(
 	return h.persistAzureRevocation(ctx, record, calcRefundAmount, calcRefundCurrency)
 }
 
-// flipRevocationInFlight sets the in-flight flag BEFORE the Azure Return
+// revocationInFlightResult is the answer for a row whose revocation another
+// request has in flight: the RECONCILE_PENDING response, which tells the
+// frontend not to retry.
+func revocationInFlightResult() *revokeReconcilePendingResult {
+	return &revokeReconcilePendingResult{
+		Code:          "RECONCILE_PENDING",
+		AzureReturned: true,
+		Message:       "Refund already issued. We will reconcile your audit record shortly. Do not retry.",
+	}
+}
+
+// claimRevocationInFlight sets the in-flight flag BEFORE the Azure Return
 // (partial-success guard, issue #290 Finding #6) so that if the subsequent
 // MarkPurchaseRevoked write fails, the row is visible to the
-// finalize_revocations sweep rather than silently stuck. Best-effort: if the
-// flip itself fails, log and continue — the in-flight flag is a safety net,
-// not a hard precondition.
+// finalize_revocations sweep rather than silently stuck.
 //
-// The one exception is a row the flip found already revoked (issue #316): a
-// concurrent request finished the revocation after this one loaded the row,
-// so no second Return is posted and the caller answers already_revoked, as
-// the up-front idempotency check would have. The result is nil otherwise.
-func (h *Handler) flipRevocationInFlight(ctx context.Context, record *config.PurchaseHistoryRecord) *revokePurchaseResult {
+// The flip is also the claim that lets exactly one request post the Return
+// (issue #1518). Both results are nil only when this request won it; every
+// other outcome stops before the Return:
+//   - the row was revoked by a concurrent request after this one loaded it
+//     (issue #316): already_revoked, as the up-front idempotency check would
+//     have answered;
+//   - another request holds the claim: RECONCILE_PENDING, as the up-front
+//     in-flight check would have answered;
+//   - the flip failed: an error (500). No Return has been posted, so a retry
+//     is safe, whereas posting one without the claim could race the holder's
+//     and, if this request's MarkPurchaseRevoked then failed too, leave a
+//     refund that no in-flight flag points the sweep at.
+//
+// Because only the claim holder reaches the Return, the ClearRevocationInFlight
+// in handleAzureReturnError only ever clears this request's own flag.
+func (h *Handler) claimRevocationInFlight(ctx context.Context, record *config.PurchaseHistoryRecord) (any, error) {
 	flipErr := h.config.FlipPurchaseRevocationInFlight(ctx, record.PurchaseID)
-	if flipErr == nil {
-		return nil
-	}
-	if !errors.Is(flipErr, config.ErrPurchaseAlreadyRevoked) {
-		logging.Warnf("revoke azure: FlipPurchaseRevocationInFlight for %s failed (continuing): %v", record.PurchaseID, flipErr)
-		return nil
+	switch {
+	case flipErr == nil:
+		return nil, nil
+	case errors.Is(flipErr, config.ErrRevocationInFlight):
+		logging.Infof("revoke azure: purchase_id=%s is already being revoked by a concurrent request; not posting a second Return", record.PurchaseID)
+		return revocationInFlightResult(), nil
+	case errors.Is(flipErr, config.ErrPurchaseAlreadyRevoked):
+		logging.Infof("revoke azure: purchase_id=%s was revoked by a concurrent request; not posting a second Return", record.PurchaseID)
+		return h.alreadyRevokedResult(ctx, record.PurchaseID), nil
+	default:
+		return nil, fmt.Errorf("revoke azure: claim revocation of %s (no Return posted): %w", record.PurchaseID, flipErr)
 	}
-	logging.Infof("revoke azure: purchase_id=%s was revoked by a concurrent request; not posting a second Return", record.PurchaseID)
+}
+
+// alreadyRevokedResult answers already_revoked for a row a concurrent request
+// revoked. Best-effort: the re-read stamp only fills in the response.
+func (h *Handler) alreadyRevokedResult(ctx context.Context, purchaseID string) *revokePurchaseResult {
 	result := &revokePurchaseResult{Status: "already_revoked"}
-	// Best-effort: the stamp only fills in the response.
-	if current, err := h.config.GetPurchaseHistoryByPurchaseID(ctx, record.PurchaseID); err == nil && current != nil && current.RevokedAt != nil {
+	if current, err := h.config.GetPurchaseHistoryByPurchaseID(ctx, purchaseID); err == nil && current != nil && current.RevokedAt != nil {
 		result.RevokedAt = current.RevokedAt.Format(time.RFC3339)
 		result.RevokedVia = current.RevokedVia
 	}
@@ -896,7 +922,8 @@ func (h *Handler) azureCalculateRefund(ctx context.Context, calcClient azureCalc
 
 // handleAzureReturnError clears the in-flight flag (no refund was issued) and
 // maps the Return error to the right status: 422 on the 7-day window edge, 400
-// on other client errors, 500 otherwise.
+// on other client errors, 500 otherwise. The caller holds the claim from
+// claimRevocationInFlight, so the flag it clears is its own (issue #1518).
 func (h *Handler) handleAzureReturnError(ctx context.Context, record *config.PurchaseHistoryRecord, err error) error {
 	// Azure Return failed. Clear the in-flight flag so the row is not left in
 	// a permanently sticky state that would mislead the finalize_revocations

internal/api/handler_purchases_revoke_claim_integration_test.go

@@ -0,0 +1,179 @@
+//go:build integration
+
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	armreservations "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/reservations/armreservations"
+	"github.com/LeanerCloud/cloud-commitments-platform/internal/config"
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// revokeRequestKey tags each concurrent revoke's context with its index, so
+// the stub Return client knows which request is calling it.
+type revokeRequestKey struct{}
+
+// barrierCalcRefundClient answers CalculateRefund with a fixed quote, but
+// only once every party has called it. Both revokes have loaded the row (not
+// revoked, not in flight) by then, so they reach the in-flight flip together,
+// as two first-time revokes would.
+type barrierCalcRefundClient struct {
+	price   *armreservations.Price
+	parties int32
+	arrived atomic.Int32
+	ready   chan struct{}
+}
+
+func (c *barrierCalcRefundClient) Post(ctx context.Context, _ string, _ armreservations.CalculateRefundRequest, _ *armreservations.CalculateRefundClientPostOptions) (armreservations.CalculateRefundClientPostResponse, error) {
+	if c.arrived.Add(1) == c.parties {
+		close(c.ready)
+	}
+	select {
+	case <-c.ready:
+	case <-ctx.Done():
+		return armreservations.CalculateRefundClientPostResponse{}, ctx.Err()
+	}
+	return armreservations.CalculateRefundClientPostResponse{
+		CalculateRefundResponse: armreservations.CalculateRefundResponse{
+			Properties: &armreservations.RefundResponseProperties{
+				SessionID:           toPtr("s-claim-race"),
+				BillingRefundAmount: c.price,
+			},
+		},
+	}, nil
+}
+
+// raceReturnClient stands in for Azure's Return API. The first Return
+// succeeds, but only after the other request has finished, so the other
+// request runs its whole path while the first one's refund is in flight.
+// Azure refuses a second Return of the same reservation, so any later call
+// gets a 400. While the first Return is in flight, it records whether the row
+// still carries the in-flight flag.
+type raceReturnClient struct {
+	store      *config.PostgresStore
+	purchaseID string
+	done       [2]chan struct{}
+	calls      atomic.Int32
+	// flagHeld is whether the row was still flagged in flight once the other
+	// request had finished, while the first Return was still in flight.
+	flagHeld atomic.Bool
+}
+
+func (c *raceReturnClient) Post(ctx context.Context, _ string, _ armreservations.RefundRequest, _ *armreservations.ReturnClientPostOptions) (armreservations.ReturnClientPostResponse, error) {
+	if c.calls.Add(1) > 1 {
+		return armreservations.ReturnClientPostResponse{}, &azcore.ResponseError{StatusCode: 400, ErrorCode: "ReservationAlreadyReturned"}
+	}
+	me, _ := ctx.Value(revokeRequestKey{}).(int)
+	select {
+	case <-c.done[1-me]:
+	case <-time.After(10 * time.Second):
+		// The other request never finished; the assertions report it.
+	case <-ctx.Done():
+		return armreservations.ReturnClientPostResponse{}, ctx.Err()
+	}
+	row, err := c.store.GetPurchaseHistoryByPurchaseID(context.Background(), c.purchaseID)
+	if err == nil && row != nil {
+		c.flagHeld.Store(row.RevocationInFlight)
+	}
+	return armreservations.ReturnClientPostResponse{}, nil
+}
+
+// TestRevokeAzure_ConcurrentFirstRevokes_OneReturn covers issue #1518 on real
+// Postgres. Two first-time revokes of the same Azure purchase load the row
+// (unrevoked, not in flight) and reach the in-flight flip together. Before
+// the fix the flip let an already-flagged row be flipped again, so both
+// posted a Return; Azure refused the second, and that request's error
+// handling cleared the flag while the first request's refund was still in
+// flight. Had the first request's MarkPurchaseRevoked then failed, the row
+// would have been left neither revoked nor flagged, out of the
+// finalize_revocations sweep's reach.
+//
+// Now exactly one request posts a Return and records the revocation, the
+// other answers RECONCILE_PENDING without posting one, and the flag stays set
+// until the holder's MarkPurchaseRevoked clears it.
+func TestRevokeAzure_ConcurrentFirstRevokes_OneReturn(t *testing.T) {
+	_, store := setupSubmitDB(t)
+	const rounds = 5
+	for round := range rounds {
+		t.Run(fmt.Sprintf("round %d", round), func(t *testing.T) {
+			runConcurrentFirstRevokes(t, store, round)
+		})
+	}
+}
+
+func runConcurrentFirstRevokes(t *testing.T, store *config.PostgresStore, round int) {
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	record := withinWindowRecord(fmt.Sprintf(
+		"/providers/Microsoft.Capacity/reservationOrders/order-1518-%d/reservations/res-1518-%d", round, round))
+	record.AccountID = fmt.Sprintf("acct-1518-%d", round)
+	record.Region = "eastus"
+	record.ResourceType = "Standard_D2s_v3"
+	require.NoError(t, store.SavePurchaseHistory(ctx, record))
+
+	calc := &barrierCalcRefundClient{
+		price:   &armreservations.Price{Amount: toPtr(10.0), CurrencyCode: toPtr("USD")},
+		parties: 2,
+		ready:   make(chan struct{}),
+	}
+	ret := &raceReturnClient{store: store, purchaseID: record.PurchaseID, done: [2]chan struct{}{make(chan struct{}), make(chan struct{})}}
+	h := &Handler{config: store, azureRevokeFactory: &azureRevokeClientFactory{
+		newCredential:            func() (azcore.TokenCredential, error) { return nil, nil },
+		newCalculateRefundClient: func(azcore.TokenCredential) (azureCalculateRefundClient, error) { return calc, nil },
+		newReturnClient:          func(azcore.TokenCredential) (azureReturnClient, error) { return ret, nil },
+	}}
+	body, err := json.Marshal(confirmedQuote(10, "USD"))
+	require.NoError(t, err)
+	req := &events.LambdaFunctionURLRequest{Body: string(body)}
+
+	results := make([]any, 2)
+	errs := make([]error, 2)
+	var wg sync.WaitGroup
+	for i := range 2 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer close(ret.done[i])
+			reqCtx := context.WithValue(ctx, revokeRequestKey{}, i)
+			results[i], errs[i] = h.loadAndRevokePurchaseHistory(reqCtx, req, revokeAdminSession(), record.PurchaseID)
+		}()
+	}
+	wg.Wait()
+
+	assert.Equal(t, int32(1), ret.calls.Load(), "exactly one of two concurrent first revokes may post an Azure Return")
+	assert.True(t, ret.flagHeld.Load(), "the other request must leave the in-flight flag set while the first one's Return is in flight")
+
+	revoked, inFlight := 0, 0
+	for i := range 2 {
+		require.NoError(t, errs[i], "request %d", i)
+		switch res := results[i].(type) {
+		case *revokePurchaseResult:
+			assert.Equal(t, "revoked", res.Status, "request %d", i)
+			revoked++
+		case *revokeReconcilePendingResult:
+			assert.Equal(t, revocationInFlightResult(), res, "request %d", i)
+			inFlight++
+		default:
+			t.Errorf("request %d: unexpected result %T %+v", i, results[i], results[i])
+		}
+	}
+	assert.Equal(t, 1, revoked, "one request revokes")
+	assert.Equal(t, 1, inFlight, "the other answers in flight (RECONCILE_PENDING)")
+
+	got, err := store.GetPurchaseHistoryByPurchaseID(ctx, record.PurchaseID)
+	require.NoError(t, err)
+	require.NotNil(t, got)
+	assert.NotNil(t, got.RevokedAt, "the winner's MarkPurchaseRevoked must land")
+	assert.False(t, got.RevocationInFlight, "MarkPurchaseRevoked clears the holder's flag")
+}

internal/api/handler_purchases_revoke_guard_test.go

@@ -183,7 +183,9 @@ func confirmedQuote(amount float64, currency string) revokeConfirmBody {
 // flip is now a compare-and-set on revoked_at IS NULL. When it reports the
 // row already revoked, a concurrent request finished the revocation after
 // this one loaded the row, so no second Return is posted and the answer is
-// already_revoked. Any other flip failure stays best-effort.
+// already_revoked. Since issue #1518 the flip is also the claim that lets one
+// request post the Return: a row another request holds answers
+// RECONCILE_PENDING, and any other flip failure stops before the Return.
 func TestCallAzureReturn_FlipFindsRowRevoked(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
@@ -215,21 +217,52 @@ func TestCallAzureReturn_FlipFindsRowRevoked(t *testing.T) {
 		mockStore.AssertNotCalled(t, "MarkPurchaseRevoked", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	})
 
-	t.Run("other flip error: still returns and marks", func(t *testing.T) {
+	// Issue #1518: another request holds the claim. This one must neither
+	// post a Return nor touch the flag (clearing it would orphan the holder's
+	// refund from the finalize sweep), and answers RECONCILE_PENDING like the
+	// up-front in-flight check.
+	t.Run("in flight: no Return, no clear, RECONCILE_PENDING", func(t *testing.T) {
 		t.Parallel()
 		mockStore := new(MockConfigStore)
 		t.Cleanup(func() { mockStore.AssertExpectations(t) })
 		r := armReservationRecord()
-		mockStore.On("FlipPurchaseRevocationInFlight", ctx, r.PurchaseID).Return(fmt.Errorf("FlipPurchaseRevocationInFlight: connection reset"))
-		mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
+		mockStore.On("FlipPurchaseRevocationInFlight", ctx, r.PurchaseID).
+			Return(fmt.Errorf("FlipPurchaseRevocationInFlight: purchase_id %q: %w", r.PurchaseID, config.ErrRevocationInFlight))
+		// Registered only so a regression fails on the assertions below
+		// rather than on an unexpected-call panic.
+		mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
 		returnClient := &stubReturnClient{}
 
 		h := &Handler{config: mockStore}
 		result, err := h.callAzureReturn(ctx, &countingCalcRefundClient{price: quote}, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
 		require.NoError(t, err)
-		assert.Equal(t, 1, returnClient.calls)
-		res, ok := result.(*revokePurchaseResult)
-		require.True(t, ok, "got %T", result)
-		assert.Equal(t, "revoked", res.Status)
+		assert.Zero(t, returnClient.calls, "the request that loses the claim must not post a Return")
+		assert.Equal(t, revocationInFlightResult(), result)
+		mockStore.AssertNotCalled(t, "ClearRevocationInFlight", mock.Anything, mock.Anything)
+		mockStore.AssertNotCalled(t, "MarkPurchaseRevoked", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
+	})
+
+	// Issue #1518: a request that could not establish the claim does not know
+	// whether another request holds it, so it posts no Return (a retry is
+	// safe) and leaves the flag alone.
+	t.Run("other flip error: no Return, no clear, 500", func(t *testing.T) {
+		t.Parallel()
+		mockStore := new(MockConfigStore)
+		t.Cleanup(func() { mockStore.AssertExpectations(t) })
+		r := armReservationRecord()
+		mockStore.On("FlipPurchaseRevocationInFlight", ctx, r.PurchaseID).Return(fmt.Errorf("FlipPurchaseRevocationInFlight: connection reset"))
+		mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+		returnClient := &stubReturnClient{}
+
+		h := &Handler{config: mockStore}
+		result, err := h.callAzureReturn(ctx, &countingCalcRefundClient{price: quote}, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+		require.Error(t, err)
+		assert.Nil(t, result)
+		assert.ErrorContains(t, err, "connection reset")
+		_, isClientErr := IsClientError(err)
+		assert.False(t, isClientErr, "a store failure is a 500, not a client error")
+		assert.Zero(t, returnClient.calls, "no Return without the claim")
+		mockStore.AssertNotCalled(t, "ClearRevocationInFlight", mock.Anything, mock.Anything)
+		mockStore.AssertNotCalled(t, "MarkPurchaseRevoked", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	})
 }

internal/api/openapi.yaml

@@ -1480,7 +1480,9 @@ paths:
         one-hour safety margin, and the body must echo the quote from GET
         /api/purchases/{id}/revoke/calculate. An already revoked row answers
         "already_revoked"; a row whose Azure refund was issued but not yet
-        recorded answers code RECONCILE_PENDING and must not be retried.
+        recorded, or is being posted by a concurrent revoke of the same
+        purchase, answers code RECONCILE_PENDING and must not be retried.
+        Only one revoke at a time posts the Azure return.
       parameters:
         - $ref: '#/components/parameters/CSRFToken'
       requestBody:
@@ -8341,7 +8343,8 @@ components:
         One of three shapes. A canceled scheduled execution: status
         "canceled" and message. A revoked purchase history row: status
         "revoked" or "already_revoked" with revoked_at and revoked_via. An
-        Azure refund that was issued but could not be recorded yet: code
+        Azure refund that was issued but could not be recorded yet, or that a
+        concurrent revoke of the same purchase is posting: code
         RECONCILE_PENDING, azure_returned and message (do not retry).
       properties:
         status: