use reliar_core::{ConversationId, CorrelationId, MessageId, MessageType};
use reliar_inbox::InboxStore;
use reliar_inbox::{
InboxClaim, InboxFailure, InboxMessage, InboxPurgeReport, InboxPurgeRequest, InboxRecord,
InboxRecordId, InboxScope,
};
use sqlx::{PgConnection, Postgres, Transaction};
use tracing::Instrument as _;
use crate::connection::session::Session;
use crate::records::truncate_last_error;
use crate::settings::PostgresInboxSettings;
use super::claim as claim_repo;
use super::claim::ClaimStateRow;
use super::error::PostgresInboxError;
use super::outcomes as outcomes_repo;
use super::purge as purge_repo;
use super::rows::InboxRow;
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct PostgresInboxStore {
pub(super) session: Session,
settings: PostgresInboxSettings,
}
impl PostgresInboxStore {
#[must_use]
pub fn new(pool: sqlx::PgPool) -> Self {
let settings = PostgresInboxSettings::default();
let session = Session::new(pool, settings.statement_timeout);
Self { session, settings }
}
pub fn with_settings(
pool: sqlx::PgPool,
settings: PostgresInboxSettings,
) -> Result<Self, PostgresInboxError> {
settings.validate()?;
let session = Session::new(pool, settings.statement_timeout);
Ok(Self { session, settings })
}
pub(super) fn max_attempts(&self) -> u32 {
self.settings.max_attempts
}
}
impl<'c> InboxStore<Transaction<'c, Postgres>> for PostgresInboxStore {
type Error = PostgresInboxError;
fn claim(
&self,
tx: &mut Transaction<'c, Postgres>,
scope: &InboxScope,
message: InboxMessage<'_>,
) -> impl Future<Output = Result<InboxClaim, Self::Error>> + Send {
let span = tracing::debug_span!(
"reliar.inbox.claim",
inbox.scope = %scope,
message.id = %message.id,
message.r#type = %message.message_type,
inbox.outcome = tracing::field::Empty,
inbox.attempt = tracing::field::Empty,
inbox.attempts = tracing::field::Empty,
inbox.record_id = tracing::field::Empty,
);
let recording_span = span.clone();
async move {
let result = claim_locked(tx, scope.as_str(), message).await;
if let Ok(claim) = &result {
record_claim_outcome(&recording_span, claim);
}
result
}
.instrument(span)
}
fn complete(
&self,
tx: &mut Transaction<'c, Postgres>,
scope: &InboxScope,
id: MessageId,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
let span = tracing::debug_span!(
"reliar.inbox.complete",
inbox.scope = %scope,
message.id = %id,
);
async move {
let scope_str = scope.as_str();
let message_id = id.as_uuid();
let affected = outcomes_repo::complete_row(
tx,
outcomes_repo::CompleteRowParams {
scope: scope_str,
message_id,
},
)
.await?;
if affected == 0 {
return Err(PostgresInboxError::NotClaimed {
scope: scope_str.to_owned(),
message_id: id,
});
}
Ok(())
}
.instrument(span)
}
fn fail(
&self,
scope: &InboxScope,
message: InboxMessage<'_>,
error: &(dyn std::error::Error + 'static),
) -> impl Future<Output = Result<InboxFailure, Self::Error>> + Send {
let last_error = format_error_chain(error);
let span = tracing::debug_span!(
"reliar.inbox.fail",
inbox.scope = %scope,
message.id = %message.id,
message.r#type = %message.message_type,
inbox.outcome = tracing::field::Empty,
inbox.attempts = tracing::field::Empty,
inbox.record_id = tracing::field::Empty,
);
let recording_span = span.clone();
async move {
let session = &self.session;
let scope_str = scope.as_str();
let last_error = truncate_last_error(last_error);
let max_attempts = i32::try_from(self.max_attempts()).unwrap_or(i32::MAX);
let id = InboxRecordId::new().as_uuid();
let message_id = message.id.as_uuid();
let message_type = message.message_type.name();
let message_version = i32::from(message.message_type.version());
let conversation_id = message.conversation_id.as_uuid();
let correlation_id = message.correlation_id.map(CorrelationId::as_str);
let causation_id = message.causation_id.map(|c| c.as_uuid());
let row = session
.run(async |conn: &mut PgConnection| {
outcomes_repo::fail_row(
&mut *conn,
outcomes_repo::FailRowParams {
id,
scope: scope_str,
message_id,
message_type,
message_version,
conversation_id,
correlation_id,
causation_id,
last_error: &last_error,
max_attempts,
},
)
.await
})
.await
.map_err(|e| session.map_err::<PostgresInboxError>(e))?;
let failure = match row {
None => InboxFailure::AlreadyCompleted,
Some(row) => {
let attempts = u32::try_from(row.attempts).unwrap_or(u32::MAX);
match row.dead_at {
Some(dead_at) => InboxFailure::Dead {
id: InboxRecordId::from_uuid(row.id),
attempts,
dead_at,
},
None => InboxFailure::Recorded { attempts },
}
}
};
record_fail_outcome(&recording_span, &failure);
Ok(failure)
}
.instrument(span)
}
async fn find(
&self,
scope: &InboxScope,
id: MessageId,
) -> Result<Option<InboxRecord>, PostgresInboxError> {
let session = &self.session;
let scope_str = scope.as_str();
let message_id = id.as_uuid();
let row = session
.run(async |conn: &mut PgConnection| {
purge_repo::find_row(
&mut *conn,
purge_repo::FindRowParams {
scope: scope_str,
message_id,
},
)
.await
})
.await
.map_err(|e| session.map_err::<PostgresInboxError>(e))?;
let Some(row) = row else {
return Ok(None);
};
Ok(Some(build_record(row)?))
}
fn purge(
&self,
request: InboxPurgeRequest,
) -> impl Future<Output = Result<InboxPurgeReport, Self::Error>> + Send {
let span = tracing::debug_span!(
"reliar.inbox.purge",
inbox.completed_deleted = tracing::field::Empty,
inbox.incomplete_deleted = tracing::field::Empty,
inbox.dead_deleted = tracing::field::Empty,
);
let recording_span = span.clone();
async move {
let session = &self.session;
let batch_size = i64::from(request.batch_size);
let completed_ms = request.completed_retention.map(to_millis);
let incomplete_ms = request.incomplete_retention.map(to_millis);
let dead_ms = request.dead_retention.map(to_millis);
let (completed_deleted, incomplete_deleted, dead_deleted) = session
.run(async |conn: &mut PgConnection| {
let completed_deleted = match completed_ms {
Some(retention_ms) => {
purge_repo::purge_completed_rows(
&mut *conn,
purge_repo::PurgeCompletedRowsParams {
retention_ms,
batch_size,
},
)
.await?
}
None => 0,
};
let incomplete_deleted = match incomplete_ms {
Some(retention_ms) => {
purge_repo::purge_incomplete_rows(
&mut *conn,
purge_repo::PurgeIncompleteRowsParams {
retention_ms,
batch_size,
},
)
.await?
}
None => 0,
};
let dead_deleted = match dead_ms {
Some(retention_ms) => {
purge_repo::purge_dead_retention_rows(
&mut *conn,
purge_repo::PurgeDeadRetentionRowsParams {
retention_ms,
batch_size,
},
)
.await?
}
None => 0,
};
Ok((completed_deleted, incomplete_deleted, dead_deleted))
})
.await
.map_err(|e| session.map_err::<PostgresInboxError>(e))?;
let report = InboxPurgeReport::new(completed_deleted, incomplete_deleted, dead_deleted);
recording_span.record("inbox.completed_deleted", report.completed_deleted);
recording_span.record("inbox.incomplete_deleted", report.incomplete_deleted);
recording_span.record("inbox.dead_deleted", report.dead_deleted);
Ok(report)
}
.instrument(span)
}
}
fn record_claim_outcome(span: &tracing::Span, claim: &InboxClaim) {
match claim {
InboxClaim::Claimed { attempt } => {
span.record("inbox.outcome", "claimed");
span.record("inbox.attempt", attempt);
}
InboxClaim::AlreadyCompleted { .. } => {
span.record("inbox.outcome", "already_completed");
}
InboxClaim::InProgress => {
span.record("inbox.outcome", "in_progress");
}
InboxClaim::Dead { id, attempts, .. } => {
span.record("inbox.outcome", "dead");
span.record("inbox.record_id", tracing::field::display(id));
span.record("inbox.attempts", attempts);
}
_ => {}
}
}
fn record_fail_outcome(span: &tracing::Span, failure: &InboxFailure) {
match failure {
InboxFailure::Recorded { attempts } => {
span.record("inbox.outcome", "recorded");
span.record("inbox.attempts", attempts);
}
InboxFailure::Dead {
id,
attempts,
dead_at: _,
} => {
span.record("inbox.outcome", "dead");
span.record("inbox.attempts", attempts);
span.record("inbox.record_id", tracing::field::display(id));
}
InboxFailure::AlreadyCompleted => {
span.record("inbox.outcome", "already_completed");
}
_ => {}
}
}
const ADVISORY_LOCK_CLASS: i32 = i32::from_be_bytes(*b"RELI");
async fn claim_locked(
tx: &mut Transaction<'_, Postgres>,
scope: &str,
message: InboxMessage<'_>,
) -> Result<InboxClaim, PostgresInboxError> {
let message_id = message.id.as_uuid();
let acquired = claim_repo::try_advisory_lock(
&mut **tx,
claim_repo::TryAdvisoryLockParams {
class: ADVISORY_LOCK_CLASS,
scope,
message_id,
},
)
.await?;
if !acquired {
return Ok(InboxClaim::InProgress);
}
let id = InboxRecordId::new();
let claim_params = claim_row_params(id, scope, &message);
if let Some(attempts) = claim_repo::insert_claim_row(&mut **tx, claim_params).await? {
return Ok(InboxClaim::Claimed {
attempt: claimed_attempt(attempts),
});
}
let row = claim_repo::select_claim_state(
&mut **tx,
claim_repo::SelectClaimStateParams { scope, message_id },
)
.await?;
let Some(row) = row else {
let upsert_params = claim_row_params(InboxRecordId::new(), scope, &message);
let row = claim_repo::upsert_claim_row(&mut **tx, upsert_params).await?;
return Ok(claim_from_state(&row));
};
Ok(claim_from_state(&row))
}
fn claim_row_params<'a>(
id: InboxRecordId,
scope: &'a str,
message: &InboxMessage<'a>,
) -> claim_repo::ClaimRowParams<'a> {
claim_repo::ClaimRowParams {
id: id.as_uuid(),
scope,
message_id: message.id.as_uuid(),
message_type: message.message_type.name(),
message_version: i32::from(message.message_type.version()),
conversation_id: message.conversation_id.as_uuid(),
correlation_id: message.correlation_id.map(CorrelationId::as_str),
causation_id: message.causation_id.map(|c| c.as_uuid()),
}
}
fn claim_from_state(row: &ClaimStateRow) -> InboxClaim {
if let Some(completed_at) = row.completed_at {
return InboxClaim::AlreadyCompleted { completed_at };
}
if let Some(dead_at) = row.dead_at {
return InboxClaim::Dead {
id: InboxRecordId::from_uuid(row.id),
attempts: claimed_attempts_recorded(row.attempts),
dead_at,
};
}
InboxClaim::Claimed {
attempt: claimed_attempt(row.attempts),
}
}
fn claimed_attempt(attempts: i32) -> u32 {
u32::try_from(attempts)
.unwrap_or(u32::MAX)
.saturating_add(1)
}
fn claimed_attempts_recorded(attempts: i32) -> u32 {
u32::try_from(attempts).unwrap_or(u32::MAX)
}
fn format_error_chain(error: &(dyn std::error::Error + 'static)) -> String {
let mut out = error.to_string();
let mut source = error.source();
while let Some(err) = source {
out.push_str(": ");
out.push_str(&err.to_string());
source = err.source();
}
out
}
pub(super) fn build_record(row: InboxRow) -> Result<InboxRecord, PostgresInboxError> {
let scope = InboxScope::new(row.scope).map_err(|err| PostgresInboxError::Database {
source: sqlx::Error::Decode(err.into()),
})?;
let correlation_id = row
.correlation_id
.map(CorrelationId::parse)
.transpose()
.map_err(|err| PostgresInboxError::Database {
source: sqlx::Error::Decode(err.into()),
})?;
let message_type = MessageType::from_parts(
row.message_type,
u16::try_from(row.message_version).unwrap_or(u16::MAX),
);
let message_id = MessageId::from_uuid(row.message_id);
let mut message = InboxMessage::new(message_id, &message_type)
.conversation(ConversationId::from_uuid(row.conversation_id));
if let Some(correlation_id) = correlation_id.as_ref() {
message = message.correlation(correlation_id);
}
if let Some(causation_id) = row.causation_id {
message = message.causation(MessageId::from_uuid(causation_id));
}
Ok(InboxRecord::builder(
InboxRecordId::from_uuid(row.id),
scope,
message,
row.received_at,
)
.updated_at(row.updated_at)
.completed_at(row.completed_at)
.dead_at(row.dead_at)
.attempts(u32::try_from(row.attempts).unwrap_or(u32::MAX))
.last_error(row.last_error)
.build())
}
fn to_millis(duration: std::time::Duration) -> i64 {
i64::try_from(duration.as_millis()).unwrap_or(i64::MAX)
}