use std::time::Duration;
use sqlx::PgPool;
use uuid::Uuid;
use crate::error::{BusError, BusResult};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Locator(pub String);
#[derive(Clone, Debug)]
pub enum Published {
Confirmed(Locator),
Retryable(String),
Fatal(String),
}
#[derive(Clone, Debug)]
pub struct Envelope {
pub message_id: Uuid,
pub conversation_id: Uuid,
pub team_id: Uuid,
pub body: String,
pub publish_key: Uuid,
}
pub trait MessagingBackend: Send + Sync {
fn name(&self) -> &'static str;
fn publish(&self, envelope: Envelope) -> impl std::future::Future<Output = Published> + Send;
fn fetch(
&self,
locator: &Locator,
message_id: Uuid,
) -> impl std::future::Future<Output = BusResult<Option<String>>> + Send;
fn retain(
&self,
before: chrono::DateTime<chrono::Utc>,
) -> impl std::future::Future<Output = BusResult<u64>> + Send;
fn reconcile(
&self,
envelope: &Envelope,
) -> impl std::future::Future<Output = BusResult<Option<Locator>>> + Send;
}
#[derive(Clone)]
pub struct PostgresBackend {
pool: PgPool,
team_id: Option<Uuid>,
faults: Faults,
retryable_left: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
#[derive(Clone, Debug, Default)]
pub struct Faults {
pub retryable: usize,
pub fatal: bool,
pub lose_confirmation: bool,
pub delay: Option<Duration>,
pub fail_reconcile: bool,
}
impl PostgresBackend {
pub fn new(pool: PgPool) -> Self {
Self {
pool,
team_id: None,
faults: Faults::default(),
retryable_left: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
pub fn with_team(mut self, team_id: Uuid) -> Self {
self.team_id = Some(team_id);
self
}
pub fn with_faults(pool: PgPool, faults: Faults) -> Self {
let left = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(faults.retryable));
Self {
pool,
team_id: None,
faults,
retryable_left: left,
}
}
pub const NAME: &'static str = "postgres";
}
impl MessagingBackend for PostgresBackend {
fn name(&self) -> &'static str {
Self::NAME
}
async fn publish(&self, envelope: Envelope) -> Published {
if let Some(delay) = self.faults.delay {
tokio::time::sleep(delay).await;
}
if self.faults.fatal {
return Published::Fatal("the backend refused this payload".into());
}
if self
.retryable_left
.fetch_update(
std::sync::atomic::Ordering::SeqCst,
std::sync::atomic::Ordering::SeqCst,
|left| (left > 0).then(|| left - 1),
)
.is_ok()
{
return Published::Retryable("the backend was unreachable".into());
}
let done = sqlx::query(
"UPDATE conversation_messages
SET canonical_locator = $2
WHERE id = $1 AND (canonical_locator IS NULL OR canonical_locator = $2)",
)
.bind(envelope.message_id)
.bind(envelope.message_id.to_string())
.execute(&self.pool)
.await;
match done {
Err(e) => {
tracing::warn!(error = %e, "publish failed");
Published::Retryable("the backend write failed".into())
}
Ok(_) if self.faults.lose_confirmation => {
Published::Retryable("the confirmation was lost".into())
}
Ok(_) => Published::Confirmed(Locator(envelope.message_id.to_string())),
}
}
async fn fetch(&self, locator: &Locator, message_id: Uuid) -> BusResult<Option<String>> {
let id: Uuid = locator
.0
.parse()
.map_err(|_| BusError::invalid("not a locator this backend issued"))?;
if id != message_id {
return Err(BusError::Forbidden(
"that locator names another message".to_owned(),
));
}
let row: Option<(String,)> = sqlx::query_as(
"SELECT m.body FROM conversation_messages m
JOIN conversations c ON c.id = m.conversation_id
WHERE m.id = $1 AND ($2::uuid IS NULL OR c.team_id = $2)",
)
.bind(id)
.bind(self.team_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|r| r.0))
}
async fn retain(&self, _before: chrono::DateTime<chrono::Utc>) -> BusResult<u64> {
Ok(0)
}
async fn reconcile(&self, envelope: &Envelope) -> BusResult<Option<Locator>> {
if self.faults.fail_reconcile {
return Err(BusError::invalid("the backend could not be asked"));
}
let row: Option<(Uuid,)> = sqlx::query_as(
"SELECT m.id FROM conversation_messages m
JOIN conversation_outbox o ON o.message_id = m.id
WHERE o.publish_key = $1 AND m.canonical_locator IS NOT NULL",
)
.bind(envelope.publish_key)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|r| Locator(r.0.to_string())))
}
}