SQLX Transaction Retry Serialization

use sqlx::postgres::PgDatabaseError;
use sqlx::{PgPool, Row};
use std::time::Duration;

async fn transfer(
    pool: &PgPool,
    from_id: i64,
    to_id: i64,
    amount: i64,
) -> Result<i64, sqlx::Error> {
    let mut attempts = 0;

    loop {
        attempts += 1;
        match transfer_once(pool, from_id, to_id, amount).await {
            Ok(balance) => return Ok(balance),
            Err(err) => {
                if is_serialization_failure(&err) && attempts < 4 {
                    let backoff = 50 * attempts;
                    tokio::time::sleep(Duration::from_millis(backoff)).await;
                    continue; // 序列化冲突：整块重放
                }
                return Err(err);
            }
        }
    }
}

fn is_serialization_failure(err: &sqlx::Error) -> bool {
    match err {
        sqlx::Error::Database(db) => db
            .downcast_ref::<PgDatabaseError>()
            .code()
            .starts_with("40"),
        _ => false,
    }
}

async fn transfer_once(
    pool: &PgPool,
    from_id: i64,
    to_id: i64,
    amount: i64,
) -> Result<i64, sqlx::Error> {
    let mut tx = pool.begin().await?;

    let balance: i64 = sqlx::query("SELECT balance FROM accounts WHERE id = $1 FOR UPDATE")
        .bind(from_id)
        .fetch_one(&mut *tx)
        .await?
        .get(0);

    if balance < amount {
        return Err(sqlx::Error::RowNotFound); // 余额不足：业务性失败不重试
    }

    sqlx::query("UPDATE accounts SET balance = balance - $1 WHERE id = $2")
        .bind(amount)
        .bind(from_id)
        .execute(&mut *tx)
        .await?;

    sqlx::query("UPDATE accounts SET balance = balance + $1 WHERE id = $2")
        .bind(amount)
        .bind(to_id)
        .execute(&mut *tx)
        .await?;

    tx.commit().await?;
    Ok(balance - amount)
}

#[tokio::main]
async fn main() -> Result<(), sqlx::Error> {
    let pool = PgPool::connect("postgres://user:pass@localhost/app").await?;
    let final_balance = transfer(&pool, 1, 2, 500).await?;
    println!("source balance now: {final_balance}");
    pool.close().await;
    Ok(())
}