tracker-redis-exporter/src/client_manager.rs

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<Pool<Client>> {
    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<Client>;

pub struct ClientManager {
    client_map: RwLock<HashMap<TrackerRedisInfo, Pool<Client>>>,
    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<RedisConnection> {
        if let Some(redis) = &info.redis {
            self.get_connection(redis)
        } else {
            self.get_master_connection()
        }
    }

    pub fn get_master_connection(&self) -> color_eyre::Result<RedisConnection> {
        self.get_connection(&self.master_redis_info)
    }

    pub fn get_connection(&self, info: &TrackerRedisInfo) -> color_eyre::Result<RedisConnection> {
        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)
            }
        }
    }
}