asynq_rs/controller/
redis.rs

1use crate::controller::redis::RedisClient::Single;
2use anyhow::{anyhow, Result};
3use gset::Getset;
4use hex;
5use r2d2::{Pool, PooledConnection};
6use redis::cluster::ClusterClient;
7use redis::{from_redis_value, Client, ConnectionLike, FromRedisValue, RedisError, RedisResult};
8use sha1::{Digest, Sha1};
9use std::collections::HashMap;
10use std::sync::{Arc, LazyLock};
11use std::time::Duration;
12use tokio::sync::{OnceCell, RwLock};
13
14#[derive(Getset, Debug, Clone)]
15pub struct RedisClientOpt {
16    // Network type to use, either tcp or unix.
17    // Default is tcp.
18
19    // Redis server address in "host:port" format.
20    #[getset(set, vis = "pub")]
21    pub addr: String,
22
23    pub cluster_addr: Vec<String>,
24
25    // Username to authenticate the current connection when Redis ACLs are used.
26    // See: https://redis.io/commands/auth.
27    // pub username: String,
28
29    // Password to authenticate the current connection.
30    // See: https://redis.io/commands/auth.
31    // pub password: String,
32
33    // Redis DB to select after connecting to a server.
34    // See: https://redis.io/commands/select.
35    // pub db: i8,
36
37    // Dial timeout for establishing new connections.
38    // Default is 24 hour.
39    pub max_lifetime: Duration,
40
41    // Timeout for socket reads.
42    // If timeout is reached, read commands will fail with a timeout error
43    // instead of blocking.
44    //
45    // Use value -1 for no timeout and 0 for default.
46    // Default is 1 hour.
47    pub idle_timeout: Duration,
48
49    // Timeout for socket writes.
50    // If timeout is reached, write commands will fail with a timeout error
51    // instead of blocking.
52    //
53    // Use value -1 for no timeout and 0 for default.
54    // Default is ReadTimout. 30s
55    pub connection_timeout: Duration,
56
57    pub connection_type: RdbType,
58
59    // Maximum number of socket connections.
60    // Default is 10 connections per every CPU as reported by runtime.NumCPU.
61    pub pool_size: u32,
62
63    pub min_pool_size: u32,
64}
65impl Default for RedisClientOpt {
66    fn default() -> Self {
67        Self {
68            addr: "redis://127.0.0.1:6379/12".to_string(),
69            cluster_addr: vec![],
70            // username: "".to_string(),
71            // password: "".to_string(),
72            // db: 0,
73            max_lifetime: Duration::from_secs(86400),
74            idle_timeout: Duration::from_secs(3600),
75            connection_timeout: Duration::from_secs(30),
76            connection_type: RdbType::RdbSingle,
77            pool_size: 10,
78            min_pool_size: 1,
79        }
80    }
81}
82
83#[derive(PartialEq, Debug, Clone)]
84pub enum RdbType {
85    RdbSingle,
86    RdbCluster,
87}
88
89pub static GLOBAL_REDIS_POOL: OnceCell<Pool<RedisConnectionManager>> = OnceCell::const_new();
90
91pub async fn init_global_redis(conf: RedisClientOpt) -> Result<()> {
92    let res = gen_redis_conn_pool(conf).await?;
93    GLOBAL_REDIS_POOL.get_or_init(|| async { res }).await;
94    Ok(())
95}
96
97/// 根据配置获取不同的链接
98pub async fn gen_redis_conn_pool(conf: RedisClientOpt) -> Result<Pool<RedisConnectionManager>> {
99    if conf.connection_type.ne(&RdbType::RdbSingle) && conf.connection_type.ne(&RdbType::RdbCluster)
100    {
101        return Err(anyhow::anyhow!("Redis connection type is error!"));
102    }
103
104    let manager: RedisConnectionManager;
105
106    if conf.connection_type == RdbType::RdbSingle {
107        manager = RedisConnectionManager {
108            redis_client: Single(get_single(&conf).await?),
109        };
110    } else if conf.connection_type == RdbType::RdbCluster {
111        manager = RedisConnectionManager {
112            redis_client: RedisClient::Cluster(get_cluster(conf.cluster_addr).await?),
113        };
114    } else {
115        return Err(anyhow::anyhow!("Redis connection type not supported!"));
116    }
117    Ok(Pool::builder()
118        .max_size(conf.pool_size)
119        .min_idle(Some(conf.min_pool_size))
120        .connection_timeout(conf.connection_timeout)
121        .max_lifetime(Some(conf.max_lifetime))
122        .idle_timeout(Some(conf.idle_timeout))
123        .build(manager)?)
124}
125
126/// get single redis client
127async fn get_single(conf: &RedisClientOpt) -> Result<Client> {
128    Ok(Client::open(&*conf.addr)?)
129}
130
131/// 获取 redis cluster 链接
132async fn get_cluster(conf: Vec<String>) -> Result<ClusterClient> {
133    Ok(ClusterClient::new(conf)?)
134}
135
136#[derive(Clone)]
137pub enum RedisClient {
138    Single(Client),
139    Cluster(ClusterClient),
140}
141
142impl RedisClient {
143    pub fn get_redis_connection(&self) -> RedisResult<RedisConnection> {
144        match self {
145            Single(s) => {
146                let conn = s.get_connection()?;
147                Ok(RedisConnection::Single(Box::new(conn)))
148            }
149            RedisClient::Cluster(c) => {
150                let conn = c.get_connection()?;
151                Ok(RedisConnection::Cluster(Box::new(conn)))
152            }
153        }
154    }
155}
156
157pub enum RedisConnection {
158    Single(Box<redis::Connection>),
159    Cluster(Box<redis::cluster::ClusterConnection>),
160}
161
162impl RedisConnection {
163    pub fn is_open(&self) -> bool {
164        match self {
165            RedisConnection::Single(sc) => sc.is_open(),
166            RedisConnection::Cluster(cc) => cc.is_open(),
167        }
168    }
169
170    pub fn query<T: FromRedisValue>(&mut self, cmd: &redis::Cmd) -> RedisResult<T> {
171        match self {
172            RedisConnection::Single(sc) => match sc.as_mut().req_command(cmd) {
173                Ok(val) => from_redis_value(&val),
174                Err(e) => Err(e),
175            },
176            RedisConnection::Cluster(cc) => match cc.req_command(cmd) {
177                Ok(val) => from_redis_value(&val),
178                Err(e) => Err(e),
179            },
180        }
181    }
182}
183
184#[derive(Clone)]
185pub struct RedisConnectionManager {
186    pub redis_client: RedisClient,
187}
188
189impl r2d2::ManageConnection for RedisConnectionManager {
190    type Connection = RedisConnection;
191    type Error = RedisError;
192
193    fn connect(&self) -> Result<RedisConnection, Self::Error> {
194        self.redis_client.get_redis_connection()
195    }
196
197    fn is_valid(&self, conn: &mut RedisConnection) -> Result<(), Self::Error> {
198        match conn {
199            RedisConnection::Single(sc) => {
200                let _: String = redis::cmd("PING").query(sc)?;
201            }
202            RedisConnection::Cluster(cc) => {
203                let _: String = redis::cmd("PING").query(cc)?;
204            }
205        }
206        Ok(())
207    }
208
209    fn has_broken(&self, conn: &mut RedisConnection) -> bool {
210        !conn.is_open()
211    }
212}
213
214///get_rdb
215pub async fn get_rdb() -> Result<PooledConnection<RedisConnectionManager>> {
216    Ok(match GLOBAL_REDIS_POOL.get() {
217        None => return Err(anyhow!("rdb client is None")),
218        Some(cli) => cli,
219    }
220    .get()?)
221}
222
223#[cfg(test)]
224mod tests {
225    use crate::common::func::{make_queue, make_task_key};
226    use crate::common::queue::PENDING;
227    use crate::common::rdb_cmd::RDB_ENQUEUE_CMD;
228    use crate::controller::redis::{init_global_redis, RedisClientOpt, GLOBAL_REDIS_POOL};
229    use crate::controller::task::Task;
230    use chrono::Utc;
231    use serde_json::json;
232    use uuid::Uuid;
233
234    #[tokio::test]
235    async fn test_redis_client_opt() {
236        let opt = RedisClientOpt::default();
237
238        let res = init_global_redis(opt).await;
239        println!("{:?}", res);
240        let mut redis = GLOBAL_REDIS_POOL.get().unwrap().get().unwrap();
241
242        let task = Task {
243            task_id: Uuid::new_v4().to_string(),
244            queue: "".to_string(),
245            payload: "".to_string(),
246            state: "".to_string(),
247            retry: 0,
248            retried: 0,
249            error_msg: "".to_string(),
250            last_failed_at: Default::default(),
251            completed_at: Default::default(),
252            todo_at: Default::default(),
253            task_life_time: 0,
254        };
255        let msg = json!(task);
256
257        let mut sql = redis::cmd("EVAL");
258        sql.arg(RDB_ENQUEUE_CMD)
259            .arg(2)
260            .arg(make_task_key(
261                task.queue.as_str(),
262                "t",
263                task.task_id.as_str(),
264            ))
265            .arg(make_queue(task.queue.as_str(), PENDING))
266            .arg(msg.to_string())
267            .arg(task.task_id)
268            .arg(Utc::now().timestamp());
269
270        let res: usize = redis.query(&sql).unwrap();
271
272        print!("{}", res);
273    }
274
275    #[tokio::test]
276    async fn test_lua_sha() {
277        let opt = RedisClientOpt::default();
278
279        let res = init_global_redis(opt).await;
280        println!("{:?}", res);
281
282        // let res = make_lua_sha(RDB_COMPLETED_CMD, "RDB_COMPLETED_CMD".to_string()).await;
283        println!("{:?}", res);
284    }
285}
286
287/////////////////////////////// lua script  /////////////////////////
288pub static GLOBAL_REDIS_SCRIPT: LazyLock<Arc<RwLock<HashMap<String, String>>>> =
289    LazyLock::new(|| Arc::new(RwLock::new(HashMap::with_capacity(4))));
290
291// 初始化函数，可以在程序启动时调用
292
293pub async fn make_lua_sha(
294    conn: &mut PooledConnection<RedisConnectionManager>,
295    cmd: &str,
296    key: String,
297) -> Result<()> {
298    // let mut conn = get_rdb().await?;
299    let sha = calculate_sha1(cmd);
300    let _: String = conn
301        .query(redis::cmd("EVAL").arg(cmd).arg(0))
302        .unwrap_or_default();
303    let res = {
304        let read_guard = GLOBAL_REDIS_SCRIPT.read().await;
305        match read_guard.get(&key) {
306            None => "".to_string(),
307            Some(res) => res.to_string(),
308        }
309    };
310    if res.is_empty() {
311        GLOBAL_REDIS_SCRIPT.write().await.insert(key, sha);
312    }
313    Ok(())
314}
315
316fn calculate_sha1(cmd: &str) -> String {
317    let mut hasher = Sha1::new();
318    hasher.update(cmd);
319    let result = hasher.finalize();
320    hex::encode(result)
321}
asynq_rs/controller/redis.rs

asynq_rs/controller/
redis.rs
