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

use serde_json::Value;
use sqlx::PgPool;
use time::OffsetDateTime;

use crate::error::DurableError;
use crate::tx::{self, PgTx};

/// Ownership token returned by claim; required for projection and release.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct OutboxClaim {
    pub message_id: i64,
    pub worker_id: String,
    pub generation: i64,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct OutboxMessage {
    pub id: i64,
    pub topic: String,
    pub payload: Value,
    pub created_at: OffsetDateTime,
    pub available_at: OffsetDateTime,
    pub attempt_count: i32,
    pub claim: OutboxClaim,
}

/// Selects which outbox rows a claim may take.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OutboxClaimFilter<'a> {
    /// Any claimable topic.
    Any,
    /// Exact topic match.
    Topic(&'a str),
}

pub async fn enqueue(tx: &mut PgTx<'_>, topic: &str, payload: Value) -> Result<i64, DurableError> {
    let (id,): (i64,) = sqlx::query_as(
        "INSERT INTO outbox_messages (topic, payload, available_at)
         VALUES ($1, $2, now())
         RETURNING id",
    )
    .bind(topic)
    .bind(payload)
    .fetch_one(tx.as_mut())
    .await?;
    Ok(id)
}

/// Enqueue work available after a delay measured by the database clock.
pub async fn enqueue_after(
    tx: &mut PgTx<'_>,
    topic: &str,
    payload: Value,
    delay: time::Duration,
) -> Result<i64, DurableError> {
    let delay_secs = delay.as_seconds_f64().max(0.0);
    let (id,): (i64,) = sqlx::query_as(
        "INSERT INTO outbox_messages (topic, payload, available_at)
         VALUES ($1, $2, now() + ($3 * interval '1 second'))
         RETURNING id",
    )
    .bind(topic)
    .bind(payload)
    .bind(delay_secs)
    .fetch_one(tx.as_mut())
    .await?;
    Ok(id)
}

/// Claim claimable outbox rows for exclusive projection using
/// `FOR UPDATE SKIP LOCKED`. Lease eligibility and expiry use the database
/// clock. Each successful claim advances `claim_generation` and returns an
/// ownership token required for projection or release.
pub async fn claim_batch(
    pool: &PgPool,
    worker_id: &str,
    filter: OutboxClaimFilter<'_>,
    limit: i64,
    lease: time::Duration,
) -> Result<Vec<OutboxMessage>, DurableError> {
    let worker_id = worker_id.to_owned();
    let topic = match filter {
        OutboxClaimFilter::Any => None,
        OutboxClaimFilter::Topic(topic) => Some(topic.to_owned()),
    };
    tx::with_transaction(pool, move |tx| {
        Box::pin(async move {
            let filter = match topic.as_deref() {
                None => OutboxClaimFilter::Any,
                Some(topic) => OutboxClaimFilter::Topic(topic),
            };
            claim_batch_in_tx(tx, &worker_id, filter, limit, lease).await
        })
    })
    .await
}

pub async fn claim_batch_in_tx(
    tx: &mut PgTx<'_>,
    worker_id: &str,
    filter: OutboxClaimFilter<'_>,
    limit: i64,
    lease: time::Duration,
) -> Result<Vec<OutboxMessage>, DurableError> {
    let lease_secs = lease.as_seconds_f64().max(0.0);
    let topic = match filter {
        OutboxClaimFilter::Any => None,
        OutboxClaimFilter::Topic(topic) => Some(topic),
    };
    let rows = sqlx::query_as::<
        _,
        (
            i64,
            String,
            Value,
            OffsetDateTime,
            OffsetDateTime,
            i32,
            i64,
            String,
        ),
    >(
        "WITH candidates AS (
            SELECT id
            FROM outbox_messages
            WHERE projected_at IS NULL
              AND available_at <= now()
              AND (claimed_until IS NULL OR claimed_until < now())
              AND ($4::text IS NULL OR topic = $4)
            ORDER BY id
            FOR UPDATE SKIP LOCKED
            LIMIT $1
         )
         UPDATE outbox_messages AS messages
         SET claimed_by = $2,
             claimed_until = now() + ($3 * interval '1 second'),
             claim_generation = messages.claim_generation + 1,
             attempt_count = messages.attempt_count + 1
         FROM candidates
         WHERE messages.id = candidates.id
         RETURNING
            messages.id,
            messages.topic,
            messages.payload,
            messages.created_at,
            messages.available_at,
            messages.attempt_count,
            messages.claim_generation,
            messages.claimed_by",
    )
    .bind(limit)
    .bind(worker_id)
    .bind(lease_secs)
    .bind(topic)
    .fetch_all(tx.as_mut())
    .await?;

    Ok(rows
        .into_iter()
        .map(
            |(
                id,
                topic,
                payload,
                created_at,
                available_at,
                attempt_count,
                claim_generation,
                claimed_by,
            )| {
                OutboxMessage {
                    id,
                    topic,
                    payload,
                    created_at,
                    available_at,
                    attempt_count,
                    claim: OutboxClaim {
                        message_id: id,
                        worker_id: claimed_by,
                        generation: claim_generation,
                    },
                }
            },
        )
        .collect())
}

/// Mark a claimed message projected. Idempotent when the row is already
/// projected. Requires the live claim token when still pending.
pub async fn mark_projected(tx: &mut PgTx<'_>, claim: &OutboxClaim) -> Result<(), DurableError> {
    let result = sqlx::query(
        "UPDATE outbox_messages
         SET projected_at = now(),
             claimed_by = NULL,
             claimed_until = NULL,
             last_error = NULL
         WHERE id = $1
           AND projected_at IS NULL
           AND claimed_by = $2
           AND claim_generation = $3",
    )
    .bind(claim.message_id)
    .bind(&claim.worker_id)
    .bind(claim.generation)
    .execute(tx.as_mut())
    .await?;
    if result.rows_affected() == 1 {
        return Ok(());
    }

    let projected: Option<(bool,)> = sqlx::query_as(
        "SELECT projected_at IS NOT NULL
         FROM outbox_messages
         WHERE id = $1",
    )
    .bind(claim.message_id)
    .fetch_optional(tx.as_mut())
    .await?;

    match projected {
        Some((true,)) => Ok(()),
        Some((false,)) => Err(DurableError::OutboxClaimMismatch {
            message_id: claim.message_id,
        }),
        None => Err(DurableError::OutboxMissing {
            message_id: claim.message_id,
        }),
    }
}

pub async fn mark_projected_pool(pool: &PgPool, claim: &OutboxClaim) -> Result<(), DurableError> {
    let claim = claim.clone();
    tx::with_transaction(pool, move |tx| {
        Box::pin(async move { mark_projected(tx, &claim).await })
    })
    .await
}

/// Release a claimed message for later retry. Fenced to the claim token so a
/// stale claimant cannot clear a reclaimed lease. Retry availability uses the
/// database clock.
pub async fn release_with_error(
    tx: &mut PgTx<'_>,
    claim: &OutboxClaim,
    error: &str,
    retry_delay: time::Duration,
) -> Result<(), DurableError> {
    let delay_secs = retry_delay.as_seconds_f64().max(0.0);
    let result = sqlx::query(
        "UPDATE outbox_messages
         SET claimed_by = NULL,
             claimed_until = NULL,
             available_at = now() + ($4 * interval '1 second'),
             last_error = $5
         WHERE id = $1
           AND projected_at IS NULL
           AND claimed_by = $2
           AND claim_generation = $3",
    )
    .bind(claim.message_id)
    .bind(&claim.worker_id)
    .bind(claim.generation)
    .bind(delay_secs)
    .bind(error)
    .execute(tx.as_mut())
    .await?;
    if result.rows_affected() != 1 {
        return Err(DurableError::OutboxClaimMismatch {
            message_id: claim.message_id,
        });
    }
    Ok(())
}

pub async fn unprojected_count(
    pool: &PgPool,
    filter: OutboxClaimFilter<'_>,
) -> Result<i64, DurableError> {
    let topic = match filter {
        OutboxClaimFilter::Any => None,
        OutboxClaimFilter::Topic(topic) => Some(topic),
    };
    let (count,): (i64,) = sqlx::query_as(
        "SELECT COUNT(*)::bigint
         FROM outbox_messages
         WHERE projected_at IS NULL
           AND ($1::text IS NULL OR topic = $1)",
    )
    .bind(topic)
    .fetch_one(pool)
    .await?;
    Ok(count)
}