reliar-store-postgres 0.6.0

//! `complete`/`fail`/`release`/`extend_lease`: the worker-guarded outcome updates every claimed
//! row eventually gets.

use std::time::Duration;

use reliar_core::Serializer;
use reliar_outbox::{CompletedMessage, FailedMessage, FailureOutcome, MessageRef, WorkerId};

use crate::error::PostgresStoreError;

use super::PostgresOutboxStore;

/// [`reliar_outbox::OutboxStore::complete`]'s body.
pub(super) async fn complete<Ser: Serializer + Send + Sync + 'static>(
    store: &PostgresOutboxStore<Ser>,
    worker: &WorkerId,
    items: &[CompletedMessage],
) -> Result<u64, PostgresStoreError> {
    if items.is_empty() {
        return Ok(0);
    }

    let ids: Vec<uuid::Uuid> = items.iter().map(|i| i.message.id.as_uuid()).collect();
    let affected = if store.settings.statement_timeout.is_zero() {
        complete_rows(&store.pool, &ids, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?
    } else {
        let mut tx = store.pool.begin().await.map_err(|e| store.map_err(e))?;

        store.set_local_timeout(&mut tx).await?;
        let affected = complete_rows(&mut *tx, &ids, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?;
        tx.commit().await.map_err(|e| store.map_err(e))?;

        affected
    };
    log_shortfall("complete", items.len(), affected);

    Ok(affected)
}

/// [`reliar_outbox::OutboxStore::fail`]'s body.
pub(super) async fn fail<Ser: Serializer + Send + Sync + 'static>(
    store: &PostgresOutboxStore<Ser>,
    worker: &WorkerId,
    items: &[FailedMessage],
) -> Result<u64, PostgresStoreError> {
    if items.is_empty() {
        return Ok(0);
    }

    let mut retry_ids = Vec::new();
    let mut retry_errors = Vec::new();
    let mut retry_delays = Vec::new();
    let mut dead_ids = Vec::new();
    let mut dead_errors = Vec::new();
    let mut dead_reasons = Vec::new();

    for item in items {
        match item.outcome {
            FailureOutcome::Retry { delay } => {
                retry_ids.push(item.message.id.as_uuid());
                retry_errors.push(item.error.clone());

                retry_delays.push(i64::try_from(delay.as_millis()).unwrap_or(i64::MAX));
            }
            FailureOutcome::Dead { reason } => {
                dead_ids.push(item.message.id.as_uuid());
                dead_errors.push(item.error.clone());

                dead_reasons.push(crate::records::encode_dead_reason(reason));
            }
            // `FailureOutcome` is `#[non_exhaustive]` from another crate; a variant this
            // build does not know how to apply is left untouched rather than guessed at —
            // it stays claimed until its lease expires and is republished, the same benign
            // outcome as any other unresolved row (ADR 0008).
            _ => tracing::error!(
                id = %item.message.id,
                "unrecognised FailureOutcome variant; row left as-is"
            ),
        }
    }

    let affected = if store.settings.statement_timeout.is_zero() {
        let mut affected = 0u64;

        if !retry_ids.is_empty() {
            affected += fail_retry_rows(
                &store.pool,
                &retry_ids,
                &retry_errors,
                &retry_delays,
                worker.as_str(),
            )
            .await
            .map_err(|e| store.map_err(e))?;
        }

        if !dead_ids.is_empty() {
            affected += fail_dead_rows(
                &store.pool,
                &dead_ids,
                &dead_errors,
                &dead_reasons,
                worker.as_str(),
            )
            .await
            .map_err(|e| store.map_err(e))?;
        }

        affected
    } else {
        let mut tx = store.pool.begin().await.map_err(|e| store.map_err(e))?;

        store.set_local_timeout(&mut tx).await?;
        let mut affected = 0u64;

        if !retry_ids.is_empty() {
            affected += fail_retry_rows(
                &mut *tx,
                &retry_ids,
                &retry_errors,
                &retry_delays,
                worker.as_str(),
            )
            .await
            .map_err(|e| store.map_err(e))?;
        }

        if !dead_ids.is_empty() {
            affected += fail_dead_rows(
                &mut *tx,
                &dead_ids,
                &dead_errors,
                &dead_reasons,
                worker.as_str(),
            )
            .await
            .map_err(|e| store.map_err(e))?;
        }

        tx.commit().await.map_err(|e| store.map_err(e))?;

        affected
    };
    log_shortfall("fail", items.len(), affected);

    Ok(affected)
}

/// [`reliar_outbox::OutboxStore::release`]'s body.
pub(super) async fn release<Ser: Serializer + Send + Sync + 'static>(
    store: &PostgresOutboxStore<Ser>,
    worker: &WorkerId,
    items: &[MessageRef],
) -> Result<u64, PostgresStoreError> {
    if items.is_empty() {
        return Ok(0);
    }

    let ids: Vec<uuid::Uuid> = items.iter().map(|i| i.id.as_uuid()).collect();
    let affected = if store.settings.statement_timeout.is_zero() {
        release_rows(&store.pool, &ids, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?
    } else {
        let mut tx = store.pool.begin().await.map_err(|e| store.map_err(e))?;

        store.set_local_timeout(&mut tx).await?;
        let affected = release_rows(&mut *tx, &ids, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?;
        tx.commit().await.map_err(|e| store.map_err(e))?;

        affected
    };
    log_shortfall("release", items.len(), affected);

    Ok(affected)
}

/// [`reliar_outbox::OutboxStore::extend_lease`]'s body.
pub(super) async fn extend_lease<Ser: Serializer + Send + Sync + 'static>(
    store: &PostgresOutboxStore<Ser>,
    worker: &WorkerId,
    items: &[MessageRef],
    lease: Duration,
) -> Result<u64, PostgresStoreError> {
    if items.is_empty() {
        return Ok(0);
    }

    let ids: Vec<uuid::Uuid> = items.iter().map(|i| i.id.as_uuid()).collect();
    let lease_ms = i64::try_from(lease.as_millis()).unwrap_or(i64::MAX);
    let affected = if store.settings.statement_timeout.is_zero() {
        extend_lease_rows(&store.pool, &ids, lease_ms, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?
    } else {
        let mut tx = store.pool.begin().await.map_err(|e| store.map_err(e))?;

        store.set_local_timeout(&mut tx).await?;
        let affected = extend_lease_rows(&mut *tx, &ids, lease_ms, worker.as_str())
            .await
            .map_err(|e| store.map_err(e))?;
        tx.commit().await.map_err(|e| store.map_err(e))?;

        affected
    };
    log_shortfall("extend_lease", items.len(), affected);

    Ok(affected)
}

/// Worker-guarded `complete`: clears the lease and sets `published_at`, only for rows this
/// worker still holds (`locked_by = $2`) — a row already reclaimed by another worker
/// contributes nothing (ADR 0008).
async fn complete_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    ids: &[uuid::Uuid],
    worker: &str,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query!(
        r#"UPDATE outbox
              SET published_at = now(),
                  attempts     = attempts + 1,
                  locked_by    = NULL,
                  locked_until = NULL,
                  updated_at   = now()
            WHERE id = ANY($1) AND locked_by = $2"#,
        ids,
        worker,
    )
    .execute(executor)
    .await?;

    Ok(result.rows_affected())
}

/// Clears the lease and resets `available_at = now()` (ADR 0040 §1): a released row must be
/// claimable at once, not stalled behind the lease it was just released from — the drain-latency
/// regression `OutboxStore::release`'s rustdoc now states as normative.
async fn release_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    ids: &[uuid::Uuid],
    worker: &str,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query!(
        r#"UPDATE outbox
              SET locked_by    = NULL,
                  locked_until = NULL,
                  available_at = now(),
                  updated_at   = now()
            WHERE id = ANY($1) AND locked_by = $2"#,
        ids,
        worker,
    )
    .execute(executor)
    .await?;

    Ok(result.rows_affected())
}

/// Advances both `locked_until` and `available_at` together (ADR 0040 §1): without moving
/// `available_at` too, a renewed row would fall back into the claim's scan window while still
/// leased, reintroducing the leased-row-at-the-head cost this migration exists to remove — for
/// exactly the long publishes lease renewal is for.
async fn extend_lease_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    ids: &[uuid::Uuid],
    lease_ms: i64,
    worker: &str,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query!(
        r#"UPDATE outbox
              SET locked_until = now() + ($2::bigint * interval '1 millisecond'),
                  available_at = now() + ($2::bigint * interval '1 millisecond'),
                  updated_at   = now()
            WHERE id = ANY($1) AND locked_by = $3"#,
        ids,
        lease_ms,
        worker,
    )
    .execute(executor)
    .await?;

    Ok(result.rows_affected())
}

async fn fail_retry_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    ids: &[uuid::Uuid],
    errors: &[String],
    delays_ms: &[i64],
    worker: &str,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query!(
        r#"UPDATE outbox o
              SET attempts     = o.attempts + 1,
                  last_error   = f.err,
                  locked_by    = NULL,
                  locked_until = NULL,
                  available_at = now() + (f.delay_ms * interval '1 millisecond'),
                  updated_at   = now()
             FROM UNNEST($1::uuid[], $2::text[], $3::bigint[]) AS f(id, err, delay_ms)
            WHERE o.id = f.id AND o.locked_by = $4"#,
        ids,
        errors,
        delays_ms,
        worker,
    )
    .execute(executor)
    .await?;

    Ok(result.rows_affected())
}

async fn fail_dead_rows<'e>(
    executor: impl sqlx::PgExecutor<'e>,
    ids: &[uuid::Uuid],
    errors: &[String],
    reasons: &[&str],
    worker: &str,
) -> Result<u64, sqlx::Error> {
    let result = sqlx::query!(
        r#"UPDATE outbox o
              SET attempts     = o.attempts + 1,
                  last_error   = f.err,
                  dead_at      = now(),
                  dead_reason  = f.reason,
                  locked_by    = NULL,
                  locked_until = NULL,
                  updated_at   = now()
             FROM UNNEST($1::uuid[], $2::text[], $3::text[]) AS f(id, err, reason)
            WHERE o.id = f.id AND o.locked_by = $4"#,
        ids,
        errors,
        reasons as &[&str],
        worker,
    )
    .execute(executor)
    .await?;

    Ok(result.rows_affected())
}

/// Logs a claimed-vs-affected shortfall at `debug` — never an error (ADR 0008): it means the
/// lease was lost to another worker or the row was already retried/completed, both benign.
fn log_shortfall(operation: &'static str, claimed: usize, affected: u64) {
    let claimed = claimed as u64;

    if affected < claimed {
        tracing::debug!(
            operation,
            claimed,
            affected,
            "fewer rows affected than claimed"
        );
    }
}