1use std::{fmt, str::FromStr, time::Duration};
4
5use crate::version::default_client_version;
6use clap::{
7 builder::{PossibleValue, TypedValueParser},
8 error::ErrorKind,
9 Arg, Args, Command, Error,
10};
11use reth_db::{mdbx::MaxReadTransactionDuration, ClientVersion};
12use reth_storage_errors::db::LogLevel;
13
14#[derive(Debug, Args, PartialEq, Eq, Default, Clone, Copy)]
16#[command(next_help_heading = "Database")]
17pub struct DatabaseArgs {
18 #[arg(long = "db.log-level", value_parser = LogLevelValueParser::default())]
20 pub log_level: Option<LogLevel>,
21 #[arg(long = "db.exclusive")]
24 pub exclusive: Option<bool>,
25 #[arg(long = "db.max-size", value_parser = parse_byte_size)]
27 pub max_size: Option<usize>,
28 #[arg(long = "db.growth-step", value_parser = parse_byte_size)]
30 pub growth_step: Option<usize>,
31 #[arg(long = "db.read-transaction-timeout")]
33 pub read_transaction_timeout: Option<u64>,
34 #[arg(long = "db.max-readers")]
36 pub max_readers: Option<u64>,
37}
38
39impl DatabaseArgs {
40 pub fn database_args(&self) -> reth_db::mdbx::DatabaseArguments {
42 self.get_database_args(default_client_version())
43 }
44
45 pub fn get_database_args(
48 &self,
49 client_version: ClientVersion,
50 ) -> reth_db::mdbx::DatabaseArguments {
51 let max_read_transaction_duration = match self.read_transaction_timeout {
52 None => None, Some(0) => Some(MaxReadTransactionDuration::Unbounded), Some(secs) => Some(MaxReadTransactionDuration::Set(Duration::from_secs(secs))),
55 };
56
57 reth_db::mdbx::DatabaseArguments::new(client_version)
58 .with_log_level(self.log_level)
59 .with_exclusive(self.exclusive)
60 .with_max_read_transaction_duration(max_read_transaction_duration)
61 .with_geometry_max_size(self.max_size)
62 .with_growth_step(self.growth_step)
63 .with_max_readers(self.max_readers)
64 }
65}
66
67#[derive(Clone, Debug, Default)]
69#[non_exhaustive]
70struct LogLevelValueParser;
71
72impl TypedValueParser for LogLevelValueParser {
73 type Value = LogLevel;
74
75 fn parse_ref(
76 &self,
77 _cmd: &Command,
78 arg: Option<&Arg>,
79 value: &std::ffi::OsStr,
80 ) -> Result<Self::Value, Error> {
81 let val =
82 value.to_str().ok_or_else(|| Error::raw(ErrorKind::InvalidUtf8, "Invalid UTF-8"))?;
83
84 val.parse::<LogLevel>().map_err(|err| {
85 let arg = arg.map(|a| a.to_string()).unwrap_or_else(|| "...".to_owned());
86 let possible_values = LogLevel::value_variants()
87 .iter()
88 .map(|v| format!("- {:?}: {}", v, v.help_message()))
89 .collect::<Vec<_>>()
90 .join("\n");
91 let msg = format!(
92 "Invalid value '{val}' for {arg}: {err}.\n Possible values:\n{possible_values}"
93 );
94 clap::Error::raw(clap::error::ErrorKind::InvalidValue, msg)
95 })
96 }
97
98 fn possible_values(&self) -> Option<Box<dyn Iterator<Item = PossibleValue> + '_>> {
99 let values = LogLevel::value_variants()
100 .iter()
101 .map(|v| PossibleValue::new(v.variant_name()).help(v.help_message()));
102 Some(Box::new(values))
103 }
104}
105
106#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
108pub struct ByteSize(pub usize);
109
110impl From<ByteSize> for usize {
111 fn from(s: ByteSize) -> Self {
112 s.0
113 }
114}
115
116impl FromStr for ByteSize {
117 type Err = String;
118
119 fn from_str(s: &str) -> Result<Self, Self::Err> {
120 let s = s.trim().to_uppercase();
121 let parts: Vec<&str> = s.split_whitespace().collect();
122
123 let (num_str, unit) = match parts.len() {
124 1 => {
125 let (num, unit) =
126 s.split_at(s.find(|c: char| c.is_alphabetic()).unwrap_or(s.len()));
127 (num, unit)
128 }
129 2 => (parts[0], parts[1]),
130 _ => {
131 return Err("Invalid format. Use '<number><unit>' or '<number> <unit>'.".to_string())
132 }
133 };
134
135 let num: usize = num_str.parse().map_err(|_| "Invalid number".to_string())?;
136
137 let multiplier = match unit {
138 "B" | "" => 1, "KB" => 1024,
140 "MB" => 1024 * 1024,
141 "GB" => 1024 * 1024 * 1024,
142 "TB" => 1024 * 1024 * 1024 * 1024,
143 _ => return Err(format!("Invalid unit: {unit}. Use B, KB, MB, GB, or TB.")),
144 };
145
146 Ok(Self(num * multiplier))
147 }
148}
149
150impl fmt::Display for ByteSize {
151 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152 const KB: usize = 1024;
153 const MB: usize = KB * 1024;
154 const GB: usize = MB * 1024;
155 const TB: usize = GB * 1024;
156
157 let (size, unit) = if self.0 >= TB {
158 (self.0 as f64 / TB as f64, "TB")
159 } else if self.0 >= GB {
160 (self.0 as f64 / GB as f64, "GB")
161 } else if self.0 >= MB {
162 (self.0 as f64 / MB as f64, "MB")
163 } else if self.0 >= KB {
164 (self.0 as f64 / KB as f64, "KB")
165 } else {
166 (self.0 as f64, "B")
167 };
168
169 write!(f, "{size:.2}{unit}")
170 }
171}
172
173fn parse_byte_size(s: &str) -> Result<usize, String> {
175 s.parse::<ByteSize>().map(Into::into)
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181 use clap::Parser;
182 use reth_db::mdbx::{GIGABYTE, KILOBYTE, MEGABYTE, TERABYTE};
183
184 #[derive(Parser)]
186 struct CommandParser<T: Args> {
187 #[command(flatten)]
188 args: T,
189 }
190
191 #[test]
192 fn test_default_database_args() {
193 let default_args = DatabaseArgs::default();
194 let args = CommandParser::<DatabaseArgs>::parse_from(["reth"]).args;
195 assert_eq!(args, default_args);
196 }
197
198 #[test]
199 fn test_command_parser_with_valid_max_size() {
200 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
201 "reth",
202 "--db.max-size",
203 "4398046511104",
204 ])
205 .unwrap();
206 assert_eq!(cmd.args.max_size, Some(TERABYTE * 4));
207 }
208
209 #[test]
210 fn test_command_parser_with_invalid_max_size() {
211 let result =
212 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.max-size", "invalid"]);
213 assert!(result.is_err());
214 }
215
216 #[test]
217 fn test_command_parser_with_valid_growth_step() {
218 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
219 "reth",
220 "--db.growth-step",
221 "4294967296",
222 ])
223 .unwrap();
224 assert_eq!(cmd.args.growth_step, Some(GIGABYTE * 4));
225 }
226
227 #[test]
228 fn test_command_parser_with_invalid_growth_step() {
229 let result =
230 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.growth-step", "invalid"]);
231 assert!(result.is_err());
232 }
233
234 #[test]
235 fn test_command_parser_with_valid_max_size_and_growth_step_from_str() {
236 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
237 "reth",
238 "--db.max-size",
239 "2TB",
240 "--db.growth-step",
241 "1GB",
242 ])
243 .unwrap();
244 assert_eq!(cmd.args.max_size, Some(TERABYTE * 2));
245 assert_eq!(cmd.args.growth_step, Some(GIGABYTE));
246
247 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
248 "reth",
249 "--db.max-size",
250 "12MB",
251 "--db.growth-step",
252 "2KB",
253 ])
254 .unwrap();
255 assert_eq!(cmd.args.max_size, Some(MEGABYTE * 12));
256 assert_eq!(cmd.args.growth_step, Some(KILOBYTE * 2));
257
258 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
260 "reth",
261 "--db.max-size",
262 "12 MB",
263 "--db.growth-step",
264 "2 KB",
265 ])
266 .unwrap();
267 assert_eq!(cmd.args.max_size, Some(MEGABYTE * 12));
268 assert_eq!(cmd.args.growth_step, Some(KILOBYTE * 2));
269
270 let cmd = CommandParser::<DatabaseArgs>::try_parse_from([
271 "reth",
272 "--db.max-size",
273 "1073741824",
274 "--db.growth-step",
275 "1048576",
276 ])
277 .unwrap();
278 assert_eq!(cmd.args.max_size, Some(GIGABYTE));
279 assert_eq!(cmd.args.growth_step, Some(MEGABYTE));
280 }
281
282 #[test]
283 fn test_command_parser_max_size_and_growth_step_from_str_invalid_unit() {
284 let result =
285 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.growth-step", "1 PB"]);
286 assert!(result.is_err());
287
288 let result =
289 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.max-size", "2PB"]);
290 assert!(result.is_err());
291 }
292
293 #[test]
294 fn test_possible_values() {
295 let parser = LogLevelValueParser;
297
298 let possible_values: Vec<PossibleValue> = parser.possible_values().unwrap().collect();
300
301 let expected_values = vec![
303 PossibleValue::new("fatal")
304 .help("Enables logging for critical conditions, i.e. assertion failures"),
305 PossibleValue::new("error").help("Enables logging for error conditions"),
306 PossibleValue::new("warn").help("Enables logging for warning conditions"),
307 PossibleValue::new("notice")
308 .help("Enables logging for normal but significant condition"),
309 PossibleValue::new("verbose").help("Enables logging for verbose informational"),
310 PossibleValue::new("debug").help("Enables logging for debug-level messages"),
311 PossibleValue::new("trace").help("Enables logging for trace debug-level messages"),
312 PossibleValue::new("extra").help("Enables logging for extra debug-level messages"),
313 ];
314
315 assert_eq!(possible_values.len(), expected_values.len());
317 for (actual, expected) in possible_values.iter().zip(expected_values.iter()) {
318 assert_eq!(actual.get_name(), expected.get_name());
319 assert_eq!(actual.get_help(), expected.get_help());
320 }
321 }
322
323 #[test]
324 fn test_command_parser_with_valid_log_level() {
325 let cmd =
326 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.log-level", "Debug"])
327 .unwrap();
328 assert_eq!(cmd.args.log_level, Some(LogLevel::Debug));
329 }
330
331 #[test]
332 fn test_command_parser_with_invalid_log_level() {
333 let result =
334 CommandParser::<DatabaseArgs>::try_parse_from(["reth", "--db.log-level", "invalid"]);
335 assert!(result.is_err());
336 }
337
338 #[test]
339 fn test_command_parser_without_log_level() {
340 let cmd = CommandParser::<DatabaseArgs>::try_parse_from(["reth"]).unwrap();
341 assert_eq!(cmd.args.log_level, None);
342 }
343}