kaspa_database/db/
conn_builder.rs

1use crate::db::DB;
2use rocksdb::{DBWithThreadMode, MultiThreaded};
3use std::{path::PathBuf, sync::Arc};
4
5#[derive(Debug)]
6pub struct Unspecified;
7
8#[derive(Debug)]
9pub struct ConnBuilder<Path, const STATS_ENABLED: bool, StatsPeriod, FDLimit> {
10    db_path: Path,
11    create_if_missing: bool,
12    parallelism: usize,
13    files_limit: FDLimit,
14    mem_budget: usize,
15    stats_period: StatsPeriod,
16}
17
18impl Default for ConnBuilder<Unspecified, false, Unspecified, Unspecified> {
19    fn default() -> Self {
20        ConnBuilder {
21            db_path: Unspecified,
22            create_if_missing: true,
23            parallelism: 1,
24            mem_budget: 64 * 1024 * 1024,
25            stats_period: Unspecified,
26            files_limit: Unspecified,
27        }
28    }
29}
30
31impl<Path, const STATS_ENABLED: bool, StatsPeriod, FDLimit> ConnBuilder<Path, STATS_ENABLED, StatsPeriod, FDLimit> {
32    pub fn with_db_path(self, db_path: PathBuf) -> ConnBuilder<PathBuf, STATS_ENABLED, StatsPeriod, FDLimit> {
33        ConnBuilder {
34            db_path,
35            files_limit: self.files_limit,
36            create_if_missing: self.create_if_missing,
37            parallelism: self.parallelism,
38            mem_budget: self.mem_budget,
39            stats_period: self.stats_period,
40        }
41    }
42    pub fn with_create_if_missing(self, create_if_missing: bool) -> ConnBuilder<Path, STATS_ENABLED, StatsPeriod, FDLimit> {
43        ConnBuilder { create_if_missing, ..self }
44    }
45    pub fn with_parallelism(self, parallelism: impl Into<usize>) -> ConnBuilder<Path, STATS_ENABLED, StatsPeriod, FDLimit> {
46        ConnBuilder { parallelism: parallelism.into(), ..self }
47    }
48    pub fn with_mem_budget(self, mem_budget: impl Into<usize>) -> ConnBuilder<Path, STATS_ENABLED, StatsPeriod, FDLimit> {
49        ConnBuilder { mem_budget: mem_budget.into(), ..self }
50    }
51    pub fn with_files_limit(self, files_limit: impl Into<i32>) -> ConnBuilder<Path, STATS_ENABLED, StatsPeriod, i32> {
52        ConnBuilder {
53            db_path: self.db_path,
54            files_limit: files_limit.into(),
55            create_if_missing: self.create_if_missing,
56            parallelism: self.parallelism,
57            mem_budget: self.mem_budget,
58            stats_period: self.stats_period,
59        }
60    }
61}
62
63impl<Path, FDLimit> ConnBuilder<Path, false, Unspecified, FDLimit> {
64    pub fn enable_stats(self) -> ConnBuilder<Path, true, Unspecified, FDLimit> {
65        ConnBuilder {
66            db_path: self.db_path,
67            create_if_missing: self.create_if_missing,
68            parallelism: self.parallelism,
69            files_limit: self.files_limit,
70            mem_budget: self.mem_budget,
71            stats_period: self.stats_period,
72        }
73    }
74}
75
76impl<Path, StatsPeriod, FDLimit> ConnBuilder<Path, true, StatsPeriod, FDLimit> {
77    pub fn disable_stats(self) -> ConnBuilder<Path, false, Unspecified, FDLimit> {
78        ConnBuilder {
79            db_path: self.db_path,
80            create_if_missing: self.create_if_missing,
81            parallelism: self.parallelism,
82            files_limit: self.files_limit,
83            mem_budget: self.mem_budget,
84            stats_period: Unspecified,
85        }
86    }
87    pub fn with_stats_period(self, stats_period: impl Into<u32>) -> ConnBuilder<Path, true, u32, FDLimit> {
88        ConnBuilder {
89            db_path: self.db_path,
90            create_if_missing: self.create_if_missing,
91            parallelism: self.parallelism,
92            files_limit: self.files_limit,
93            mem_budget: self.mem_budget,
94            stats_period: stats_period.into(),
95        }
96    }
97}
98
99macro_rules! default_opts {
100    ($self: expr) => {{
101        let mut opts = rocksdb::Options::default();
102        if $self.parallelism > 1 {
103            opts.increase_parallelism($self.parallelism as i32);
104        }
105
106        opts.optimize_level_style_compaction($self.mem_budget);
107        let guard = kaspa_utils::fd_budget::acquire_guard($self.files_limit)?;
108        opts.set_max_open_files($self.files_limit);
109        opts.create_if_missing($self.create_if_missing);
110        Ok((opts, guard))
111    }};
112}
113
114impl ConnBuilder<PathBuf, false, Unspecified, i32> {
115    pub fn build(self) -> Result<Arc<DB>, kaspa_utils::fd_budget::Error> {
116        let (opts, guard) = default_opts!(self)?;
117        let db = Arc::new(DB::new(<DBWithThreadMode<MultiThreaded>>::open(&opts, self.db_path.to_str().unwrap()).unwrap(), guard));
118        Ok(db)
119    }
120}
121
122impl ConnBuilder<PathBuf, true, Unspecified, i32> {
123    pub fn build(self) -> Result<Arc<DB>, kaspa_utils::fd_budget::Error> {
124        let (mut opts, guard) = default_opts!(self)?;
125        opts.enable_statistics();
126        let db = Arc::new(DB::new(<DBWithThreadMode<MultiThreaded>>::open(&opts, self.db_path.to_str().unwrap()).unwrap(), guard));
127        Ok(db)
128    }
129}
130
131impl ConnBuilder<PathBuf, true, u32, i32> {
132    pub fn build(self) -> Result<Arc<DB>, kaspa_utils::fd_budget::Error> {
133        let (mut opts, guard) = default_opts!(self)?;
134        opts.enable_statistics();
135        opts.set_report_bg_io_stats(true);
136        opts.set_stats_dump_period_sec(self.stats_period);
137        let db = Arc::new(DB::new(<DBWithThreadMode<MultiThreaded>>::open(&opts, self.db_path.to_str().unwrap()).unwrap(), guard));
138        Ok(db)
139    }
140}