use std::time::Duration;
use rudb_common::{Error, Result};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Config {
memory_limit: Option<u64>,
threads: usize,
query_timeout: Option<Duration>,
}
impl Default for Config {
fn default() -> Self {
let threads = std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get);
Self { memory_limit: None, threads, query_timeout: None }
}
}
impl Config {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn memory_limit(&self) -> Option<u64> {
self.memory_limit
}
#[must_use]
pub fn threads(&self) -> usize {
self.threads
}
#[must_use]
pub fn query_timeout(&self) -> Option<Duration> {
self.query_timeout
}
#[must_use]
pub fn with_memory_limit(mut self, bytes: u64) -> Self {
self.memory_limit = Some(bytes);
self
}
#[must_use]
pub fn with_no_memory_limit(mut self) -> Self {
self.memory_limit = None;
self
}
pub fn with_memory_limit_text(mut self, text: &str) -> Result<Self> {
self.memory_limit = Some(parse_size(text)?);
Ok(self)
}
pub fn with_threads(mut self, threads: usize) -> Result<Self> {
if threads == 0 {
return Err(Error::invalid_input("threads must be at least 1"));
}
self.threads = threads;
Ok(self)
}
#[must_use]
pub fn with_query_timeout(mut self, timeout: Duration) -> Self {
self.query_timeout = Some(timeout);
self
}
#[must_use]
pub fn with_no_query_timeout(mut self) -> Self {
self.query_timeout = None;
self
}
#[must_use]
pub fn settings(&self) -> Vec<(&'static str, String)> {
vec![
(
"memory-limit",
self.memory_limit.map_or_else(|| "unlimited".to_string(), format_size),
),
("threads", self.threads.to_string()),
(
"query-timeout",
self.query_timeout.map_or_else(
|| "none".to_string(),
|timeout| format!("{}ms", timeout.as_millis()),
),
),
]
}
}
pub fn parse_size(text: &str) -> Result<u64> {
let text = text.trim();
let digits = text.trim_end_matches(|c: char| c.is_ascii_alphabetic() || c.is_whitespace());
let unit = text[digits.len()..].trim().to_ascii_uppercase();
let number: u64 = digits
.trim()
.parse()
.map_err(|_| Error::invalid_input(format!("\"{text}\" is not a size")))?;
let power = match unit.as_str() {
"" | "B" => 0,
"KB" | "KIB" => 1,
"MB" | "MIB" => 2,
"GB" | "GIB" => 3,
"TB" | "TIB" => 4,
other => {
return Err(Error::invalid_input(format!(
"\"{other}\" is not a unit, which is one of B, KB, MB, GB and TB"
)));
}
};
number
.checked_mul(1024u64.pow(power))
.ok_or_else(|| Error::invalid_input(format!("\"{text}\" is larger than a 64 bit size")))
}
fn format_size(bytes: u64) -> String {
for (power, unit) in [(4, "TB"), (3, "GB"), (2, "MB"), (1, "KB")] {
let scale = 1024u64.pow(power);
if bytes >= scale && bytes % scale == 0 {
return format!("{}{unit}", bytes / scale);
}
}
format!("{bytes}B")
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{Config, format_size, parse_size};
#[test]
fn the_defaults_are_no_limits_and_every_core() {
let config = Config::new();
assert_eq!(config.memory_limit(), None);
assert_eq!(config.query_timeout(), None);
assert!(config.threads() >= 1);
}
#[test]
fn a_setting_reads_back_as_what_it_was_set_to() {
let config = Config::new()
.with_memory_limit(1024)
.with_threads(4)
.expect("four is a thread count")
.with_query_timeout(Duration::from_secs(30));
assert_eq!(config.memory_limit(), Some(1024));
assert_eq!(config.threads(), 4);
assert_eq!(config.query_timeout(), Some(Duration::from_secs(30)));
}
#[test]
fn a_limit_can_be_taken_off_again() {
let config = Config::new()
.with_memory_limit(1024)
.with_no_memory_limit()
.with_query_timeout(Duration::from_secs(1))
.with_no_query_timeout();
assert_eq!(config.memory_limit(), None);
assert_eq!(config.query_timeout(), None);
}
#[test]
fn no_threads_at_all_is_refused_rather_than_silently_made_one() {
let error = Config::new().with_threads(0).expect_err("zero threads runs nothing");
assert!(error.to_string().contains("at least 1"), "{error}");
}
#[test]
fn a_size_with_a_unit_is_that_many_powers_of_1024() {
assert_eq!(parse_size("1024").expect("a number is bytes"), 1024);
assert_eq!(parse_size("1KB").expect("a kilobyte"), 1024);
assert_eq!(parse_size("1MB").expect("a megabyte"), 1024 * 1024);
assert_eq!(parse_size("10GB").expect("ten gigabytes"), 10 * 1024 * 1024 * 1024);
assert_eq!(parse_size("1TB").expect("a terabyte"), 1024u64.pow(4));
}
#[test]
fn the_two_spellings_of_a_gigabyte_are_the_same_number_here() {
assert_eq!(parse_size("1GB").expect("a gigabyte"), parse_size("1GiB").expect("a gibibyte"));
}
#[test]
fn case_and_a_space_before_the_unit_are_both_allowed() {
let expected = 512 * 1024 * 1024;
assert_eq!(parse_size("512mb").expect("lower case"), expected);
assert_eq!(parse_size("512 MB").expect("a space"), expected);
assert_eq!(parse_size(" 512MiB ").expect("surrounding space"), expected);
}
#[test]
fn something_that_is_not_a_size_says_so() {
assert!(parse_size("").is_err());
assert!(parse_size("lots").is_err());
assert!(parse_size("-1").is_err());
assert!(parse_size("1.5GB").is_err());
let error = parse_size("5PB").expect_err("petabytes are not a unit here");
assert!(error.to_string().contains("not a unit"), "{error}");
let error = parse_size("16777216TB").expect_err("that does not fit in a u64");
assert!(error.to_string().contains("64 bit"), "{error}");
}
#[test]
fn a_size_prints_back_as_the_unit_it_was_written_in() {
assert_eq!(format_size(1024), "1KB");
assert_eq!(format_size(10 * 1024 * 1024 * 1024), "10GB");
assert_eq!(format_size(0), "0B");
assert_eq!(format_size(1025), "1025B");
}
#[test]
fn the_settings_list_is_what_print_config_prints() {
let config = Config::new()
.with_memory_limit_text("2GB")
.expect("two gigabytes")
.with_threads(8)
.expect("eight is a thread count")
.with_query_timeout(Duration::from_millis(1500));
let settings = config.settings();
assert_eq!(settings[0], ("memory-limit", "2GB".to_string()));
assert_eq!(settings[1], ("threads", "8".to_string()));
assert_eq!(settings[2], ("query-timeout", "1500ms".to_string()));
}
#[test]
fn the_defaults_print_as_the_absence_of_a_limit_rather_than_as_a_number() {
let settings = Config::new().settings();
assert_eq!(settings[0].1, "unlimited");
assert_eq!(settings[2].1, "none");
}
}