fix(revoke): record an accepted Azure Return even if the request is canceled right after it

internal/api/handler_purchases_revoke.go

@@ -33,6 +33,7 @@ import (
 	"github.com/LeanerCloud/cloud-commitments-go/pkg/logging"
 	"github.com/LeanerCloud/cloud-commitments-platform/internal/auth"
 	"github.com/LeanerCloud/cloud-commitments-platform/internal/config"
+	"github.com/LeanerCloud/cloud-commitments-platform/internal/ctxutil"
 	"github.com/aws/aws-lambda-go/events"
 	"github.com/jackc/pgx/v5"
 )
@@ -125,6 +126,16 @@ var revokeMarkRetryBackoffs = []time.Duration{
 	9 * time.Second,
 }
 
+// revokePersistTimeout bounds the writes that record an Azure Return Azure has
+// already accepted (issue #1525). They run detached from the request, so a
+// client disconnect or upstream timeout right after the Return does not lose
+// the record. 30s covers the whole retry schedule (13s of backoff across
+// revokeMarkRetryBackoffs) plus four attempts of two single-row UPDATEs, while
+// still capping how long a hung database can hold the handler after its
+// caller has gone. If it runs out, the result is the same 207 as exhausted
+// retries.
+const revokePersistTimeout = 30 * time.Second
+
 // revokePurchase handles POST /api/purchases/{purchaseId}/revoke.
 //
 // Authorization: session required + revoke-own:purchases (or revoke-any for
@@ -954,7 +965,14 @@ func (h *Handler) handleAzureReturnError(ctx context.Context, record *config.Pur
 // Each attempt first records that Azure accepted the Return
 // (MarkRevocationReturnPosted, until one lands): the finalize_revocations
 // sweep only finalizes rows with that record (issue #1523).
-func (h *Handler) persistAzureRevocation(ctx context.Context, record *config.PurchaseHistoryRecord, calcRefundAmount *float64, calcRefundCurrency string) (any, error) {
+//
+// Azure has already refunded by the time this runs, so the writes ignore the
+// request's cancellation and run under their own revokePersistTimeout (issue
+// #1525): a request canceled right after the Return still records it. Nothing
+// before the Return is detached.
+func (h *Handler) persistAzureRevocation(reqCtx context.Context, record *config.PurchaseHistoryRecord, calcRefundAmount *float64, calcRefundCurrency string) (any, error) {
+	ctx, cancel := ctxutil.DetachedWithTimeout(reqCtx, revokePersistTimeout)
+	defer cancel()
 	now := time.Now().UTC()
 	returnRecorded := false
 	markErr := h.recordAzureReturnAndMark(ctx, record.PurchaseID, now, calcRefundAmount, calcRefundCurrency, &returnRecorded)

internal/api/handler_purchases_revoke_claim_integration_test.go

@@ -177,3 +177,64 @@ func runConcurrentFirstRevokes(t *testing.T, store *config.PostgresStore, round
 	assert.NotNil(t, got.RevokedAt, "the winner's MarkPurchaseRevoked must land")
 	assert.False(t, got.RevocationInFlight, "MarkPurchaseRevoked clears the holder's flag")
 }
+
+// TestRevokeAzure_RequestCanceledAfterReturn_RecordsAndRevokes covers issue
+// #1525 on real Postgres. The request is canceled right after Azure accepts
+// the Return (a client disconnect). Before the fix every write that follows
+// ran under the canceled request context, so pgx refused each attempt: the
+// Return went unrecorded, the row was never revoked, and the
+// finalize_revocations sweep would release it after 2h as if nothing had been
+// refunded. Now those writes run detached, so the row ends up with the Return
+// recorded and revoked.
+func TestRevokeAzure_RequestCanceledAfterReturn_RecordsAndRevokes(t *testing.T) {
+	// NOT parallel: mutates the package-global revokeMarkRetryBackoffs, so a
+	// regression fails fast instead of sleeping through the retries.
+	orig := revokeMarkRetryBackoffs
+	revokeMarkRetryBackoffs = []time.Duration{0, 0, 0}
+	t.Cleanup(func() { revokeMarkRetryBackoffs = orig })
+
+	pg, store := setupSubmitDB(t)
+	ctx, cancelTest := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancelTest()
+
+	record := withinWindowRecord("/providers/Microsoft.Capacity/reservationOrders/order-1525/reservations/res-1525")
+	record.AccountID = "acct-1525"
+	record.Region = "eastus"
+	record.ResourceType = "Standard_D2s_v3"
+	require.NoError(t, store.SavePurchaseHistory(ctx, record))
+
+	reqCtx, cancelReq := context.WithCancel(ctx)
+	defer cancelReq()
+	calc := &barrierCalcRefundClient{
+		price:   &armreservations.Price{Amount: toPtr(10.0), CurrencyCode: toPtr("USD")},
+		parties: 1,
+		ready:   make(chan struct{}),
+	}
+	ret := &cancelAfterReturnClient{cancel: cancelReq}
+	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)}
+
+	result, err := h.loadAndRevokePurchaseHistory(reqCtx, req, revokeAdminSession(), record.PurchaseID)
+	require.NoError(t, err)
+	require.Equal(t, 1, ret.calls, "exactly one Return")
+	require.ErrorIs(t, reqCtx.Err(), context.Canceled, "the request was canceled after the Return")
+	res, ok := result.(*revokePurchaseResult)
+	require.True(t, ok, "want revoked, got %T %+v", result, result)
+	assert.Equal(t, "revoked", res.Status)
+
+	var returnPostedAt, revokedAt *time.Time
+	var inFlight bool
+	require.NoError(t, pg.DB.Pool().QueryRow(ctx, `
+		SELECT return_posted_at, revoked_at, revocation_in_flight
+		  FROM purchase_history WHERE purchase_id = $1`, record.PurchaseID).
+		Scan(&returnPostedAt, &revokedAt, &inFlight))
+	assert.NotNil(t, returnPostedAt, "the accepted Return must be recorded")
+	assert.NotNil(t, revokedAt, "the purchase must be marked revoked")
+	assert.False(t, inFlight, "MarkPurchaseRevoked clears the in-flight flag")
+}

internal/api/handler_purchases_revoke_guard_test.go

@@ -133,7 +133,7 @@ func TestRevokePurchase_AzureRefundGuard(t *testing.T) {
 			mockStore.On("GetExecutionByID", ctx, r.PurchaseID).Return(nil, fmt.Errorf("%w: execution %s", config.ErrNotFound, r.PurchaseID))
 			mockStore.On("GetPurchaseHistoryByPurchaseID", ctx, r.PurchaseID).Return(r, nil)
 			mockStore.On("FlipPurchaseRevocationInFlight", ctx, r.PurchaseID).Return(nil).Maybe()
-			mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil).Maybe()
+			mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil).Maybe()
 
 			calcClient := &countingCalcRefundClient{price: tc.quote}
 			returnClient := &stubReturnClient{}

internal/api/handler_purchases_revoke_return_posted_test.go

@@ -46,8 +46,8 @@ func TestCallAzureReturn_RecordsReturnBeforeMark(t *testing.T) {
 
 	r := armReservationRecord()
 	order := &callOrder{}
-	mockStore.On("MarkRevocationReturnPosted", ctx, r.PurchaseID).Return(nil).Run(order.add("return")).Once()
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"),
+	mockStore.On("MarkRevocationReturnPosted", mock.Anything, r.PurchaseID).Return(nil).Run(order.add("return")).Once()
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"),
 		"direct-api", "", mock.Anything, mock.Anything).Return(nil).Run(order.add("mark")).Once()
 
 	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-ok"}
@@ -93,9 +93,9 @@ func TestCallAzureReturn_RetriesReturnRecordUntilItLands(t *testing.T) {
 
 	r := armReservationRecord()
 	dbDown := errors.New("db down")
-	mockStore.On("MarkRevocationReturnPosted", ctx, r.PurchaseID).Return(dbDown).Twice()
-	mockStore.On("MarkRevocationReturnPosted", ctx, r.PurchaseID).Return(nil).Once()
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"),
+	mockStore.On("MarkRevocationReturnPosted", mock.Anything, r.PurchaseID).Return(dbDown).Twice()
+	mockStore.On("MarkRevocationReturnPosted", mock.Anything, r.PurchaseID).Return(nil).Once()
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"),
 		"direct-api", "", mock.Anything, mock.Anything).Return(dbDown).Times(4)
 
 	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-retry"}
@@ -109,3 +109,152 @@ func TestCallAzureReturn_RetriesReturnRecordUntilItLands(t *testing.T) {
 	mockStore.AssertNumberOfCalls(t, "MarkRevocationReturnPosted", 3)
 	mockStore.AssertNumberOfCalls(t, "MarkPurchaseRevoked", 4)
 }
+
+// Issue #1525: once Azure has accepted the Return, the request being canceled
+// (client disconnect, an upstream timeout) must not stop the handler from
+// recording it. The writes before the Return must still honor cancellation.
+
+// writeCtxState is the state of a write's context when the write ran.
+type writeCtxState struct {
+	err       error
+	bounded   bool
+	remaining time.Duration
+}
+
+// ctxHonoringStore fails the revoke writes with ctx.Err() once ctx is done, as
+// pgx does, and records the state of the context each write ran under.
+type ctxHonoringStore struct {
+	*MockConfigStore
+	mu     sync.Mutex
+	writes []writeCtxState
+}
+
+func (s *ctxHonoringStore) seen(ctx context.Context) error {
+	deadline, bounded := ctx.Deadline()
+	s.mu.Lock()
+	s.writes = append(s.writes, writeCtxState{err: ctx.Err(), bounded: bounded, remaining: time.Until(deadline)})
+	s.mu.Unlock()
+	return ctx.Err()
+}
+
+func (s *ctxHonoringStore) seenWrites() []writeCtxState {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return append([]writeCtxState(nil), s.writes...)
+}
+
+func (s *ctxHonoringStore) FlipPurchaseRevocationInFlight(ctx context.Context, purchaseID string) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	return s.MockConfigStore.FlipPurchaseRevocationInFlight(ctx, purchaseID)
+}
+
+func (s *ctxHonoringStore) MarkRevocationReturnPosted(ctx context.Context, purchaseID string) error {
+	if err := s.seen(ctx); err != nil {
+		return err
+	}
+	return s.MockConfigStore.MarkRevocationReturnPosted(ctx, purchaseID)
+}
+
+func (s *ctxHonoringStore) MarkPurchaseRevoked(ctx context.Context, purchaseID string, revokedAt time.Time, revokedVia, supportCaseID string, calcRefundAmount *float64, calcRefundCurrency string) error {
+	if err := s.seen(ctx); err != nil {
+		return err
+	}
+	return s.MockConfigStore.MarkPurchaseRevoked(ctx, purchaseID, revokedAt, revokedVia, supportCaseID, calcRefundAmount, calcRefundCurrency)
+}
+
+// cancelAfterReturnClient accepts the Return, then cancels the request, as a
+// client disconnecting right after Azure answered would. It records whether
+// the Return itself ran under the request's (cancelable) context.
+type cancelAfterReturnClient struct {
+	cancel        context.CancelFunc
+	calls         int
+	sawRequestCtx bool
+}
+
+func (c *cancelAfterReturnClient) Post(ctx context.Context, _ string, _ armreservations.RefundRequest, _ *armreservations.ReturnClientPostOptions) (armreservations.ReturnClientPostResponse, error) {
+	c.calls++
+	before := ctx.Err()
+	c.cancel()
+	c.sawRequestCtx = before == nil && ctx.Err() != nil
+	return armreservations.ReturnClientPostResponse{}, nil
+}
+
+func TestCallAzureReturn_RequestCanceledAfterReturn_StillRecordsAndMarks(t *testing.T) {
+	// NOT parallel: mutates the package-global revokeMarkRetryBackoffs, so a
+	// regression fails in milliseconds instead of sleeping through the retries.
+	orig := revokeMarkRetryBackoffs
+	revokeMarkRetryBackoffs = []time.Duration{0, 0, 0}
+	t.Cleanup(func() { revokeMarkRetryBackoffs = orig })
+
+	reqCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(t) })
+	store := &ctxHonoringStore{MockConfigStore: mockStore}
+
+	r := armReservationRecord()
+	mockStore.On("MarkRevocationReturnPosted", mock.Anything, r.PurchaseID).Return(nil).Once()
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"),
+		"direct-api", "", mock.Anything, mock.Anything).Return(nil).Once()
+
+	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-cancel"}
+	returnClient := &cancelAfterReturnClient{cancel: cancel}
+	h := &Handler{config: store}
+	result, err := h.callAzureReturn(reqCtx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+	require.NoError(t, err)
+
+	require.Equal(t, 1, returnClient.calls)
+	assert.True(t, returnClient.sawRequestCtx, "the Return itself must run under the request's context")
+	require.ErrorIs(t, reqCtx.Err(), context.Canceled)
+	rr, ok := result.(*revokePurchaseResult)
+	require.True(t, ok, "a canceled request must still record the accepted Return and revoke, got %T %+v", result, result)
+	assert.Equal(t, "revoked", rr.Status)
+
+	writes := store.seenWrites()
+	require.Len(t, writes, 2, "one record, one mark")
+	for i, w := range writes {
+		require.NoError(t, w.err, "write %d ran under a canceled context", i)
+		require.True(t, w.bounded, "write %d must have its own bounded timeout", i)
+		assert.LessOrEqual(t, w.remaining, revokePersistTimeout, "write %d", i)
+	}
+}
+
+// cancelingCalcRefundClient quotes a refund, then cancels the request before
+// the handler gets to the claim and the Return.
+type cancelingCalcRefundClient struct {
+	stubCalcRefundClientWithAmount
+	cancel context.CancelFunc
+}
+
+func (c *cancelingCalcRefundClient) Post(ctx context.Context, orderID string, body armreservations.CalculateRefundRequest, opts *armreservations.CalculateRefundClientPostOptions) (armreservations.CalculateRefundClientPostResponse, error) {
+	resp, err := c.stubCalcRefundClientWithAmount.Post(ctx, orderID, body, opts)
+	c.cancel()
+	return resp, err
+}
+
+// A request canceled before the Return must stop there: no write before the
+// Return ignores cancellation, so no claim lands and no Return is posted.
+func TestCallAzureReturn_RequestCanceledBeforeReturn_PostsNothing(t *testing.T) {
+	t.Parallel()
+	reqCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	mockStore := new(MockConfigStore)
+	store := &ctxHonoringStore{MockConfigStore: mockStore}
+
+	r := armReservationRecord()
+	calcClient := &cancelingCalcRefundClient{
+		stubCalcRefundClientWithAmount: stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-early"},
+		cancel:                         cancel,
+	}
+	returnClient := &stubReturnClient{}
+	h := &Handler{config: store}
+	_, err := h.callAzureReturn(reqCtx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Zero(t, returnClient.calls, "no Return may be posted for a canceled request")
+	assert.Empty(t, store.seenWrites())
+	mockStore.AssertNotCalled(t, "FlipPurchaseRevocationInFlight")
+	mockStore.AssertNotCalled(t, "MarkRevocationReturnPosted")
+	mockStore.AssertNotCalled(t, "MarkPurchaseRevoked")
+}

internal/api/handler_purchases_revoke_test.go

@@ -251,7 +251,7 @@ func TestRevokePurchase_AzureSuccess(t *testing.T) {
 	t.Cleanup(func() { mockStore.AssertExpectations(t) })
 
 	r := armReservationRecord()
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
 
 	sessID := "test-session"
 	calcClient := &stubCalcRefundClient{
@@ -1172,7 +1172,7 @@ func TestCallAzureReturn_TOCTOUWithinEpsilonSucceeds(t *testing.T) {
 	// User confirmed $100.00; Azure now quotes $100.005 (within $0.01).
 	userConfirmed := 100.0
 	calcClient := &stubCalcRefundClientWithAmount{amount: 100.005, currency: "USD", sessID: "s-1"}
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
 
 	h := &Handler{config: mockStore}
@@ -1197,7 +1197,7 @@ func TestCallAzureReturn_AuditRowPopulatedWithQuote(t *testing.T) {
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
 
 	// Assert MarkPurchaseRevoked receives the quote amount and currency.
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "",
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "",
 		mock.MatchedBy(func(v *float64) bool { return v != nil && *v == 42.50 }),
 		"EUR",
 	).Return(nil)
@@ -1273,7 +1273,7 @@ func TestCallAzureReturn_MarkPurchaseRevokedFailAllRetries(t *testing.T) {
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
 
 	// MarkPurchaseRevoked is called 1 + len(backoffs) = 4 times total (initial + 3 retries).
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).
 		Return(errors.New("db down")).Times(4)
 
 	h := &Handler{config: mockStore}
@@ -1382,7 +1382,7 @@ func TestCallAzureReturn_JustOutsideSafetyMargin(t *testing.T) {
 
 	calcClient := &stubCalcRefundClientWithAmount{amount: 50.0, currency: "USD", sessID: "s-margin"}
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
 
 	h := &Handler{config: mockStore}
 	result, err := h.callAzureReturn(ctx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(50, "USD"))
@@ -1550,7 +1550,7 @@ func TestCallAzureReturn_Success_DoesNotClearInFlight(t *testing.T) {
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
 
 	r := armReservationRecord()
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"),
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"),
 		"direct-api", "", mock.Anything, mock.Anything).Return(nil).Once()
 
 	h := &Handler{config: mockStore}
@@ -1752,7 +1752,7 @@ func TestRevokeAzurePurchase_Success(t *testing.T) {
 	t.Cleanup(func() { mockStore.AssertExpectations(t) })
 
 	r := armReservationRecord()
-	mockStore.On("MarkPurchaseRevoked", ctx, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
+	mockStore.On("MarkPurchaseRevoked", mock.Anything, r.PurchaseID, mock.AnythingOfType("time.Time"), "direct-api", "", mock.Anything, mock.Anything).Return(nil)
 
 	calcClient := &stubCalcRefundClientWithAmount{amount: 12.34, currency: "USD", sessID: "s-revoke"}
 	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}