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
use cyfs_base::*;

use rusqlite::{Connection, OpenFlags};
use std::cell::RefCell;
use std::path::PathBuf;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use thread_local::ThreadLocal;

pub struct SqliteConnectionHolder {
    data_file: PathBuf,

    /* SQLite does not support multiple writers. */
    read_conn: ThreadLocal<RefCell<Connection>>,
    write_conn: ThreadLocal<RefCell<Connection>>,
    conn_rw_lock: RwLock<u32>,
}

impl SqliteConnectionHolder {
    pub fn new(data_file: PathBuf) -> Self {
        Self {
            data_file,
            read_conn: ThreadLocal::new(),
            write_conn: ThreadLocal::new(),
            conn_rw_lock: RwLock::new(0),
        }
    }

    pub fn get_write_conn(
        &self,
    ) -> BuckyResult<(std::cell::RefMut<Connection>, RwLockWriteGuard<u32>)> {
        let conn = self.write_conn.get_or_try(|| {
            let ret = self.create_new_conn(false)?;
            Ok::<RefCell<Connection>, BuckyError>(RefCell::new(ret))
        })?;

        let lock = self.conn_rw_lock.write().unwrap();
        Ok((conn.borrow_mut(), lock))
    }

    pub fn get_read_conn(&self) -> BuckyResult<(std::cell::Ref<Connection>, RwLockReadGuard<u32>)> {
        let conn = self.read_conn.get_or_try(|| {
            let ret = self.create_new_conn(true)?;
            Ok::<RefCell<Connection>, BuckyError>(RefCell::new(ret))
        })?;

        let lock = self.conn_rw_lock.read().unwrap();
        Ok((conn.borrow(), lock))
    }

    fn create_new_conn(&self, read_only: bool) -> BuckyResult<Connection> {
        let flags = if read_only {
            OpenFlags::SQLITE_OPEN_READ_ONLY
                | OpenFlags::SQLITE_OPEN_NO_MUTEX
                | OpenFlags::SQLITE_OPEN_URI
        } else {
            OpenFlags::default()
        };

        let conn = Connection::open_with_flags(&self.data_file, flags).map_err(|e| {
            let msg = format!("open db failed, db={}, {}", self.data_file.display(), e);
            error!("{}", msg);

            BuckyError::new(BuckyErrorCode::SqliteError, msg)
        })?;

        info!(
            "open db for thread={:?}, file={}",
            std::thread::current().id(),
            self.data_file.display()
        );
        assert!(conn.is_autocommit());

        // 设置一个30s的锁重试
        if let Err(e) = conn.busy_timeout(std::time::Duration::from_secs(30)) {
            error!("init sqlite busy_timeout error! {}", e);
        }

        Ok(conn)
    }
}