PGX Serialization Retry Transaction

package main

import (
	"context"
	"errors"
	"fmt"
	"time"

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

func isSerializationFailure(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && pgErr.Code == "40001"
}

func runWithRetry(ctx context.Context, pool *pgxpool.Pool, attempts int) (int, error) {
	var balance int

	for attempt := 1; attempt <= attempts; attempt++ {
		err := pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error {
			row := tx.QueryRow(ctx,
				"select balance from accounts where id = 1 for update")
			if err := row.Scan(&balance); err != nil {
				return fmt.Errorf("lock read: %w", err)
			}

			if balance < 100 {
				return errors.New("insufficient funds")
			}

			_, err := tx.Exec(ctx,
				"update accounts set balance = balance - 100 where id = 1")
			return err
		})

		if err == nil {
			return balance - 100, nil
		}
		if isSerializationFailure(err) && attempt < attempts {
			time.Sleep(time.Duration(attempt*50) * time.Millisecond)
			continue // 序列化冲突：整块重放
		}
		return 0, err
	}
	return 0, errors.New("unreachable")
}

func main() {
	config, _ := pgxpool.ParseConfig("postgres://user:pass@localhost/app")
	pool, err := pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer pool.Close()

	final, err := runWithRetry(context.Background(), pool, 4)
	fmt.Println("final balance:", final, "err:", err)
}