use std::collections::HashMap; use parking_lot::RwLock; use r2d2::{Pool, PooledConnection}; use redis::Client; use crate::models::{TrackerInfo, TrackerRedisInfo}; fn pool_from_redisinfo(info: &TrackerRedisInfo) -> color_eyre::Result> { let url = if let Some(pass) = &info.pass { format!("redis://:{}@{}:{}", pass, info.host, info.port) } else { format!("redis://{}:{}", info.host, info.port) }; let new_client = Client::open(url)?; let max_pool_size = std::env::var("MAX_REDIS_POOL_SIZE") .unwrap_or_else(|_| "32".to_string()) .parse() .unwrap_or(32); let pool = Pool::builder() .max_size(max_pool_size) .build(new_client)?; Ok(pool) } pub type RedisConnection = PooledConnection; pub struct ClientManager { client_map: RwLock>>, master_redis_info: TrackerRedisInfo, } impl ClientManager { pub fn new(master_redis_info: TrackerRedisInfo) -> Self { Self { client_map: RwLock::new(HashMap::new()), master_redis_info, } } pub fn get_connection_from_tracker(&self, info: &TrackerInfo) -> color_eyre::Result { if let Some(redis) = &info.redis { self.get_connection(redis) } else { self.get_master_connection() } } pub fn get_master_connection(&self) -> color_eyre::Result { self.get_connection(&self.master_redis_info) } pub fn get_connection(&self, info: &TrackerRedisInfo) -> color_eyre::Result { let d = self.client_map.read(); if let Some(client) = d.get(info) { let con = client.get()?; Ok(con) } else { drop(d); let mut d = self.client_map.write(); if let Some(client) = d.get(info) { let con = client.get()?; Ok(con) } else { let new_client = pool_from_redisinfo(info)?; let con = new_client.get()?; d.insert(info.clone(), new_client); Ok(con) } } } }