fix(revoke): finalize only revocations whose Azure Return was recorded; release stale flags

internal/api/handler_purchases_revoke.go

@@ -950,25 +950,31 @@ func (h *Handler) handleAzureReturnError(ctx context.Context, record *config.Pur
 // backoff retries. If every attempt fails, Azure has already refunded but the
 // DB write could not land, so it returns a 207 RECONCILE_PENDING result (no
 // retry) for the finalize_revocations sweep to reconcile, rather than a 500.
+//
+// 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) {
 	now := time.Now().UTC()
-	markErr := h.config.MarkPurchaseRevoked(ctx, record.PurchaseID, now, "direct-api", "", calcRefundAmount, calcRefundCurrency)
+	returnRecorded := false
+	markErr := h.recordAzureReturnAndMark(ctx, record.PurchaseID, now, calcRefundAmount, calcRefundCurrency, &returnRecorded)
 	for attempt, backoff := range revokeMarkRetryBackoffs {
 		if markErr == nil {
 			break
 		}
 		logging.Warnf("revoke azure: MarkPurchaseRevoked attempt %d failed for %s: %v (retrying in %s)",
 			attempt+1, record.PurchaseID, markErr, backoff)
 		time.Sleep(backoff)
-		markErr = h.config.MarkPurchaseRevoked(ctx, record.PurchaseID, now, "direct-api", "", calcRefundAmount, calcRefundCurrency)
+		markErr = h.recordAzureReturnAndMark(ctx, record.PurchaseID, now, calcRefundAmount, calcRefundCurrency, &returnRecorded)
 	}
 	if markErr != nil {
 		// All retries failed. Azure has already refunded but we cannot persist
 		// the revocation state. Return 207 Multi-Status so the frontend knows
 		// the refund issued but not to retry — the finalize_revocations sweep
-		// will reconcile the DB write on its next tick.
-		logging.Errorf("revoke azure: MarkPurchaseRevoked failed for %s after %d attempts (Azure already returned): %v",
-			record.PurchaseID, len(revokeMarkRetryBackoffs)+1, markErr)
+		// will reconcile the DB write on its next tick, provided the Return
+		// was recorded.
+		logging.Errorf("revoke azure: MarkPurchaseRevoked failed for %s after %d attempts (Azure already returned, return recorded=%t): %v",
+			record.PurchaseID, len(revokeMarkRetryBackoffs)+1, returnRecorded, markErr)
 		return &revokeReconcilePendingResult{
 			Code:          "RECONCILE_PENDING",
 			AzureReturned: true,
@@ -986,6 +992,20 @@ func (h *Handler) persistAzureRevocation(ctx context.Context, record *config.Pur
 	}, nil
 }
 
+// recordAzureReturnAndMark is one persist attempt after Azure accepted the
+// Return: stamp return_posted_at unless an earlier attempt already did (a
+// failure is logged; the mark may still land), then MarkPurchaseRevoked.
+func (h *Handler) recordAzureReturnAndMark(ctx context.Context, purchaseID string, now time.Time, calcRefundAmount *float64, calcRefundCurrency string, returnRecorded *bool) error {
+	if !*returnRecorded {
+		if err := h.config.MarkRevocationReturnPosted(ctx, purchaseID); err != nil {
+			logging.Warnf("revoke azure: MarkRevocationReturnPosted for %s failed (continuing to the mark): %v", purchaseID, err)
+		} else {
+			*returnRecorded = true
+		}
+	}
+	return h.config.MarkPurchaseRevoked(ctx, purchaseID, now, "direct-api", "", calcRefundAmount, calcRefundCurrency)
+}
+
 // parseAzureReservationIDs extracts the reservation order ID and reservation ID
 // from an Azure ARM resource path. The purchase_id is stored as the ARM
 // resource ID at purchase time.

internal/api/handler_purchases_revoke_return_posted_test.go

@@ -0,0 +1,111 @@
+package api
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	armreservations "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/reservations/armreservations"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+// Issue #1523: the finalize_revocations sweep only finalizes rows whose Return
+// was recorded as accepted, so the handler must record it after Azure accepts
+// the Return and before (or while retrying) MarkPurchaseRevoked. Stub Azure
+// clients and a mock store; nothing here calls Azure.
+
+// callOrder records the store calls in the order the handler makes them.
+type callOrder struct {
+	mu    sync.Mutex
+	calls []string
+}
+
+func (o *callOrder) add(name string) func(mock.Arguments) {
+	return func(mock.Arguments) {
+		o.mu.Lock()
+		defer o.mu.Unlock()
+		o.calls = append(o.calls, name)
+	}
+}
+
+func (o *callOrder) get() []string {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	return append([]string(nil), o.calls...)
+}
+
+func TestCallAzureReturn_RecordsReturnBeforeMark(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(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"),
+		"direct-api", "", mock.Anything, mock.Anything).Return(nil).Run(order.add("mark")).Once()
+
+	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-ok"}
+	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
+	h := &Handler{config: mockStore}
+	result, err := h.callAzureReturn(ctx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+	require.NoError(t, err)
+	rr, ok := result.(*revokePurchaseResult)
+	require.True(t, ok)
+	assert.Equal(t, "revoked", rr.Status)
+	assert.Equal(t, []string{"return", "mark"}, order.get(), "the accepted Return must be recorded before the mark")
+}
+
+// A Return Azure refused must not be recorded: the sweep would otherwise
+// finalize a purchase that was never refunded.
+func TestCallAzureReturn_ReturnErrorRecordsNothing(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	mockStore := new(MockConfigStore)
+
+	r := armReservationRecord()
+	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-fail"}
+	returnClient := &stubReturnClient{err: errors.New("dial tcp: connection refused")}
+	h := &Handler{config: mockStore}
+	_, err := h.callAzureReturn(ctx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+	require.Error(t, err)
+	mockStore.AssertNotCalled(t, "MarkRevocationReturnPosted")
+	mockStore.AssertNotCalled(t, "MarkPurchaseRevoked")
+}
+
+// When the stamp and the mark both fail, every retry retries the stamp too,
+// until it lands; after that only the mark is retried. If the stamp lands and
+// the mark never does, the 207 is honest: the sweep will finalize the row.
+func TestCallAzureReturn_RetriesReturnRecordUntilItLands(t *testing.T) {
+	// NOT parallel: mutates the package-global revokeMarkRetryBackoffs.
+	orig := revokeMarkRetryBackoffs
+	revokeMarkRetryBackoffs = []time.Duration{0, 0, 0}
+	t.Cleanup(func() { revokeMarkRetryBackoffs = orig })
+
+	ctx := context.Background()
+	mockStore := new(MockConfigStore)
+	t.Cleanup(func() { mockStore.AssertExpectations(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"),
+		"direct-api", "", mock.Anything, mock.Anything).Return(dbDown).Times(4)
+
+	calcClient := &stubCalcRefundClientWithAmount{amount: 10.0, currency: "USD", sessID: "s-retry"}
+	returnClient := &stubReturnClient{resp: armreservations.ReturnClientPostResponse{}}
+	h := &Handler{config: mockStore}
+	result, err := h.callAzureReturn(ctx, calcClient, returnClient, r, "order-abc", "res-xyz", confirmedQuote(10, "USD"))
+	require.NoError(t, err)
+	pending, ok := result.(*revokeReconcilePendingResult)
+	require.True(t, ok)
+	assert.Equal(t, "RECONCILE_PENDING", pending.Code)
+	mockStore.AssertNumberOfCalls(t, "MarkRevocationReturnPosted", 3)
+	mockStore.AssertNumberOfCalls(t, "MarkPurchaseRevoked", 4)
+}

internal/api/nullable_purchase_cost_integration_test.go

@@ -112,7 +112,9 @@ func TestNullablePurchaseCostPostgres(t *testing.T) {
 			require.Len(t, points, 2)
 		}
 	}
-	_, err = pool.Exec(ctx, `UPDATE purchase_history SET revocation_in_flight=true WHERE purchase_id='unknown-0'`)
+	// The finalize queue reads return_posted_at (000148, issue #1523).
+	require.NoError(t, migrations.MigrateToVersion(ctx, pool, path, 148))
+	_, err = pool.Exec(ctx, `UPDATE purchase_history SET revocation_in_flight=true, return_posted_at=now() WHERE purchase_id='unknown-0'`)
 	require.NoError(t, err)
 	inFlight, err := store.GetPurchaseHistoryInFlight(ctx)
 	require.NoError(t, err)

internal/config/interfaces.go

@@ -439,12 +439,28 @@ type StoreInterface interface {
 	// a flag another request set (issue #1518).
 	ClearRevocationInFlight(ctx context.Context, purchaseID string) error
 
-	// GetPurchaseHistoryInFlight returns all purchase_history rows with
-	// revocation_in_flight=true and revoked_at IS NULL.  These are rows where
-	// the Azure Return call succeeded but MarkPurchaseRevoked failed; the
-	// finalize_revocations scheduled sweep calls this to retry the DB write.
+	// GetPurchaseHistoryInFlight returns the purchase_history rows whose
+	// Return was recorded by MarkRevocationReturnPosted but whose revoked_at
+	// is still NULL: the Azure Return call succeeded but MarkPurchaseRevoked
+	// failed. The finalize_revocations scheduled sweep calls this to retry the
+	// DB write. A flagged row with no recorded Return is not returned
+	// (issue #1523): no refund is known to have gone out for it.
 	GetPurchaseHistoryInFlight(ctx context.Context) ([]*PurchaseHistoryRecord, error)
 
+	// MarkRevocationReturnPosted records that the provider accepted the
+	// Return for purchaseID (return_posted_at, migration 000148). Called by
+	// the revoke handler after the Azure Return succeeds and before
+	// MarkPurchaseRevoked. The first stamp stands; a revoked row is a no-op.
+	// Returns a not-found error when no row matches (issue #1523).
+	MarkRevocationReturnPosted(ctx context.Context, purchaseID string) error
+
+	// ReleaseStaleRevocationsInFlight clears revocation_in_flight on unrevoked
+	// rows flagged for longer than after with no Return recorded, so the
+	// purchase can be revoked again, and returns their purchase IDs. A flagged
+	// row with no flip timestamp is stamped now and aged from then. Called by
+	// the finalize_revocations sweep (issue #1523).
+	ReleaseStaleRevocationsInFlight(ctx context.Context, after time.Duration) ([]string, error)
+
 	// UpdatePurchaseHistoryListing stamps the AWS marketplace listing_id and
 	// listing_state onto a purchase_history row. Called after
 	// CreateReservedInstancesListing succeeds (listing_state="active") and

internal/config/store_postgres.go

@@ -3060,9 +3060,12 @@ func (s *PostgresStore) MarkPurchaseRevoked(ctx context.Context, purchaseID stri
 // row, so exactly one flip of an unrevoked, unflagged row succeeds. A missing
 // row keeps the not-found error.
 func (s *PostgresStore) FlipPurchaseRevocationInFlight(ctx context.Context, purchaseID string) error {
+	// revocation_in_flight_at (migration 000148, issue #1523) is the database
+	// clock at the flip; the finalize_revocations sweep uses it to release a
+	// flag whose Return was never recorded.
 	tag, err := s.db.Exec(ctx, `
 		UPDATE purchase_history
-		   SET revocation_in_flight = true
+		   SET revocation_in_flight = true, revocation_in_flight_at = now()
 		 WHERE purchase_id = $1
 		   AND revoked_at IS NULL
 		   AND revocation_in_flight = false
@@ -3117,19 +3120,27 @@ func (s *PostgresStore) ClearRevocationInFlight(ctx context.Context, purchaseID
 	return nil
 }
 
-// GetPurchaseHistoryInFlight returns all purchase_history rows with
-// revocation_in_flight=true and revoked_at IS NULL. Used by the
+// GetPurchaseHistoryInFlight returns the purchase_history rows whose Azure
+// Return was recorded as accepted (return_posted_at set by
+// MarkRevocationReturnPosted) but whose revoked_at is still NULL. Used by the
 // finalize_revocations scheduled sweep to retry MarkPurchaseRevoked for rows
 // where the Azure Return succeeded but the subsequent DB write failed
 // (issue #290 Finding #6).
+//
+// revocation_in_flight alone is not enough (issue #1523): a row is also left
+// flagged when a request dies between the flip and the Return, and marking
+// that row revoked would record a refund that never happened. The flag is not
+// required either: once the Return is recorded the purchase is refunded, even
+// if a concurrent request's failed Return cleared the flag.
 func (s *PostgresStore) GetPurchaseHistoryInFlight(ctx context.Context) ([]*PurchaseHistoryRecord, error) {
 	rows, err := s.db.Query(ctx, `
 		SELECT account_id, purchase_id, timestamp, provider, service, region,
 		       resource_type, count, term, payment, upfront_cost, monthly_cost,
 		       estimated_savings, plan_id, plan_name, ramp_step, cloud_account_id,
-		       revocation_window_closes_at, revoked_at, revoked_via, support_case_id
+		       revocation_window_closes_at, revoked_at, revoked_via, support_case_id,
+		       revocation_in_flight
 		FROM purchase_history
-		WHERE revocation_in_flight = true
+		WHERE return_posted_at IS NOT NULL
 		  AND revoked_at IS NULL
 	`)
 	if err != nil {
@@ -3167,6 +3178,7 @@ func (s *PostgresStore) GetPurchaseHistoryInFlight(ctx context.Context) ([]*Purc
 			&revokedAt,
 			&revokedVia,
 			&supportCaseID,
+			&r.RevocationInFlight,
 		); err != nil {
 			return nil, fmt.Errorf("GetPurchaseHistoryInFlight scan: %w", err)
 		}
@@ -3189,7 +3201,6 @@ func (s *PostgresStore) GetPurchaseHistoryInFlight(ctx context.Context) ([]*Purc
 		if supportCaseID.Valid {
 			r.SupportCaseID = supportCaseID.String
 		}
-		r.RevocationInFlight = true
 		result = append(result, &r)
 	}
 	return result, rows.Err()

internal/config/store_postgres_revocation_return.go

@@ -0,0 +1,96 @@
+package config
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+// MarkRevocationReturnPosted records that the provider accepted the Return for
+// purchaseID (issue #1523, migration 000148). The revoke handler calls it right
+// after Azure accepts the Return and before MarkPurchaseRevoked, so the
+// finalize_revocations sweep can tell "refunded, audit write failed" (finalize)
+// from "flagged, but no Return ever went out" (never finalize).
+//
+// The first stamp stands (COALESCE), so a retry does not move it. A revoked
+// row is left alone and returns nil: there is nothing left to finalize. A
+// missing row returns a not-found error.
+func (s *PostgresStore) MarkRevocationReturnPosted(ctx context.Context, purchaseID string) error {
+	tag, err := s.db.Exec(ctx, `
+		UPDATE purchase_history
+		   SET return_posted_at = COALESCE(return_posted_at, now())
+		 WHERE purchase_id = $1
+		   AND revoked_at IS NULL
+	`, purchaseID)
+	if err != nil {
+		return fmt.Errorf("MarkRevocationReturnPosted: %w", err)
+	}
+	if tag.RowsAffected() > 0 {
+		return nil
+	}
+	var exists bool
+	if err := s.db.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM purchase_history WHERE purchase_id = $1)`, purchaseID).Scan(&exists); err != nil {
+		return fmt.Errorf("MarkRevocationReturnPosted existence check: %w", err)
+	}
+	if !exists {
+		return fmt.Errorf("MarkRevocationReturnPosted: purchase_id %q not found", purchaseID)
+	}
+	return nil
+}
+
+// ReleaseStaleRevocationsInFlight clears revocation_in_flight on unrevoked
+// rows that have been flagged for longer than after with no Return recorded
+// (issue #1523), and returns their purchase IDs. Such a row is one whose
+// request died between the flip and the Return, or whose Return was never
+// recorded as accepted. Clearing the flag lets the user revoke it again; if
+// Azure did take a Return that went unrecorded, Azure refuses the second one.
+//
+// A flagged row with no revocation_in_flight_at (flagged before migration 000148,
+// or by a binary that predates it) is stamped now and released once after has
+// passed from then, so its age is never overstated.
+//
+// Both statements use the database clock. Each UPDATE re-checks its predicate
+// on the locked row under READ COMMITTED, so a concurrent
+// MarkRevocationReturnPosted or a fresh flip that commits first makes the row
+// no longer match; one that commits after the release still records its
+// Return, which GetPurchaseHistoryInFlight finds with or without the flag.
+func (s *PostgresStore) ReleaseStaleRevocationsInFlight(ctx context.Context, after time.Duration) ([]string, error) {
+	if after <= 0 {
+		return nil, fmt.Errorf("ReleaseStaleRevocationsInFlight: timeout must be positive, got %s", after)
+	}
+	if _, err := s.db.Exec(ctx, `
+		UPDATE purchase_history
+		   SET revocation_in_flight_at = now()
+		 WHERE revocation_in_flight
+		   AND revoked_at IS NULL
+		   AND return_posted_at IS NULL
+		   AND revocation_in_flight_at IS NULL
+	`); err != nil {
+		return nil, fmt.Errorf("ReleaseStaleRevocationsInFlight stamp: %w", err)
+	}
+	rows, err := s.db.Query(ctx, `
+		UPDATE purchase_history
+		   SET revocation_in_flight = false, revocation_in_flight_at = NULL
+		 WHERE revocation_in_flight
+		   AND revoked_at IS NULL
+		   AND return_posted_at IS NULL
+		   AND revocation_in_flight_at < now() - make_interval(secs => $1)
+		RETURNING purchase_id
+	`, after.Seconds())
+	if err != nil {
+		return nil, fmt.Errorf("ReleaseStaleRevocationsInFlight: %w", err)
+	}
+	defer rows.Close()
+	var released []string
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			return nil, fmt.Errorf("ReleaseStaleRevocationsInFlight scan: %w", err)
+		}
+		released = append(released, id)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("ReleaseStaleRevocationsInFlight: %w", err)
+	}
+	return released, nil
+}