remora/crates/remora-durable/src/tx.rs

use std::future::Future;
use std::pin::Pin;

use sha2::{Digest, Sha256};
use sqlx::{PgPool, Postgres, Transaction};

use crate::error::DurableError;

pub type PgTx<'a> = Transaction<'a, Postgres>;
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;

pub async fn with_transaction<T, F>(pool: &PgPool, work: F) -> Result<T, DurableError>
where
    F: for<'c> FnOnce(&'c mut PgTx<'_>) -> BoxFuture<'c, Result<T, DurableError>>,
{
    let mut tx = pool.begin().await?;
    match work(&mut tx).await {
        Ok(value) => {
            tx.commit().await?;
            Ok(value)
        }
        Err(error) => {
            let _ = tx.rollback().await;
            Err(error)
        }
    }
}

/// Bound how long this transaction waits for row or advisory locks.
pub async fn set_lock_timeout(tx: &mut PgTx<'_>) -> Result<(), DurableError> {
    sqlx::query("SET LOCAL lock_timeout = '5s'")
        .execute(tx.as_mut())
        .await?;
    Ok(())
}

/// Advisory lock class reserved for mutation ledger identity locks.
/// Feature packages must not use this class.
pub const ADVISORY_CLASS_MUTATION: i32 = 1;

/// Advisory lock class for application and feature-package locks.
/// Prefer this (or higher package-owned classes) over the mutation class.
pub const ADVISORY_CLASS_APP: i32 = 2;

/// Mixes the class into a 64-bit key to separate lock domains without reducing
/// the PostgreSQL `bigint` identity space.
#[must_use]
pub fn advisory_key(class: i32, key: i32) -> i64 {
    let mut hasher = Sha256::new();
    hasher.update(b"remora.advisory.v1");
    hasher.update(class.to_be_bytes());
    hasher.update(key.to_be_bytes());
    let digest = hasher.finalize();
    i64::from_be_bytes([
        digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7],
    ])
}

/// Transaction-scoped advisory lock on a domain-separated 64-bit key.
pub async fn advisory_xact_lock(tx: &mut PgTx<'_>, key: i64) -> Result<(), DurableError> {
    sqlx::query("SELECT pg_advisory_xact_lock($1)")
        .bind(key)
        .execute(tx.as_mut())
        .await?;
    Ok(())
}