reliar-store-postgres 0.7.0

//! `acquire`: the single-statement `FOR UPDATE SKIP LOCKED` claim (ADR 0006) and its best-effort
//! poison sweep (ADR 0039 §4).

use reliar_core::Serializer;
use reliar_outbox::{AcquireRequest, AcquiredBatch, PoisonedRow};

use super::error::PostgresOutboxError;
use crate::records::{RawRow, decode_row};

use super::PostgresOutboxStore;

/// [`reliar_outbox::OutboxStore::acquire`]'s body — see that trait method's rustdoc for the
/// full contract.
pub(super) async fn acquire<Ser: Serializer + Send + Sync + 'static>(
    store: &PostgresOutboxStore<Ser>,
    request: AcquireRequest,
) -> Result<AcquiredBatch, PostgresOutboxError> {
    let batch_size = i64::from(request.batch_size);
    let lease_ms = i64::try_from(request.lease.as_millis()).unwrap_or(i64::MAX);
    let worker = request.worker.as_str();

    // `Duration::ZERO` (the default) issues nothing and runs the claim as the single
    // implicit-transaction statement ADR 0006 relies on; a non-zero `statement_timeout`
    // costs a `BEGIN`/`SET LOCAL`/statement/`COMMIT` round trip instead, which is why it is
    // opt-in.
    let rows = if store.settings.statement_timeout.is_zero() {
        claim_rows(&store.pool, batch_size, worker, lease_ms)
            .await
            .map_err(|e| store.map_err(e))?
    } else {
        let mut tx = store.pool.begin().await.map_err(|e| store.map_err(e))?;
        let timeout_ms = i64::try_from(store.settings.statement_timeout.as_millis())
            .unwrap_or(i64::MAX)
            .to_string();

        sqlx::query_scalar!(
            "SELECT set_config('statement_timeout', $1, true)",
            timeout_ms
        )
        .fetch_one(&mut *tx)
        .await
        .map_err(|e| store.map_err(e))?;
        let rows = claim_rows(&mut *tx, batch_size, worker, lease_ms)
            .await
            .map_err(|e| store.map_err(e))?;
        tx.commit().await.map_err(|e| store.map_err(e))?;

        rows
    };

    let mut records = Vec::with_capacity(rows.len());
    let mut poisoned = Vec::new();
    let mut poisoned_ids = Vec::new();
    let mut poisoned_errors = Vec::new();

    for raw in rows {
        match decode_row(raw) {
            Ok(record) => records.push(record),
            Err(err) => {
                poisoned_ids.push(err.id.as_uuid());
                poisoned_errors.push(crate::records::truncate_last_error(err.detail.clone()));

                poisoned.push(PoisonedRow::new(err.id, err.message_id, err.detail));
            }
        }
    }

    if !poisoned_ids.is_empty() {
        // Not an observed publish attempt, so `attempts` is untouched (ADR 0009: `attempts`
        // counts outcomes, never claims) — only the lease clears and the row goes dead. Runs
        // under the same `statement_timeout` policy as the claim itself: it must run inside
        // the claim's own `SET LOCAL` wrap rather than directly on the pool, so a slow poison
        // sweep stays bounded by a non-zero `statement_timeout` too.
        //
        // **Best-effort (ADR 0039 §4): a sweep failure never turns a committed claim into an
        // `Err`.** The claim above has already committed and its rows are already leased to
        // this caller; failing the whole batch here would strand the N healthy rows for a
        // full lease over a problem with the poisoned ones. On failure this only logs — the
        // poisoned rows keep their lease and are re-attempted (sweep or publish) once it
        // lapses, so `poisoned` means "could not decode and an attempt was made to deaden",
        // not "is dead".
        let undecodable =
            crate::records::encode_dead_reason(reliar_outbox::DeadReason::Undecodable);
        let sweep_result = if store.settings.statement_timeout.is_zero() {
            poison_sweep_rows(
                &store.pool,
                &poisoned_ids,
                &poisoned_errors,
                worker,
                undecodable,
            )
            .await
        } else {
            async {
                let mut tx = store.pool.begin().await?;

                store.set_local_timeout_raw(&mut tx).await?;
                poison_sweep_rows(
                    &mut *tx,
                    &poisoned_ids,
                    &poisoned_errors,
                    worker,
                    undecodable,
                )
                .await?;

                tx.commit().await
            }
            .await
        };

        if let Err(err) = sweep_result {
            // Plain snake_case fields, not the usual dotted `worker.id`/`poisoned.count`
            // convention: `tracing`'s event macro hits a `macro_rules!` parsing
            // ambiguity ("multiple parsing options: built-in NTs tt ('field') or 1 other
            // option") when an explicit `target:` is followed by a dotted field path — a
            // `tracing` macro limitation, not a style choice.
            tracing::warn!(
                target: "reliar.outbox.acquire",
                worker_id = %worker,
                poisoned_count = poisoned_ids.len(),
                error = %store.map_err(err),
                "poison sweep failed; the claimed batch is returned and the undecodable rows \
                 stay leased until their lease lapses"
            );
        }
    }

    Ok(AcquiredBatch::new(records, poisoned))
}

/// `acquire`'s poison sweep: moves every row `decode_row` couldn't reconstruct to dead with
/// `DeadReason::Undecodable`, worker-guarded the same way every other outcome update is
/// (`o.locked_by = $3`) so a row already reclaimed by a different worker after this one's lease
/// lapsed is left alone.
async fn poison_sweep_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    poisoned_ids: &[uuid::Uuid],
    poisoned_errors: &[String],
    worker: &str,
    undecodable: &str,
) -> Result<(), sqlx::Error> {
    sqlx::query!(
        r#"UPDATE outbox o
              SET dead_at      = now(),
                  dead_reason  = $4,
                  last_error   = f.err,
                  locked_by    = NULL,
                  locked_until = NULL,
                  updated_at   = now()
             FROM UNNEST($1::uuid[], $2::text[]) AS f(id, err)
            WHERE o.id = f.id AND o.locked_by = $3"#,
        poisoned_ids,
        poisoned_errors,
        worker,
        undecodable,
    )
    .execute(executor)
    .await?;

    Ok(())
}

/// The canonical single-statement claim (ADR 0006): a CTE
/// `SELECT … FOR UPDATE SKIP LOCKED` feeding an `UPDATE … RETURNING`, so the row lock is
/// released before the call returns and no network I/O to a publisher can ever happen while it
/// is held. Named against [`RawRow`] via `query_as!` (never `FromRow`) so both the plain-pool
/// and `statement_timeout`-wrapped-transaction call sites in [`acquire`] share one macro
/// invocation instead of two structurally distinct anonymous row types.
///
/// **`available_at` moves to the lease end, alongside `locked_until` (ADR 0040 §1).**
/// `available_at` already means "the next time this row may be claimed" — retry backoff writes
/// exactly that — so a lease is the same statement about the same thing, not a new meaning. The
/// consequence: the claim's own boundary condition (`available_at <= now()`, in the `SELECT`
/// above) now *stops the index scan* at the first leased row instead of walking past every leased
/// row still ahead of it in `(available_at, sequence)` order, which is what made a claim's cost
/// scale with how much was in flight rather than with the batch it returned. The
/// `(locked_until IS NULL OR locked_until < now())` predicate stays: it is redundant for a row
/// this version wrote (`available_at = locked_until` while leased) but keeps the claim correct
/// for a row leased before the upgrade, at no cost once the scan already stops early.
async fn claim_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    batch_size: i64,
    worker: &str,
    lease_ms: i64,
) -> Result<Vec<RawRow>, sqlx::Error> {
    sqlx::query_as!(
        RawRow,
        r#"WITH claimed AS (
               SELECT id FROM outbox
                WHERE published_at IS NULL AND dead_at IS NULL
                  AND available_at <= now()
                  AND (locked_until IS NULL OR locked_until < now())
                  AND (expires_at IS NULL OR expires_at > now())
                ORDER BY available_at, sequence
                LIMIT $1
                FOR UPDATE SKIP LOCKED
           )
           UPDATE outbox o
              SET locked_by    = $2,
                  locked_until = now() + ($3::bigint * interval '1 millisecond'),
                  available_at = now() + ($3::bigint * interval '1 millisecond'),
                  updated_at   = now()
             FROM claimed
            WHERE o.id = claimed.id
           RETURNING o.id, o.message_id, o.message_type, o.message_version,
                     o.correlation_id, o.conversation_id, o.causation_id, o.request_id,
                     o.content_type, o.payload, o.tenant_id, o.expires_at, o.ordering_key,
                     o.metadata, o.headers, o.metadata_version,
                     o.created_at, o.available_at,
                     o.attempts, o.locked_by, o.locked_until,
                     o.published_at, o.dead_at, o.dead_reason, o.last_error"#,
        batch_size,
        worker,
        lease_ms,
    )
    .fetch_all(executor)
    .await
}