use async_trait::async_trait;
use redis::AsyncCommands;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tokio_stream::StreamExt;
use crate::errors::{Result, unavailable};
use crate::pubsub::{PubSub, PubSubHandler, match_topic};
pub struct RedisPubSub {
client: redis::Client,
subscriptions: Arc<RwLock<HashMap<String, Vec<PubSubHandler>>>>,
listeners: Arc<RwLock<HashMap<String, JoinHandle<()>>>>,
}
impl RedisPubSub {
pub async fn new(client: redis::Client) -> Result<Self> {
let mut conn = client
.get_multiplexed_async_connection()
.await
.map_err(|e| unavailable("redis", e.to_string()))?;
let _: String = redis::cmd("PING")
.query_async(&mut conn)
.await
.map_err(|e| unavailable("redis", e.to_string()))?;
Ok(Self {
client,
subscriptions: Arc::new(RwLock::new(HashMap::new())),
listeners: Arc::new(RwLock::new(HashMap::new())),
})
}
async fn ensure_listener(&self, pattern: String) -> Result<()> {
if self.listeners.read().await.contains_key(&pattern) {
return Ok(());
}
let client = self.client.clone();
let subscriptions = self.subscriptions.clone();
let listener_pattern = pattern.clone();
let handle = tokio::spawn(async move {
let Ok(mut pubsub) = client.get_async_pubsub().await else {
return;
};
if pubsub
.psubscribe(to_redis_pattern(&listener_pattern))
.await
.is_err()
{
return;
}
let mut stream = pubsub.on_message();
while let Some(msg) = stream.next().await {
let topic = msg.get_channel_name().to_owned();
let Ok(payload) = msg.get_payload::<Vec<u8>>() else {
continue;
};
let snapshot = subscriptions.read().await.clone();
for (local_pattern, handlers) in snapshot {
if match_topic(&local_pattern, &topic) {
for handler in handlers {
handler(topic.clone(), payload.clone());
}
}
}
}
});
let mut listeners = self.listeners.write().await;
if let Some(existing) = listeners.insert(pattern, handle) {
existing.abort();
}
Ok(())
}
}
#[async_trait]
impl PubSub for RedisPubSub {
async fn subscribe(&self, pattern: &str, handler: PubSubHandler) -> Result<()> {
let mut subs = self.subscriptions.write().await;
let first = !subs.contains_key(pattern);
subs.entry(pattern.to_owned()).or_default().push(handler);
drop(subs);
if first {
self.ensure_listener(pattern.to_owned()).await?;
}
Ok(())
}
async fn unsubscribe(&self, pattern: &str) -> Result<()> {
self.subscriptions.write().await.remove(pattern);
if let Some(handle) = self.listeners.write().await.remove(pattern) {
handle.abort();
}
Ok(())
}
async fn publish(&self, topic: &str, data: Vec<u8>) -> Result<()> {
let mut conn = self
.client
.get_multiplexed_async_connection()
.await
.map_err(|e| unavailable("redis", e.to_string()))?;
let _: usize = conn
.publish(topic, data)
.await
.map_err(|e| unavailable("redis", e.to_string()))?;
Ok(())
}
async fn close(&self) -> Result<()> {
self.subscriptions.write().await.clear();
for (_, handle) in self.listeners.write().await.drain() {
handle.abort();
}
Ok(())
}
}
fn to_redis_pattern(pattern: &str) -> String {
pattern
.strip_suffix(".*")
.map_or_else(|| pattern.to_owned(), |prefix| format!("{prefix}*"))
}