149bec6
use std::future::Future;
use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;
use redis::aio::ConnectionManager;
use redis::cluster_async::ClusterConnection;
use serde::Serialize;
use serde::de::DeserializeOwned;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, warn};
use crate::error::EmailError;
#[derive(Clone)]
pub enum QueueConnection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
impl QueueConnection {
pub fn is_cluster(&self) -> bool {
matches!(self, QueueConnection::Cluster(_))
}
}
macro_rules! dispatch_conn {
($conn:expr, |$c:ident| $body:expr) => {
match &mut $conn {
QueueConnection::Standalone($c) => $body,
QueueConnection::Cluster($c) => $body,
}
};
}
pub struct EmailQueue {
connection: QueueConnection,
config: QueueSettings,
pubsub_client: redis::Client,
}
#[derive(Clone)]
pub struct QueueSettings {
pub channel_prefix: String,
pub max_retries: u32,
pub retry_backoff: Duration,
pub enabled: bool,
}
impl EmailQueue {
pub async fn connect(url: &str, config: QueueSettings) -> Result<Self, EmailError> {
let connection = open_connection(url).await?;
let pubsub_client = open_pubsub_client(url)?;
Ok(Self {
connection,
config,
pubsub_client,
})
}
pub fn enabled(&self) -> bool {
self.config.enabled
}
pub fn channel(&self, name: &str) -> String {
format!("{}.{}", self.config.channel_prefix, name)
}
pub fn config(&self) -> &QueueSettings {
&self.config
}
pub async fn publish<T: Serialize + ?Sized>(
&self,
channel: &str,
msg: &T,
) -> Result<(), EmailError> {
let payload = serde_json::to_vec(msg)?;
let mut conn = self.connection.clone();
let receiver_count: i64 = dispatch_conn!(conn, |c| {
Box::pin(
redis::cmd("PUBLISH")
.arg(channel)
.arg(payload.as_slice())
.query_async(c),
)
.await
})?;
debug!(
channel = %channel,
receivers = receiver_count,
bytes = payload.len(),
"published email message"
);
Ok(())
}
pub async fn publish_to<T: Serialize + ?Sized>(
&self,
logical: &str,
msg: &T,
) -> Result<(), EmailError> {
let channel = self.channel(logical);
self.publish(&channel, msg).await
}
pub async fn subscribe<T, F, Fut>(
&self,
logical: &str,
group_id: &str,
handler: F,
) -> Result<SubscribeHandle<T>, EmailError>
where
T: DeserializeOwned + Send + Clone + 'static,
F: Fn(T, TypedMessageMeta) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), EmailError>> + Send + 'static,
{
let _ = group_id;
let channel = self.channel(logical);
let max_retries = self.config.max_retries;
let retry_backoff = self.config.retry_backoff;
let cancel = CancellationToken::new();
let cancel_child = cancel.clone();
let channel_for_log = channel.clone();
let mut pubsub = self.pubsub_client.get_async_pubsub().await?;
pubsub.subscribe(&channel).await?;
let join: JoinHandle<Result<(), EmailError>> = tokio::spawn(async move {
let mut stream = pubsub.on_message();
loop {
tokio::select! {
_ = cancel_child.cancelled() => {
debug!(channel = %channel_for_log, "subscriber cancelled");
return Ok(());
}
msg = futures::StreamExt::next(&mut stream) => {
let Some(msg) = msg else {
warn!(channel = %channel_for_log, "pubsub stream ended");
return Ok(());
};
let raw: String = match msg.get_payload() {
Ok(s) => s,
Err(err) => {
warn!(channel = %channel_for_log, error = %err, "invalid pubsub payload type");
continue;
}
};
let meta = TypedMessageMeta {
channel: channel_for_log.clone(),
raw: raw.clone(),
};
let value: T = match serde_json::from_str(&raw) {
Ok(v) => v,
Err(err) => {
warn!(channel = %channel_for_log, error = %err, "message deserialize failed, dropping");
continue;
}
};
run_with_retries(
&channel_for_log,
value,
meta,
&handler,
max_retries,
retry_backoff,
&cancel_child,
)
.await;
}
}
}
});
Ok(SubscribeHandle {
cancel,
join: Some(Arc::new(tokio::sync::Mutex::new(join))),
_marker: PhantomData,
})
}
}
#[derive(Clone)]
pub struct TypedMessageMeta {
pub channel: String,
pub raw: String,
}
/// The join handle of a running subscriber loop, shared so `shutdown` can await
/// it from outside the task.
type SharedJoinHandle = Arc<tokio::sync::Mutex<JoinHandle<Result<(), EmailError>>>>;
pub struct SubscribeHandle<T> {
cancel: CancellationToken,
join: Option<SharedJoinHandle>,
_marker: PhantomData<T>,
}
impl<T> SubscribeHandle<T> {
pub async fn shutdown(mut self, timeout: Duration) -> Result<(), EmailError> {
self.cancel.cancel();
if let Some(handle) = self.join.take() {
let join = {
let mut guard = handle.lock().await;
std::mem::replace(&mut *guard, tokio::spawn(async { Ok(()) }))
};
match tokio::time::timeout(timeout, join).await {
Ok(Ok(res)) => res,
Ok(Err(join_err)) => {
error!(error = %join_err, "subscriber task join failed");
Err(EmailError::Other(format!(
"subscriber join failed: {join_err}"
)))
}
Err(_) => {
warn!(
timeout_secs = timeout.as_secs(),
"subscriber shutdown timed out"
);
Err(EmailError::Cancelled)
}
}
} else {
Ok(())
}
}
}
impl<T> Drop for SubscribeHandle<T> {
fn drop(&mut self) {
self.cancel.cancel();
}
}
async fn open_connection(url: &str) -> Result<QueueConnection, EmailError> {
if url.split(',').count() > 1 {
let nodes: Vec<String> = url
.split(',')
.map(str::trim)
.filter(|s| !s.is_empty())
.map(String::from)
.collect();
let client = redis::cluster::ClusterClient::new(nodes)?;
let conn = client.get_async_connection().await?;
Ok(QueueConnection::Cluster(conn))
} else {
let client = redis::Client::open(url)?;
let manager = ConnectionManager::new(client).await?;
Ok(QueueConnection::Standalone(manager))
}
}
fn open_pubsub_client(url: &str) -> Result<redis::Client, EmailError> {
if url.split(',').count() > 1 {
let first = url
.split(',')
.map(str::trim)
.find(|s| !s.is_empty())
.ok_or_else(|| EmailError::Config("empty redis url".into()))?;
Ok(redis::Client::open(first)?)
} else {
Ok(redis::Client::open(url)?)
}
}
async fn run_with_retries<T, F, Fut>(
channel: &str,
value: T,
meta: TypedMessageMeta,
handler: &F,
max_retries: u32,
retry_backoff: Duration,
cancel: &CancellationToken,
) where
T: Send + Clone,
F: Fn(T, TypedMessageMeta) -> Fut + Send + Sync,
Fut: Future<Output = Result<(), EmailError>> + Send,
{
let mut attempt: u32 = 0;
loop {
attempt += 1;
match handler(value.clone(), meta.clone()).await {
Ok(()) => {
debug!(channel = %channel, attempt, "message processed");
return;
}
Err(err) if attempt >= max_retries => {
error!(
channel = %channel,
attempt,
error = %err,
"handler exhausted retries, dropping message"
);
return;
}
Err(err) => {
warn!(
channel = %channel,
attempt,
error = %err,
"handler failed, will retry"
);
tokio::select! {
_ = cancel.cancelled() => return,
_ = tokio::time::sleep(retry_backoff) => {}
}
}
}
}
}