use std::time::Duration; use eyre::{Result, WrapErr}; use futures::StreamExt; use redis::ToRedisArgs; use tokio::sync::mpsc; use tracing::{info, warn}; use super::RedisPool; const PUBSUB_BUFFER_SIZE: usize = 1024; const INITIAL_RECONNECT_BACKOFF: Duration = Duration::from_millis(250); const MAX_RECONNECT_BACKOFF: Duration = Duration::from_secs(30); enum SubscriptionOutcome { SubscriberClosed, Disconnected, } impl RedisPool { pub fn subscribe(&self, channel: &'static str) -> mpsc::Receiver> { let seed_urls = self.config.seed_urls().to_vec(); let (sender, receiver) = mpsc::channel(PUBSUB_BUFFER_SIZE); tokio::spawn(run_subscription(seed_urls, channel, sender)); receiver } #[doc(hidden)] #[cfg(feature = "test")] pub fn subscribe_with_seed_urls( seed_urls: Vec, channel: &'static str, ) -> mpsc::Receiver> { let (sender, receiver) = mpsc::channel(PUBSUB_BUFFER_SIZE); tokio::spawn(run_subscription(seed_urls, channel, sender)); receiver } pub async fn publish(&self, channel: &str, message: M) -> Result<()> where M: ToRedisArgs + Send + Sync, { let mut connection = self .connect() .await .wrap_err("connecting to Redis for publishing")?; let _: usize = redis::cmd("PUBLISH") .arg(channel) .arg(message) .query_async(&mut connection) .await .wrap_err("publishing to Redis channel")?; Ok(()) } } async fn run_subscription( seed_urls: Vec, channel: &'static str, sender: mpsc::Sender>, ) { let mut next_seed = 0; let mut reconnect_backoff = INITIAL_RECONNECT_BACKOFF; loop { let mut connected = false; for _ in 0..seed_urls.len() { let seed_url = &seed_urls[next_seed]; next_seed = (next_seed + 1) % seed_urls.len(); match forward_from_seed(seed_url, channel, &sender).await { Ok(SubscriptionOutcome::SubscriberClosed) => return, Ok(SubscriptionOutcome::Disconnected) => { warn!(channel, "Redis Pub/Sub connection disconnected"); connected = true; break; } Err(error) => { warn!( ?error, channel, "Failed to establish Redis Pub/Sub subscription" ); } } } if sender.is_closed() { return; } let delay = if connected { INITIAL_RECONNECT_BACKOFF } else { reconnect_backoff }; tokio::time::sleep(delay).await; reconnect_backoff = if connected { INITIAL_RECONNECT_BACKOFF } else { reconnect_backoff .saturating_mul(2) .min(MAX_RECONNECT_BACKOFF) }; } } async fn forward_from_seed( seed_url: &str, channel: &'static str, sender: &mpsc::Sender>, ) -> Result { let client = redis::Client::open(seed_url) .wrap_err("configuring Redis Pub/Sub client")?; let mut pubsub = client .get_async_pubsub() .await .wrap_err("connecting to Redis Pub/Sub")?; pubsub .subscribe(channel) .await .wrap_err("subscribing to Redis channel")?; info!(channel, "Established Redis Pub/Sub subscription"); let mut stream = pubsub.into_on_message(); while let Some(message) = stream.next().await { if message.get_channel_name() != channel { continue; } if sender .send(message.get_payload_bytes().to_vec()) .await .is_err() { return Ok(SubscriptionOutcome::SubscriberClosed); } } Ok(SubscriptionOutcome::Disconnected) }