modrinth/packages/xredis/src/pubsub.rs

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<Vec<u8>> {
        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<String>,
        channel: &'static str,
    ) -> mpsc::Receiver<Vec<u8>> {
        let (sender, receiver) = mpsc::channel(PUBSUB_BUFFER_SIZE);
        tokio::spawn(run_subscription(seed_urls, channel, sender));
        receiver
    }

    pub async fn publish<M>(&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<String>,
    channel: &'static str,
    sender: mpsc::Sender<Vec<u8>>,
) {
    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<Vec<u8>>,
) -> Result<SubscriptionOutcome> {
    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)
}