236 lines
11 KiB
Go
236 lines
11 KiB
Go
package pgstore
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
// idempotency.go: the storage half of `Idempotency-Key`.
|
|
|
|
// ErrKeyReused is the same key presented for a DIFFERENT request. Told apart from a replay because
|
|
// the remedies differ: one is "your retry worked", the other is "your client has a bug".
|
|
var ErrKeyReused = errors.New("pgstore: this idempotency key was used for another request")
|
|
|
|
// ErrKeyInFlight is a repeat that arrived while the first attempt is still running. Waiting is the
|
|
// remedy, which is why the answer carries Retry-After.
|
|
var ErrKeyInFlight = errors.New("pgstore: the first request under this idempotency key is still running")
|
|
|
|
// IdempotencyKey identifies one claim: (principal, method, path), as the canon ratified — the same
|
|
// key on another operation is ANOTHER key, not a conflict.
|
|
//
|
|
// The path is KEYED by its digest and stored beside it: an address is unbounded, and an over-long
|
|
// one overflowed the index tuple, so the claim failed and the caller was answered 500 where the
|
|
// contract promises 409.
|
|
type IdempotencyKey struct {
|
|
UserID string
|
|
Method string
|
|
Path string
|
|
Key string
|
|
Fingerprint []byte
|
|
}
|
|
|
|
// pathHash is what the index is keyed by. Derived rather than carried, so it cannot disagree with
|
|
// the path beside it.
|
|
func (k IdempotencyKey) pathHash() []byte {
|
|
sum := sha256.Sum256([]byte(k.Path))
|
|
return sum[:]
|
|
}
|
|
|
|
// IdempotentResponse is a completed first attempt, replayed verbatim.
|
|
type IdempotentResponse struct {
|
|
Status int
|
|
Location string
|
|
Body []byte
|
|
// ContentSHA256 is the digest of the body that attempt accepted, for a request whose identity its
|
|
// FINGERPRINT cannot settle: a multipart upload declares its parts and not its bytes, so two
|
|
// different books can arrive under one key with the same declaration. A caller holding one of
|
|
// these may not replay until it has read the repeat's body and found the same digest.
|
|
//
|
|
// Absent where the fingerprint IS the whole request — a JSON write is compared whole.
|
|
ContentSHA256 []byte
|
|
}
|
|
|
|
// ClaimStale is how long a claim may be outstanding before another attempt may take it over; without
|
|
// it a dead process locks the key until the retention sweep.
|
|
//
|
|
// ⚠ Must exceed the longest write here — an upload, bounded by its route's read deadline. Enforced
|
|
// at boot (config.Load): under a longer deadline a retry takes the claim from a LIVE attempt and the
|
|
// user gets the second book the header exists to prevent.
|
|
const ClaimStale = 30 * time.Minute
|
|
|
|
// ClaimToken says WHICH attempt holds a claim. The key identifies the request; this identifies the
|
|
// try, and without it every writer addressed the row by identity alone — so a late attempt released
|
|
// its successor's claim, and a superseded one wrote its own receipt over the row.
|
|
type ClaimToken string
|
|
|
|
// NewClaimToken mints one. It is an owner marker rather than a secret: what matters is that two
|
|
// attempts never share one.
|
|
func NewClaimToken() ClaimToken { return ClaimToken(newID("clm")) }
|
|
|
|
// errKeyVanished is the row disappearing between an attempt's insert and its re-read. Not an answer
|
|
// to anybody: the key is simply free again, and the claim is taken from the top.
|
|
var errKeyVanished = errors.New("pgstore: the idempotency key was released mid-claim")
|
|
|
|
// claimRounds bounds that retry. Two attempts of one key can hand it back and forth at most as fast
|
|
// as they can fail, and a caller that loses three rounds is meeting something other than a race.
|
|
const claimRounds = 3
|
|
|
|
// ClaimIdempotency takes the right to perform a request once.
|
|
//
|
|
// It answers one of four things, and the four are the whole semantics: a token to go ahead with, a
|
|
// stored response to replay, ErrKeyInFlight while the first attempt is still running, ErrKeyReused
|
|
// when the key was used for a different request. The token is what the caller must present to
|
|
// complete or release: it is the proof that the claim is still this attempt's.
|
|
func (s *Store) ClaimIdempotency(ctx context.Context, k IdempotencyKey, now time.Time) (*IdempotentResponse, ClaimToken, error) {
|
|
// ⚠ Retried, because losing a race is not one of the four answers. Two attempts arrive together;
|
|
// one wins the insert and then FAILS FAST — a 4xx gives the key straight back — and the loser's
|
|
// re-read, on a fresh READ COMMITTED snapshot, finds nothing at all. Answered rather than retried
|
|
// that was a 500 on precisely the case this header exists for, which is the same 500 the
|
|
// `do nothing` below was written to close. Reproduced on a live database.
|
|
for range claimRounds {
|
|
out, token, err := s.claim(ctx, k, now)
|
|
if errors.Is(err, errKeyVanished) {
|
|
continue
|
|
}
|
|
return out, token, err
|
|
}
|
|
return nil, "", fmt.Errorf("pgstore: the idempotency key was claimed and released under us %d times", claimRounds)
|
|
}
|
|
|
|
func (s *Store) claim(ctx context.Context, k IdempotencyKey, now time.Time) (*IdempotentResponse, ClaimToken, error) {
|
|
var out *IdempotentResponse
|
|
token := NewClaimToken()
|
|
err := s.inTx(ctx, func(tx pgx.Tx) error {
|
|
var fingerprint, content []byte
|
|
var status *int
|
|
var location string
|
|
var body []byte
|
|
var claimedAt time.Time
|
|
var finishedAt *time.Time
|
|
const read = `
|
|
select request_sha256, status, location, body, claimed_at, finished_at, content_sha256
|
|
from idempotency_keys
|
|
where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4 for update`
|
|
err := tx.QueryRow(ctx, read, k.UserID, k.Method, k.pathHash(), k.Key).
|
|
Scan(&fingerprint, &status, &location, &body, &claimedAt, &finishedAt, &content)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
// ⚠ `for update` on a missing row locks nothing, so two simultaneous first attempts both
|
|
// land here. `do nothing` plus a re-read decides between them; a bare insert made the loser
|
|
// raise a unique violation — a 500 on the one case this header exists for.
|
|
tag, err := tx.Exec(ctx, `
|
|
insert into idempotency_keys (user_id, method, path, path_sha256, key, request_sha256, claimed_at, claim_token)
|
|
values ($1, $2, $3, $4, $5, $6, $7, $8)
|
|
on conflict (user_id, method, path_sha256, key) do nothing`,
|
|
k.UserID, k.Method, k.Path, k.pathHash(), k.Key, k.Fingerprint, now, token)
|
|
if err != nil {
|
|
return fmt.Errorf("pgstore: claim idempotency key: %w", err)
|
|
}
|
|
if tag.RowsAffected() == 1 {
|
|
return nil
|
|
}
|
|
if err := tx.QueryRow(ctx, read, k.UserID, k.Method, k.pathHash(), k.Key).
|
|
Scan(&fingerprint, &status, &location, &body, &claimedAt, &finishedAt, &content); err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return errKeyVanished // the winner released it; the key is free and this attempt starts over
|
|
}
|
|
return fmt.Errorf("pgstore: re-read idempotency key: %w", err)
|
|
}
|
|
} else if err != nil {
|
|
return fmt.Errorf("pgstore: read idempotency key: %w", err)
|
|
}
|
|
if !bytes.Equal(fingerprint, k.Fingerprint) {
|
|
return ErrKeyReused
|
|
}
|
|
if finishedAt != nil && status != nil {
|
|
out = &IdempotentResponse{Status: *status, Location: location, Body: body, ContentSHA256: content}
|
|
return nil
|
|
}
|
|
if now.Sub(claimedAt) < ClaimStale {
|
|
return ErrKeyInFlight
|
|
}
|
|
// The first attempt is gone: its process died, or it answered an error and released nothing.
|
|
// Taking the claim over is what keeps a retry from meeting `key_in_flight` forever, and the NEW
|
|
// token is what stops the old attempt — which may still be alive and merely slow — from
|
|
// completing or releasing what is no longer its claim.
|
|
if _, err := tx.Exec(ctx, `
|
|
update idempotency_keys set claimed_at = $5, request_sha256 = $6, claim_token = $7
|
|
where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4`,
|
|
k.UserID, k.Method, k.pathHash(), k.Key, now, k.Fingerprint, token); err != nil {
|
|
return fmt.Errorf("pgstore: re-claim idempotency key: %w", err)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, "", err
|
|
}
|
|
if out != nil {
|
|
return out, "", nil // a replay holds no claim
|
|
}
|
|
return nil, token, nil
|
|
}
|
|
|
|
// ErrClaimLost is a completion whose row is gone or already finished. The work IS done, but nothing
|
|
// will replay it — a fact worth reporting rather than losing in a nil.
|
|
var ErrClaimLost = errors.New("pgstore: this attempt no longer holds the idempotency key")
|
|
|
|
// CompleteIdempotency records what the first attempt answered, so a repeat can be given the same
|
|
// thing without doing the work again.
|
|
//
|
|
// Only the attempt that still HOLDS the claim may write the receipt. An attempt whose claim was
|
|
// taken over has had its work done by somebody else; letting it answer would replay a `Location` for
|
|
// a book that was never created.
|
|
func (s *Store) CompleteIdempotency(ctx context.Context, k IdempotencyKey, token ClaimToken, resp IdempotentResponse, now time.Time) error {
|
|
tag, err := s.pool.Exec(ctx, `
|
|
update idempotency_keys set status = $5, location = $6, body = $7, finished_at = $8,
|
|
content_sha256 = $10
|
|
where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4
|
|
and claim_token = $9 and finished_at is null`,
|
|
k.UserID, k.Method, k.pathHash(), k.Key, resp.Status, resp.Location, resp.Body, now, token,
|
|
resp.ContentSHA256)
|
|
if err != nil {
|
|
return fmt.Errorf("pgstore: complete idempotency key: %w", err)
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return ErrClaimLost
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ReleaseIdempotency gives a key back after an attempt that did NOT complete.
|
|
//
|
|
// A failed attempt is not a completed one, and the contract names the case that matters: a `408` is
|
|
// not a completed attempt, so the same key may be presented again and the retry is a repeat of the
|
|
// first call rather than a second book.
|
|
//
|
|
// On the token for the same reason as the completion, and this is the half that cost a duplicate
|
|
// book: a slow attempt that finally failed deleted the row of the successor which had taken the key
|
|
// over and was already doing the work.
|
|
func (s *Store) ReleaseIdempotency(ctx context.Context, k IdempotencyKey, token ClaimToken) error {
|
|
tag, err := s.pool.Exec(ctx, `
|
|
delete from idempotency_keys
|
|
where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4
|
|
and claim_token = $5 and finished_at is null`,
|
|
k.UserID, k.Method, k.pathHash(), k.Key, token)
|
|
if err != nil {
|
|
return fmt.Errorf("pgstore: release idempotency key: %w", err)
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return ErrClaimLost
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SweepIdempotency forgets records past the window the contract promises ("at least 24 hours").
|
|
func (s *Store) SweepIdempotency(ctx context.Context, before time.Time) (int64, error) {
|
|
tag, err := s.pool.Exec(ctx, `delete from idempotency_keys where claimed_at < $1`, before)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("pgstore: sweep idempotency keys: %w", err)
|
|
}
|
|
return tag.RowsAffected(), nil
|
|
}
|