294 lines
8.2 KiB
Rust
294 lines
8.2 KiB
Rust
|
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)
|
||
|
|
}
|