sidekiq/
redis.rs

1use bb8::{CustomizeConnection, ManageConnection, Pool};
2use redis::aio::MultiplexedConnection as Connection;
3use redis::AsyncCommands;
4pub use redis::RedisError;
5use redis::ToRedisArgs;
6pub use redis::Value as RedisValue;
7use redis::{Client, IntoConnectionInfo};
8use std::future::Future;
9use std::ops::DerefMut;
10use std::pin::Pin;
11
12pub type RedisPool = Pool<RedisConnectionManager>;
13
14#[derive(Debug)]
15pub struct NamespaceCustomizer {
16    namespace: String,
17}
18
19impl CustomizeConnection<RedisConnection, RedisError> for NamespaceCustomizer {
20    fn on_acquire<'a>(
21        &'a self,
22        connection: &'a mut RedisConnection,
23    ) -> Pin<Box<dyn Future<Output = Result<(), RedisError>> + Send + 'a>> {
24        Box::pin(async {
25            // All redis operations used by the sidekiq lib will use this as a prefix.
26            connection.set_namespace(self.namespace.clone());
27
28            Ok(())
29        })
30    }
31}
32
33#[must_use]
34pub fn with_custom_namespace(namespace: String) -> Box<NamespaceCustomizer> {
35    Box::new(NamespaceCustomizer { namespace })
36}
37
38/// A `bb8::ManageConnection` for `redis::Client::get_async_connection` wrapped in a helper type
39/// for namespacing.
40#[derive(Clone, Debug)]
41pub struct RedisConnectionManager {
42    client: Client,
43}
44
45impl RedisConnectionManager {
46    /// Create a new `RedisConnectionManager`.
47    /// See `redis::Client::open` for a description of the parameter types.
48    pub fn new<T: IntoConnectionInfo>(info: T) -> Result<Self, RedisError> {
49        Ok(Self {
50            client: Client::open(info.into_connection_info()?)?,
51        })
52    }
53}
54
55impl ManageConnection for RedisConnectionManager {
56    type Connection = RedisConnection;
57    type Error = RedisError;
58
59    async fn connect(&self) -> Result<Self::Connection, Self::Error> {
60        // Disable response timeout because this connection is used for blocking
61        // commands like BRPOP which can legitimately block for seconds.
62        let config = redis::AsyncConnectionConfig::new().set_response_timeout(None);
63        Ok(RedisConnection::new(
64            self.client
65                .get_multiplexed_async_connection_with_config(&config)
66                .await?,
67        ))
68    }
69
70    async fn is_valid(&self, mut conn: &mut Self::Connection) -> Result<(), Self::Error> {
71        let pong: String = redis::cmd("PING")
72            .query_async(&mut conn.deref_mut().connection)
73            .await?;
74        match pong.as_str() {
75            "PONG" => Ok(()),
76            _ => Err(redis::RedisError::from((
77                redis::ErrorKind::Server(redis::ServerErrorKind::ResponseError),
78                "ping request",
79            ))),
80        }
81    }
82
83    fn has_broken(&self, _conn: &mut Self::Connection) -> bool {
84        false
85    }
86}
87
88/// A wrapper type for making the redis crate compatible with namespacing.
89pub struct RedisConnection {
90    connection: Connection,
91    namespace: Option<String>,
92}
93
94impl RedisConnection {
95    #[must_use]
96    pub fn new(connection: Connection) -> Self {
97        Self {
98            connection,
99            namespace: None,
100        }
101    }
102
103    pub fn set_namespace(&mut self, namespace: String) {
104        self.namespace = Some(namespace);
105    }
106
107    #[must_use]
108    pub fn with_namespace(self, namespace: String) -> Self {
109        Self {
110            connection: self.connection,
111            namespace: Some(namespace),
112        }
113    }
114
115    fn namespaced_key(&self, key: String) -> String {
116        if let Some(ref namespace) = self.namespace {
117            return format!("{namespace}:{key}");
118        }
119
120        key
121    }
122
123    fn namespaced_keys(&self, keys: Vec<String>) -> Vec<String> {
124        if let Some(ref namespace) = self.namespace {
125            let keys: Vec<String> = keys
126                .iter()
127                .map(|key| format!("{namespace}:{key}"))
128                .collect();
129
130            return keys;
131        }
132
133        keys
134    }
135
136    /// This allows you to borrow the raw redis connection without any namespacing support.
137    pub fn unnamespaced_borrow_mut(&mut self) -> &mut Connection {
138        &mut self.connection
139    }
140
141    pub async fn brpop(
142        &mut self,
143        keys: Vec<String>,
144        timeout: usize,
145    ) -> Result<Option<(String, String)>, RedisError> {
146        self.connection
147            .brpop(self.namespaced_keys(keys), timeout as f64)
148            .await
149    }
150
151    pub fn cmd_with_key(&mut self, cmd: &str, key: String) -> redis::Cmd {
152        let mut c = redis::cmd(cmd);
153        c.arg(self.namespaced_key(key));
154        c
155    }
156
157    pub async fn del(&mut self, key: String) -> Result<usize, RedisError> {
158        self.connection.del(self.namespaced_key(key)).await
159    }
160
161    pub async fn expire(&mut self, key: String, value: usize) -> Result<usize, RedisError> {
162        self.connection
163            .expire(self.namespaced_key(key), value as i64)
164            .await
165    }
166
167    pub async fn lpush<V>(&mut self, key: String, value: V) -> Result<(), RedisError>
168    where
169        V: ToRedisArgs + Send + Sync,
170    {
171        self.connection.lpush(self.namespaced_key(key), value).await
172    }
173
174    pub async fn sadd<V>(&mut self, key: String, value: V) -> Result<(), RedisError>
175    where
176        V: ToRedisArgs + Send + Sync,
177    {
178        self.connection.sadd(self.namespaced_key(key), value).await
179    }
180
181    pub async fn hset(
182        &mut self,
183        key: String,
184        field: String,
185        value: String,
186    ) -> Result<(), RedisError> {
187        self.connection
188            .hset(self.namespaced_key(key), field, value)
189            .await
190    }
191
192    pub async fn hdel(&mut self, key: String, field: String) -> Result<(), RedisError> {
193        self.connection.hdel(self.namespaced_key(key), field).await
194    }
195
196    pub async fn srem<V>(&mut self, key: String, value: V) -> Result<(), RedisError>
197    where
198        V: ToRedisArgs + Send + Sync,
199    {
200        self.connection.srem(self.namespaced_key(key), value).await
201    }
202
203    pub async fn unlink(&mut self, key: String) -> Result<(), RedisError> {
204        let _: i64 = self.connection.unlink(self.namespaced_key(key)).await?;
205        Ok(())
206    }
207
208    /// Pipeline SREM + UNLINK in a single round-trip. Both keys are namespaced.
209    /// Equivalent to Ruby's `conn.pipelined { |p| p.srem(set, member); p.unlink(hash_key) }`.
210    pub async fn srem_and_unlink(
211        &mut self,
212        set_key: String,
213        member: String,
214        hash_key: String,
215    ) -> Result<(), RedisError> {
216        redis::pipe()
217            .cmd("SREM")
218            .arg(self.namespaced_key(set_key))
219            .arg(&member)
220            .cmd("UNLINK")
221            .arg(self.namespaced_key(hash_key))
222            .query_async::<()>(self.unnamespaced_borrow_mut())
223            .await
224    }
225
226    pub async fn set_nx_ex<V>(
227        &mut self,
228        key: String,
229        value: V,
230        ttl_in_seconds: usize,
231    ) -> Result<RedisValue, RedisError>
232    where
233        V: ToRedisArgs + Send + Sync,
234    {
235        redis::cmd("SET")
236            .arg(self.namespaced_key(key))
237            .arg(value)
238            .arg("NX")
239            .arg("EX")
240            .arg(ttl_in_seconds)
241            .query_async(self.unnamespaced_borrow_mut())
242            .await
243    }
244
245    pub async fn zrange(
246        &mut self,
247        key: String,
248        lower: isize,
249        upper: isize,
250    ) -> Result<Vec<String>, RedisError> {
251        self.connection
252            .zrange(self.namespaced_key(key), lower, upper)
253            .await
254    }
255
256    pub async fn zrangebyscore_limit<
257        L: redis::ToSingleRedisArg + Send + Sync,
258        U: redis::ToSingleRedisArg + Sync + Send,
259    >(
260        &mut self,
261        key: String,
262        lower: L,
263        upper: U,
264        offset: isize,
265        limit: isize,
266    ) -> Result<Vec<String>, RedisError> {
267        self.connection
268            .zrangebyscore_limit(self.namespaced_key(key), lower, upper, offset, limit)
269            .await
270    }
271
272    pub async fn zadd<
273        V: redis::ToSingleRedisArg + Send + Sync,
274        S: redis::ToSingleRedisArg + Send + Sync,
275    >(
276        &mut self,
277        key: String,
278        value: V,
279        score: S,
280    ) -> Result<usize, RedisError> {
281        self.connection
282            .zadd(self.namespaced_key(key), value, score)
283            .await
284    }
285
286    pub async fn zadd_ch<V: ToRedisArgs + Send + Sync, S: ToRedisArgs + Send + Sync>(
287        &mut self,
288        key: String,
289        value: V,
290        score: S,
291    ) -> Result<bool, RedisError> {
292        redis::cmd("ZADD")
293            .arg(self.namespaced_key(key))
294            .arg("CH")
295            .arg(score)
296            .arg(value)
297            .query_async(self.unnamespaced_borrow_mut())
298            .await
299    }
300
301    pub async fn zrem<V>(&mut self, key: String, value: V) -> Result<usize, RedisError>
302    where
303        V: ToRedisArgs + Send + Sync,
304    {
305        self.connection.zrem(self.namespaced_key(key), value).await
306    }
307}
sidekiq/redis.rs

sidekiq/
redis.rs
