1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use std::collections::HashMap;
use super::DbMysql;
use std::sync::Mutex;

use time;
use mysql::{self, Opts};

static mut el: *mut DbPool = 0 as *mut _;

const MAX_KEEP_CONN: f64 = 3600f64;

pub struct DbPool {
    pub db_mysql: HashMap<String, Vec<DbMysql>>,
    pub db_info: HashMap<String, String>,
    pub mutex: Mutex<i32>,
}

pub trait PoolTrait: Sized {
    fn get_db_trait(pool: &mut DbPool, db_name: &String) -> Option<Self>;
    fn release_db_trait(pool: &mut DbPool, db_name: &String, db: Self);
    fn init_db_trait(pool: &mut DbPool, db_name: &String) -> Option<Self>;
}

impl DbPool {
    pub fn new() -> DbPool {
        DbPool {
            db_mysql: HashMap::new(),
            db_info: HashMap::new(),
            mutex: Mutex::new(0),
        }
    }

    pub fn instance() -> &'static mut DbPool {
        unsafe {
            if el == 0 as *mut _ {
                el = Box::into_raw(Box::new(DbPool::new()));
            }
            &mut *el
        }
    }

    pub fn set_db_info(&mut self, db_info: HashMap<String, String>) -> bool {
        self.db_info = db_info;
        true
    }

    pub fn check_connect_timeout(&mut self) {
        let _guard = self.mutex.lock().unwrap();
        let cur_time = time::precise_time_s();
        for (_, list) in self.db_mysql.iter_mut() {
            let val: Vec<DbMysql> = list.drain(..).collect();
            for v in val {
                if cur_time - v.last_use_time < MAX_KEEP_CONN {
                    list.push(v);
                }
            }
        }
    }
}

impl PoolTrait for DbMysql {
    fn get_db_trait(pool: &mut DbPool, db_name: &String) -> Option<DbMysql> {
        let db = {
            let _guard = pool.mutex.lock().unwrap();
            let mut list = match pool.db_mysql.contains_key(db_name) {
                true => pool.db_mysql.get_mut(db_name).unwrap(),
                false => {
                    pool.db_mysql.entry(db_name.to_string()).or_insert(vec![]);
                    pool.db_mysql.get_mut(db_name).unwrap()
                }
            };
            if list.is_empty() {
                None
            } else {
                list.pop()
            }
        };

        match db {
            Some(_) => db,
            None => PoolTrait::init_db_trait(pool, db_name),
        }
    }

    fn release_db_trait(pool: &mut DbPool, db_name: &String, mut db: DbMysql) {
        // db is lose connection, not need add to pool
        if !db.is_connect {
            return;
        }
        db.last_use_time = time::precise_time_s();
        let _guard = pool.mutex.lock().unwrap();
        let mut list = match pool.db_mysql.contains_key(db_name) {
            true => pool.db_mysql.get_mut(db_name).unwrap(),
            false => {
                pool.db_mysql.entry(db_name.to_string()).or_insert(vec![]);
                pool.db_mysql.get_mut(db_name).unwrap()
            }
        };
        list.push(db);
    }

    fn init_db_trait(pool: &mut DbPool, db_name: &String) -> Option<DbMysql> {
        let mut info = pool.db_info.get(&db_name.to_string());
        if info.is_none() {
            info = pool.db_info.get(&"mysql".to_string());
        }
        let info = unwrap_or!(info, return None);
        let mut opts: Opts = DbMysql::from_url_basic(&**info).unwrap();
        opts.db_name = Some(db_name.clone());
        let pool = unwrap_or!(mysql::Conn::new(opts).ok(), return None);
        Some(DbMysql::new(pool))
    }
}