vsdb_core/common/
mod.rs

1//!
2//! # Common components
3//!
4
5pub(crate) mod engines;
6
7use engines::Engine;
8use parking_lot::Mutex;
9use ruc::*;
10use std::{
11    env, fs,
12    mem::size_of,
13    path::{Path, PathBuf},
14    sync::{
15        LazyLock,
16        atomic::{AtomicBool, Ordering},
17    },
18};
19use threadpool::ThreadPool;
20
21/////////////////////////////////////////////////////////////////////////////
22/////////////////////////////////////////////////////////////////////////////
23
24pub const NULL: &[u8] = &[];
25
26pub type RawBytes = Vec<u8>;
27pub type RawKey = RawBytes;
28pub type RawValue = RawBytes;
29
30pub type Pre = u64;
31pub const PREFIX_SIZE: usize = size_of::<Pre>();
32pub type PreBytes = [u8; PREFIX_SIZE];
33
34pub const KB: u64 = 1 << 10;
35pub const MB: u64 = 1 << 20;
36pub const GB: u64 = 1 << 30;
37
38const RESERVED_ID_CNT: Pre = 4096_0000;
39pub const BIGGEST_RESERVED_ID: Pre = RESERVED_ID_CNT - 1;
40
41/////////////////////////////////////////////////////////////////////////////
42/////////////////////////////////////////////////////////////////////////////
43
44const BASE_DIR_VAR: &str = "VSDB_BASE_DIR";
45
46static VSDB_BASE_DIR: LazyLock<Mutex<PathBuf>> =
47    LazyLock::new(|| Mutex::new(gen_data_dir()));
48
49static VSDB_CUSTOM_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
50    let mut d = VSDB_BASE_DIR.lock().clone();
51    d.push("__CUSTOM__");
52    pnk!(fs::create_dir_all(&d));
53    unsafe { env::set_var("VSDB_CUSTOM_DIR", d.as_os_str()) }
54    d
55});
56
57#[cfg(feature = "rocks_backend")]
58pub static VSDB: LazyLock<VsDB<engines::RocksDB>> = LazyLock::new(|| pnk!(VsDB::new()));
59
60#[cfg(feature = "parity_backend")]
61pub static VSDB: LazyLock<VsDB<engines::ParityDB>> = LazyLock::new(|| pnk!(VsDB::new()));
62
63/// Clean orphan instances in background.
64pub static TRASH_CLEANER: LazyLock<Mutex<ThreadPool>> = LazyLock::new(|| {
65    let pool = threadpool::Builder::new()
66        .num_threads(1)
67        .thread_stack_size(512 * MB as usize) // use large stack size
68        .build();
69    Mutex::new(pool)
70});
71
72/////////////////////////////////////////////////////////////////////////////
73/////////////////////////////////////////////////////////////////////////////
74
75/// Parse bytes to a specified integer type.
76#[macro_export]
77macro_rules! parse_int {
78    ($bytes: expr, $ty: ty) => {{
79        let array: [u8; std::mem::size_of::<$ty>()] = $bytes[..].try_into().unwrap();
80        <$ty>::from_be_bytes(array)
81    }};
82}
83
84/// Parse bytes to a `Pre` type.
85#[macro_export]
86macro_rules! parse_prefix {
87    ($bytes: expr) => {
88        $crate::parse_int!($bytes, $crate::common::Pre)
89    };
90}
91
92/////////////////////////////////////////////////////////////////////////////
93/////////////////////////////////////////////////////////////////////////////
94
95pub struct VsDB<T: Engine> {
96    db: T,
97}
98
99impl<T: Engine> VsDB<T> {
100    #[inline(always)]
101    fn new() -> Result<Self> {
102        Ok(Self {
103            db: T::new().c(d!())?,
104        })
105    }
106
107    #[inline(always)]
108    fn flush(&self) {
109        self.db.flush()
110    }
111}
112
113/////////////////////////////////////////////////////////////////////////////
114/////////////////////////////////////////////////////////////////////////////
115
116#[inline(always)]
117fn gen_data_dir() -> PathBuf {
118    // Compatible with Windows OS?
119    let d = env::var(BASE_DIR_VAR)
120        .or_else(|_| env::var("HOME").map(|h| format!("{}/.vsdb", h)))
121        .unwrap_or_else(|_| "/tmp/.vsdb".to_owned());
122    pnk!(fs::create_dir_all(&d));
123    PathBuf::from(d)
124}
125
126/// ${VSDB_CUSTOM_DIR}
127#[inline(always)]
128pub fn vsdb_get_custom_dir() -> &'static Path {
129    VSDB_CUSTOM_DIR.as_path()
130}
131
132/// ${VSDB_BASE_DIR}
133#[inline(always)]
134pub fn vsdb_get_base_dir() -> PathBuf {
135    VSDB_BASE_DIR.lock().clone()
136}
137
138/// Set ${VSDB_BASE_DIR} manually.
139#[inline(always)]
140pub fn vsdb_set_base_dir(dir: impl AsRef<Path>) -> Result<()> {
141    static HAS_INITED: AtomicBool = AtomicBool::new(false);
142
143    if HAS_INITED.swap(true, Ordering::Relaxed) {
144        Err(eg!("VSDB has been initialized !!"))
145    } else {
146        unsafe { env::set_var(BASE_DIR_VAR, dir.as_ref().as_os_str()) }
147        *VSDB_BASE_DIR.lock() = dir.as_ref().to_path_buf();
148        Ok(())
149    }
150}
151
152/// Flush data to disk, may take a long time.
153#[inline(always)]
154pub fn vsdb_flush() {
155    VSDB.flush();
156}