use std::time::Duration;
use rudb_common::{Error, Result};
#[derive(Debug, Clone, Copy, 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: rudb_io::default_memory_limit(), 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: f64 =
digits.trim().parse().map_err(|_| Error::parser("Memory must have a number (e.g. 1GB)"))?;
let scale: f64 = match unit.as_str() {
"" | "B" => 1.0,
"KB" => 1e3,
"MB" => 1e6,
"GB" => 1e9,
"TB" => 1e12,
"KIB" => 1024.0,
"MIB" => 1024f64.powi(2),
"GIB" => 1024f64.powi(3),
"TIB" => 1024f64.powi(4),
other => {
let other = other.to_ascii_lowercase();
return Err(Error::parser(format!(
"Unknown unit for memory: '{other}' (expected: KB, MB, GB, TB for 1000^i units or KiB, MiB, GiB, TiB for 1024^i units)"
)));
}
};
let bytes = number * scale;
if !bytes.is_finite() || bytes < 0.0 {
return Err(Error::parser(format!("\"{text}\" is not a size")));
}
if bytes >= SIZE_CEILING {
return Err(Error::parser(format!("\"{text}\" is larger than a 64 bit size")));
}
#[expect(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
reason = "the range is checked on the line above and a size is a whole number of bytes"
)]
Ok(bytes as u64)
}
const SIZE_CEILING: f64 = 18_446_744_073_709_551_616.0;
fn format_size(bytes: u64) -> String {
for power in (1..=4).rev() {
for (scale, unit) in [(1024u64, "iB"), (1000u64, "B")] {
let Some(scale) = scale.checked_pow(power) else { continue };
if bytes >= scale && bytes % scale == 0 {
let prefix = ["K", "M", "G", "T"][power as usize - 1];
return format!("{}{prefix}{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_most_of_the_machine_and_every_core() {
let config = Config::new();
assert_eq!(config.memory_limit(), rudb_io::default_memory_limit());
assert_eq!(config.query_timeout(), None);
assert!(config.threads() >= 1);
}
#[test]
fn the_default_limit_leaves_the_machine_something() {
let Some(machine) = rudb_io::physical_memory() else {
eprintln!("skipping, this platform does not say how much memory it has");
return;
};
let limit = Config::new().memory_limit().expect("a machine that says its size has one");
assert!(limit < machine, "{limit} is not under {machine}");
assert!(limit > machine / 2, "{limit} is a smaller share of {machine} than intended");
}
#[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_decimal_unit_is_a_power_of_a_thousand_and_a_binary_one_a_power_of_1024() {
assert_eq!(parse_size("1024").expect("a number is bytes"), 1024);
assert_eq!(parse_size("1KB").expect("a kilobyte"), 1000);
assert_eq!(parse_size("1MB").expect("a megabyte"), 1_000_000);
assert_eq!(parse_size("10GB").expect("ten gigabytes"), 10_000_000_000);
assert_eq!(parse_size("1TB").expect("a terabyte"), 1_000_000_000_000);
assert_eq!(parse_size("1KiB").expect("a kibibyte"), 1024);
assert_eq!(parse_size("1MiB").expect("a mebibyte"), 1024 * 1024);
assert_eq!(parse_size("1GiB").expect("a gibibyte"), 1024u64.pow(3));
assert_eq!(parse_size("1TiB").expect("a tebibyte"), 1024u64.pow(4));
}
#[test]
fn the_two_spellings_of_a_gigabyte_are_two_different_numbers() {
let decimal = parse_size("1GB").expect("a gigabyte");
let binary = parse_size("1GiB").expect("a gibibyte");
assert_eq!(decimal, 1_000_000_000);
assert_eq!(binary, 1_073_741_824);
assert_eq!(rudb_common::human(decimal), "953.7 MiB");
}
#[test]
fn case_and_a_space_before_the_unit_are_both_allowed() {
assert_eq!(parse_size("512mb").expect("lower case"), 512_000_000);
assert_eq!(parse_size("512 MB").expect("a space"), 512_000_000);
assert_eq!(parse_size(" 512MiB ").expect("surrounding space"), 512 * 1024 * 1024);
assert_eq!(parse_size("512 gib").expect("both at once"), 512 * 1024u64.pow(3));
}
#[test]
fn a_fraction_is_a_size_because_duckdb_takes_one() {
assert_eq!(parse_size("1.5GB").expect("a gigabyte and a half"), 1_500_000_000);
assert_eq!(parse_size("0.5MiB").expect("half a mebibyte"), 512 * 1024);
}
#[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());
let error = parse_size("5PB").expect_err("petabytes are not a unit here");
assert!(error.to_string().contains("Unknown unit"), "{error}");
let error = parse_size("16777216TiB").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), "1KiB");
assert_eq!(format_size(10 * 1024u64.pow(3)), "10GiB");
assert_eq!(format_size(2_000_000_000), "2GB");
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 a_limit_taken_off_prints_as_the_absence_of_one_rather_than_as_a_number() {
let settings = Config::new().with_no_memory_limit().settings();
assert_eq!(settings[0].1, "unlimited");
assert_eq!(settings[2].1, "none");
}
#[test]
fn the_default_limit_prints_as_a_size() {
if rudb_io::physical_memory().is_none() {
eprintln!("skipping, this platform does not say how much memory it has");
return;
}
let printed = Config::new().settings()[0].1.clone();
assert_ne!(printed, "unlimited");
assert!(printed.ends_with('B'), "{printed}");
}
}