Crussell/backend/handlers/payments/locks.go

package payments

import (
	"context"
	"log"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// advisoryLockAttempts and advisoryLockRetryDelay bound the total time a
// try-lock retry loop waits for a contended advisory lock (~3s). This is the
// core defence against pool exhaustion: every payment handler pins a pool
// connection and would otherwise block on `pg_advisory_lock` for the FULL
// Square round-trip (up to ~30s) of whichever request holds the lock, so a
// handful of concurrent same-key requests can exhaust the whole pool
// (max(4, numCPU)) and hang the app. With a bounded try-lock loop the waiter
// gives up after ~3s and surfaces "operation in progress" instead of holding a
// pool connection hostage.
const (
	advisoryLockAttempts   = 30
	advisoryLockRetryDelay = 100 * time.Millisecond
)

// queryRower is satisfied by both *pgxpool.Conn (session-level locks) and
// pgx.Tx (transaction-scoped locks) so the try-lock helpers work on pinned
// pool connections and inside transactions alike.
type queryRower interface {
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

// acquireAdvisoryLock acquires a session advisory lock on the pinned conn with
// bounded retries (defends the pool against concurrent same-key lock waiters).
// Returns (true, nil) once the lock is held; (false, nil) if the lock could
// not be acquired within the bound — the caller must surface a 409/503
// "operation in progress, try again" instead of blocking. The unlock is still
// the caller's responsibility (pg_advisory_unlock on the same conn via defer).
func acquireAdvisoryLock(ctx context.Context, conn *pgxpool.Conn, key string) (bool, error) {
	return tryAdvisoryLock(ctx, conn, key, "pg_try_advisory_lock")
}

// releasePaymentLock releases a session advisory lock acquired by
// acquireAdvisoryLock on the SAME pinned pool connection (pg_advisory_unlock
// only releases locks held by the calling session). It is the generic release
// counterpart to acquireAdvisoryLock, and callers defer it immediately after a
// successful acquire so the unlock runs before the deferred pinConn.Release().
// Errors are logged and otherwise ignored — exactly what the inline
// `pg_advisory_unlock(hashtext('crussell:...:' || $1))` blocks it replaces did
// — and the key must be the FULL "crussell:..." string that was hashed at
// acquire time so the two hashtext() calls produce the same lock bigint.
func releasePaymentLock(pinConn *pgxpool.Conn, lockKey string) {
	if _, err := pinConn.Exec(context.Background(), `
		SELECT pg_advisory_unlock(hashtext($1))
	`, lockKey); err != nil {
		log.Printf("Failed to release payment serialization lock %s: %v", lockKey, err)
	}
}

// acquireAdvisoryXactLockBlocking is the transaction-scoped BLOCKING variant
// of acquireAdvisoryLock: it issues `SELECT pg_advisory_xact_lock(...)`
// ONCE and waits for as long as the key is contended — there is no 3s bound.
// The lock is transaction-scoped, so it is auto-released at the caller's
// commit/rollback (never unlocked explicitly).
//
// This stays deliberately blocking for ONE site only: the admin cancellation
// path (lockCancellationPayments). Admin cancellations are rare and there is
// only ONE such transaction at a time, so the pool-exhaustion rationale that
// justifies the bounded try-lock everywhere else does not apply here. A bound
// would be actively harmful: if a manual RefundPayment holds the same
// "crussell:refund:" key across its up-to-30s Square round-trip, a timed-out
// cancellation would abort, and the caller (manage.go) would commit the
// cancellation with ZERO refund rows created — no sweep retry is possible
// because the rows never existed, so the refund would be silently lost. The
// blocking acquire guarantees the cancellation refund runs, waiting for the
// manual refund to finish rather than dropping the money.
func acquireAdvisoryXactLockBlocking(ctx context.Context, tx pgx.Tx, key string) error {
	// Exec, not QueryRow.Scan: pg_advisory_xact_lock returns void, and Exec
	// discards the result set (the same pattern the payment handlers use for
	// their blocking pg_advisory_lock) while still blocking server-side until
	// the lock is granted.
	_, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtext($1))`, key)
	return err
}

// tryAdvisoryLock is the shared try-lock retry loop. fn is the Postgres
// advisory-lock function to call (pg_try_advisory_lock or
// pg_try_advisory_xact_lock). The key is passed to hashtext() so it is hashed
// to a bigint exactly like the blocking `pg_advisory_lock(hashtext($1))` calls
// it replaces — the two acquire the same locks.
func tryAdvisoryLock(ctx context.Context, q queryRower, key, fn string) (bool, error) {
	// ONE reusable timer for the whole retry loop instead of time.After per
	// attempt: a contended lock spins ~30 times, and allocating a fresh timer
	// (with its own goroutine) on every attempt is wasteful. The timer is
	// Stop+drained before each Reset so a previously-fired tick can never make
	// the next wait return early (missed-tick semantics). Timing is preserved:
	// 100ms between attempts, ~3s total bound.
	timer := time.NewTimer(advisoryLockRetryDelay)
	defer timer.Stop()
	for attempt := 0; attempt < advisoryLockAttempts; attempt++ { // ~3s total
		var acquired bool
		if err := q.QueryRow(ctx, `SELECT `+fn+`(hashtext($1))`, key).Scan(&acquired); err != nil {
			return false, err
		}
		if acquired {
			return true, nil
		}
		if !timer.Stop() {
			select {
			case <-timer.C:
			default:
			}
		}
		timer.Reset(advisoryLockRetryDelay)
		select {
		case <-timer.C:
		case <-ctx.Done():
			return false, ctx.Err()
		}
	}
	return false, nil
}