use std::fmt::{self, Debug};
use std::marker::PhantomData;
use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, LazyLock, Mutex, PoisonError};
use std::time::Duration;
use ruststream::codec::CodecError;
use ruststream::{AckError, Carries, HeaderMap, IncomingMessage};
use sync_wrapper::SyncWrapper;
use thiserror::Error;
use tokio::runtime::Handle;
#[cfg(feature = "testing")]
use super::broker::Shared;
use super::claims::InWork;
use super::database::QueueDatabase;
use super::engine::{Claimed, Events, Now};
use super::form::advisory::LockHold;
use super::form::lease::settle::release;
#[cfg(feature = "testing")]
use super::form::lease::settle::release_in_process;
use super::form::lease::{LeaseBook, Slot};
use super::form::row_lock::BatchTx;
use super::headers::HeaderCell;
use super::queue::Queue;
#[cfg(feature = "testing")]
use super::testing::off_clock;
use super::transactional::{InboxMode, Plain, TxHold};
use super::tx::PoolTx;
use super::{Lane, QueueRow, RowLane};
use crate::inbox::pools::ServicePool;
pub(crate) mod settle;
use settle::Outcome;
pub(super) enum Hold<DB: QueueDatabase, Row: Events<DB>> {
Own(SyncWrapper<PoolTx<DB>>),
Batch(Arc<BatchTx<DB>>),
Lease {
book: &'static LeaseBook<DB, Row>,
slot: Slot,
tx: Option<TxHold<DB>>,
},
Advisory(LockHold<DB>),
Lent(TxHold<DB>),
}
pub struct InboxDelivery<DB: QueueDatabase, Row: Events<DB>, Mode = Plain> {
claimed: Claimed<Row>,
headers: Row::Headers,
pub(super) hold: Option<Hold<DB, Row>>,
queue: &'static Queue,
pub(super) pool: &'static ServicePool<DB>,
row_lent: <Row::Lane as Lane<Row>>::Lent,
#[cfg(feature = "testing")]
in_process: Option<Arc<Shared<DB>>>,
in_work: Option<InWork>,
_mode: PhantomData<fn() -> Mode>,
}
static ROW_GONE: LazyLock<CodecError> = LazyLock::new(|| CodecError::Decode(Box::new(RowGone)));
#[derive(Debug, Error)]
#[error("the claimed row is gone: the fetch returned no row for its id")]
struct RowGone;
fn row_gone(queue: &Queue, id: &impl Debug) {
tracing::warn!(
target: "ruststream_sqlx",
subscription = queue.name,
table = queue.table,
row = queue.row,
?id,
"the fetch returned no row for a claimed id; the decode-failure policy settles its \
delivery",
);
}
fn unlent_payload(row: &str) -> String {
format!(
"the table is in row mode: its deliveries lend the row itself and carry no payload, so a \
handler that decodes one fails each delivery by the decode policy; take `&{row}` as the \
handler's input, `&[{row}]` in a batch"
)
}
static UNLENT_LOGGED: Mutex<Vec<&'static Queue>> = Mutex::new(Vec::new());
#[cold]
#[inline(never)]
fn read_unlent(queue: &'static Queue) {
let mut logged = UNLENT_LOGGED.lock().unwrap_or_else(PoisonError::into_inner);
if logged.iter().any(|seen| ptr::eq(*seen, queue)) {
return;
}
logged.push(queue);
drop(logged);
tracing::warn!(
target: "ruststream_sqlx",
subscription = queue.name,
table = queue.table,
row = queue.row,
"{}",
unlent_payload(queue.row),
);
}
impl<DB: QueueDatabase, Row: Events<DB>, Mode> Debug for InboxDelivery<DB, Row, Mode> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("InboxDelivery")
.field("subscription", &self.queue.name)
.field("id", self.claimed.id::<DB>())
.finish_non_exhaustive()
}
}
impl<DB, Row, Mode> InboxDelivery<DB, Row, Mode>
where
DB: QueueDatabase,
Row: Events<DB>,
Mode: InboxMode,
{
pub(crate) fn own(
claimed: Claimed<Row>,
tx: PoolTx<DB>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
Self::held(claimed, Hold::Own(SyncWrapper::new(tx)), queue, pool)
}
pub(crate) fn lent(
claimed: Claimed<Row>,
hold: TxHold<DB>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
Self::held(claimed, Hold::Lent(hold), queue, pool)
}
pub(crate) fn batched(
claimed: Claimed<Row>,
batch: Arc<BatchTx<DB>>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
Self::held(claimed, Hold::Batch(batch), queue, pool)
}
pub(crate) fn leased(
claimed: Claimed<Row>,
book: &'static LeaseBook<DB, Row>,
lease: Row::Token,
tx: Option<TxHold<DB>>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
let slot = book.enter(claimed.id::<DB>(), lease);
Self::held(claimed, Hold::Lease { book, slot, tx }, queue, pool)
}
pub(crate) fn advised(
claimed: Claimed<Row>,
hold: LockHold<DB>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
Self::held(claimed, Hold::Advisory(hold), queue, pool)
}
pub(crate) const fn pool(&self) -> &'static ServicePool<DB> {
self.pool
}
fn held(
mut claimed: Claimed<Row>,
hold: Hold<DB, Row>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
let headers = match &mut claimed {
Claimed::Row(row) => HeaderCell::take(row),
Claimed::Missing(id) => {
row_gone(queue, id);
Row::Headers::default()
}
Claimed::Undecodable { id, attempt, error } => {
tracing::warn!(
target: "ruststream_sqlx",
subscription = queue.name,
table = queue.table,
row = queue.row,
?id,
attempt,
%error,
"the row does not decode into its struct; the decode-failure policy settles \
its delivery",
);
Row::Headers::default()
}
};
Self {
claimed,
headers,
hold: Some(hold),
queue,
pool,
row_lent: Default::default(),
#[cfg(feature = "testing")]
in_process: None,
in_work: None,
_mode: PhantomData,
}
}
pub(crate) fn counted(mut self, in_work: InWork) -> Self {
self.in_work = Some(in_work);
self
}
#[cfg(feature = "testing")]
pub(crate) fn on(mut self, connection: &Arc<Shared<DB>>) -> Self {
self.in_process = connection
.harness
.in_process()
.then(|| Arc::clone(connection));
self
}
}
impl<DB, Row> InboxDelivery<DB, Row>
where
DB: QueueDatabase,
Row: Events<DB> + QueueRow<Lane = RowLane>,
{
pub(super) fn of_batch(
row: Row,
headers: Row::Headers,
hold: Hold<DB, Row>,
queue: &'static Queue,
pool: &'static ServicePool<DB>,
) -> Self {
Self {
claimed: Claimed::Row(row),
headers,
hold: Some(hold),
queue,
pool,
row_lent: AtomicBool::new(true),
#[cfg(feature = "testing")]
in_process: None,
in_work: None,
_mode: PhantomData,
}
}
}
impl<DB, Row, Mode> Carries<Row> for InboxDelivery<DB, Row, Mode>
where
DB: QueueDatabase,
Row: Events<DB> + QueueRow<Lane = RowLane>,
Mode: InboxMode,
{
fn carried(&self) -> Option<&Row> {
match &self.claimed {
Claimed::Row(row) => {
self.row_lent.store(true, Ordering::Relaxed);
Some(row)
}
Claimed::Missing(_) | Claimed::Undecodable { .. } => None,
}
}
}
impl<DB, Row, Mode> IncomingMessage for InboxDelivery<DB, Row, Mode>
where
DB: QueueDatabase,
Row: Events<DB>,
Mode: InboxMode,
{
fn payload(&self) -> &[u8] {
match &self.claimed {
Claimed::Row(row) => {
if <Row::Lane as Lane<Row>>::unlent(&self.row_lent) {
read_unlent(self.queue);
}
<Row::Lane as Lane<Row>>::payload(row)
}
Claimed::Missing(_) | Claimed::Undecodable { .. } => &[],
}
}
fn headers(&self) -> &HeaderMap {
let row = match &self.claimed {
Claimed::Row(row) => Some(row),
Claimed::Missing(_) | Claimed::Undecodable { .. } => None,
};
self.headers.read(row)
}
fn decode_error(&self) -> Option<&CodecError> {
match &self.claimed {
Claimed::Row(_) => None,
Claimed::Missing(_) => Some(LazyLock::force(&ROW_GONE)),
Claimed::Undecodable { error, .. } => Some(error),
}
}
fn partition_key(&self) -> Option<&[u8]> {
match &self.claimed {
Claimed::Row(row) => Row::partition_key(row),
Claimed::Missing(_) | Claimed::Undecodable { .. } => None,
}
}
fn redelivery_count(&self) -> Option<u64> {
let carried = match &self.claimed {
Claimed::Row(row) => Row::attempt(row),
Claimed::Undecodable { attempt, .. } => *attempt,
Claimed::Missing(_) => None,
};
if self.queue.counted_attempt {
carried.map(|attempt| attempt.saturating_sub(1))
} else {
carried
}
}
async fn ack(self) -> Result<(), AckError> {
self.settle(Outcome::Ack).await
}
async fn nack(self, requeue: bool) -> Result<(), AckError> {
self.settle(if requeue {
Outcome::Retry
} else {
Outcome::Discard
})
.await
}
fn supports_nack_after(&self) -> bool {
self.queue.native_retry_after && self.spent().is_none()
}
async fn nack_after(self, delay: Duration) -> Result<(), AckError> {
if !self.supports_nack_after() {
return Err(AckError::Unsupported);
}
self.settle(Outcome::RetryAfter(delay)).await
}
}
impl<DB: QueueDatabase, Row: Events<DB>, Mode> Drop for InboxDelivery<DB, Row, Mode> {
fn drop(&mut self) {
let Some(hold) = self.hold.take() else {
return;
};
#[cfg(feature = "testing")]
let in_process = self.in_process.take();
match hold {
#[cfg(feature = "testing")]
Hold::Own(tx) if in_process.is_some() => {
if let Ok(runtime) = Handle::try_current() {
drop(runtime.spawn(off_clock(tx.into_inner().rollback())));
}
}
Hold::Own(tx) => {
if let Ok(runtime) = Handle::try_current() {
drop(runtime.spawn(tx.into_inner().rollback()));
}
}
Hold::Batch(batch) => batch.release(self.queue),
Hold::Lent(hold) => drop(hold),
Hold::Advisory(hold) => {
#[cfg(feature = "testing")]
if let Some(connection) = in_process {
connection.harness.expect(self.queue.name);
drop(hold);
connection.harness.released();
return;
}
drop(hold);
}
Hold::Lease { book, slot, tx } => {
drop(tx);
let id = self.claimed.id::<DB>().clone();
#[cfg(feature = "testing")]
if let Some(connection) = in_process {
release_in_process(connection, book, slot, self.queue, id);
return;
}
drop(book.runtime().spawn(release::<DB, Row>(
book,
slot,
self.queue,
id,
Now::default(),
)));
}
}
#[cfg(feature = "testing")]
if let Some(connection) = in_process {
connection.harness.returned(self.queue.name);
}
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::fmt::Debug;
use std::mem;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use ruststream_sqlx_dialect::{Column, Form, TableSpec};
use tracing::field::{Field, Visit};
use tracing::span::{Attributes, Id, Record};
use tracing::{Event, Level, Metadata, Subscriber};
use super::{row_gone, unlent_payload};
use crate::inbox::engine::{IdAt, Prepared};
use crate::inbox::queue::Queue;
#[derive(Default)]
struct Logged(Mutex<Vec<(Level, BTreeMap<String, String>)>>);
struct Fields<'a>(&'a mut BTreeMap<String, String>);
impl Visit for Fields<'_> {
fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
self.0.insert(field.name().to_owned(), format!("{value:?}"));
}
fn record_str(&mut self, field: &Field, value: &str) {
self.0.insert(field.name().to_owned(), value.to_owned());
}
}
impl Subscriber for Logged {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &Attributes<'_>) -> Id {
Id::from_u64(1)
}
fn record(&self, _: &Id, _: &Record<'_>) {}
fn record_follows_from(&self, _: &Id, _: &Id) {}
fn event(&self, event: &Event<'_>) {
let mut fields = BTreeMap::new();
event.record(&mut Fields(&mut fields));
self.0
.lock()
.expect("the log is not poisoned")
.push((*event.metadata().level(), fields));
}
fn enter(&self, _: &Id) {}
fn exit(&self, _: &Id) {}
}
#[test]
fn a_claimed_id_without_a_row_is_logged_with_its_id() {
let queue = Queue {
name: "mail",
table: "mail_jobs",
row: "app::FetchedMail",
spec: TableSpec::new("mail_jobs", Column::new("job_id"), Form::RowLock),
id_at: IdAt::First,
native_retry_after: false,
kinds: None,
prepared: Prepared::default(),
begin_claim: None,
counted_attempt: false,
one_writer: false,
poll_interval: Duration::from_secs(1),
lease: None,
cap: None,
};
let logged = Arc::new(Logged::default());
tracing::subscriber::with_default(Arc::clone(&logged), || row_gone(&queue, &42_i64));
let logged = mem::take(&mut *logged.0.lock().expect("the log is not poisoned"));
let [(level, fields)] = logged.as_slice() else {
panic!("one event, not {}", logged.len());
};
assert_eq!(*level, Level::WARN);
assert_eq!(fields.get("id").map(String::as_str), Some("42"));
assert_eq!(fields.get("subscription").map(String::as_str), Some("mail"));
assert_eq!(fields.get("table").map(String::as_str), Some("mail_jobs"));
assert_eq!(
fields.get("row").map(String::as_str),
Some("app::FetchedMail")
);
}
#[test]
fn a_payload_read_from_a_row_mode_delivery_names_the_row_to_take() {
assert_eq!(
unlent_payload("app::SendEmail"),
"the table is in row mode: its deliveries lend the row itself and carry no payload, so \
a handler that decodes one fails each delivery by the decode policy; take \
`&app::SendEmail` as the handler's input, `&[app::SendEmail]` in a batch",
);
}
}