db_redis/
pubsub.rs

1use crate::connection::RedisConnection;
2use redis::{
3    Client, PushInfo, RedisResult, ToRedisArgs,
4    aio::{ConnectionManager, ConnectionManagerConfig},
5    cluster::ClusterClientBuilder,
6};
7use std::fmt::Display;
8
9pub struct RedisPubsub {
10    pub connection: RedisConnection,
11    pub stream: tokio::sync::mpsc::UnboundedReceiver<PushInfo>,
12    pub db_name: String,
13}
14
15impl RedisPubsub {
16    pub async fn init(uri: &str, db_name: &str) -> RedisResult<RedisPubsub> {
17        let nodes = uri.split(",").collect::<Vec<&str>>();
18        let (tx, stream) = tokio::sync::mpsc::unbounded_channel();
19        if nodes.len() == 1 {
20            let client = Client::open(uri)?;
21            let config = ConnectionManagerConfig::default()
22                .set_number_of_retries(usize::MAX)
23                .set_factor(3000)
24                .set_exponent_base(1)
25                .set_push_sender(tx)
26                .set_automatic_resubscription();
27            let connection = ConnectionManager::new_with_config(client.clone(), config).await?;
28            let connection = RedisConnection::Single(connection);
29            Ok(RedisPubsub {
30                connection,
31                stream,
32                db_name: db_name.to_string(),
33            })
34        } else {
35            let client = ClusterClientBuilder::new(nodes)
36                .use_protocol(redis::ProtocolVersion::RESP3)
37                .push_sender(tx)
38                .retries(u32::MAX)
39                .retry_wait_formula(3000, 1)
40                .build()?;
41            let connection = client.get_async_connection().await?;
42            let connection = RedisConnection::Cluster(connection);
43            Ok(RedisPubsub {
44                connection,
45                stream,
46                db_name: db_name.to_string(),
47            })
48        }
49    }
50
51    pub async fn subscribe(&mut self, channel_name: impl ToRedisArgs + Display) -> RedisResult<()> {
52        let channel_name = format!("{}:{}", self.db_name, channel_name);
53        match &mut self.connection {
54            RedisConnection::Single(c) => c.subscribe(channel_name).await,
55            RedisConnection::Cluster(c) => c.subscribe(channel_name).await,
56        }
57    }
58
59    pub async fn unsubscribe(
60        &mut self,
61        channel_name: impl ToRedisArgs + Display,
62    ) -> RedisResult<()> {
63        let channel_name = format!("{}:{}", self.db_name, channel_name);
64        match &mut self.connection {
65            RedisConnection::Single(c) => c.unsubscribe(channel_name).await,
66            RedisConnection::Cluster(c) => c.unsubscribe(channel_name).await,
67        }
68    }
69}
db_redis/pubsub.rs

db_redis/
pubsub.rs
