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}