rskit_workload/
resources.rs

1//! CPU and memory quantity parsing and formatting.
2//!
3//! Mirrors gokit's `workload` resource helpers. Memory parsing is delegated to
4//! the canonical [`rskit_util::bytes`] parser (binary suffixes `k`/`ki` …
5//! `t`/`ti`); CPU is expressed in cores or millicores and normalized to
6//! nanocores. Formatting produces the compact single-letter representation the
7//! workload vocabulary uses (`512m`, `2g`).
8
9use rskit_errors::{AppError, AppResult};
10
11const KIB: i64 = 1024;
12const MIB: i64 = 1024 * 1024;
13const GIB: i64 = 1024 * 1024 * 1024;
14const TIB: i64 = 1024 * 1024 * 1024 * 1024;
15const PIB: i64 = 1024 * 1024 * 1024 * 1024 * 1024;
16
17const NANOS_PER_CORE: f64 = 1e9;
18const NANOS_PER_MILLICORE: f64 = 1e6;
19
20/// Parse a human-readable memory string into bytes.
21///
22/// Supported suffixes (case-insensitive): `k`/`ki` (KiB), `m`/`mi` (MiB),
23/// `g`/`gi` (GiB), `t`/`ti` (TiB), `p`/`pi` (PiB). A bare number is treated as
24/// bytes. All units are binary (1024-based).
25///
26/// # Errors
27///
28/// Returns [`rskit_errors::ErrorCode::InvalidFormat`] when the string is empty,
29/// not a valid quantity, negative, or larger than [`i64::MAX`] bytes.
30pub fn parse_memory(s: &str) -> AppResult<i64> {
31    let bytes = rskit_util::bytes::parse_bytes(s).ok_or_else(|| {
32        AppError::invalid_format("memory", "quantity with optional binary suffix")
33    })?;
34    i64::try_from(bytes)
35        .map_err(|_| AppError::invalid_format("memory", "quantity within i64 range"))
36}
37
38/// Parse a human-readable CPU string into nanocores.
39///
40/// Supported formats (case-insensitive): cores (`"0.5"`, `"1"`) and millicores
41/// (`"500m"`).
42///
43/// # Errors
44///
45/// Returns [`rskit_errors::ErrorCode::InvalidFormat`] when the string is empty,
46/// not a valid number, negative, or at or above the [`i64::MAX`] nanocore
47/// boundary (values that would saturate a float-to-int cast).
48pub fn parse_cpu(s: &str) -> AppResult<i64> {
49    let lower = s.trim().to_lowercase();
50    if lower.is_empty() {
51        return Err(AppError::invalid_format("cpu", "non-empty quantity"));
52    }
53
54    let (value, scale) = lower.strip_suffix('m').map_or_else(
55        || (lower.as_str(), NANOS_PER_CORE),
56        |millis| (millis, NANOS_PER_MILLICORE),
57    );
58    let parsed: f64 = value
59        .parse()
60        .map_err(|_| AppError::invalid_format("cpu", "number in cores or millicores"))?;
61    if parsed < 0.0 || !parsed.is_finite() {
62        return Err(AppError::invalid_format(
63            "cpu",
64            "non-negative finite quantity",
65        ));
66    }
67    let nanocores = parsed * scale;
68    // `as i64` saturates silently, so reject out-of-range values before casting.
69    #[allow(clippy::cast_precision_loss)]
70    if nanocores >= i64::MAX as f64 {
71        return Err(AppError::invalid_format(
72            "cpu",
73            "quantity within i64 nanocore range",
74        ));
75    }
76    #[allow(clippy::cast_possible_truncation)]
77    // validated finite and in range; fractional nanocores are dropped by design
78    Ok(nanocores as i64)
79}
80
81/// Format a byte count as a human-readable memory string using binary suffixes.
82///
83/// Picks the largest binary unit that divides the value exactly so a
84/// non-negative result round-trips through [`parse_memory`]; a value that is not
85/// an exact multiple of any unit (e.g. `1536`) is rendered as raw bytes rather
86/// than truncated. Negative inputs (which [`parse_memory`] rejects) render as
87/// raw signed integers.
88#[must_use]
89pub fn format_memory(bytes: i64) -> String {
90    for (unit, suffix) in [(PIB, 'p'), (TIB, 't'), (GIB, 'g'), (MIB, 'm'), (KIB, 'k')] {
91        if bytes >= unit && bytes % unit == 0 {
92            return format!("{}{suffix}", bytes / unit);
93        }
94    }
95    bytes.to_string()
96}
97
98/// Format a nanocore count as a human-readable CPU string.
99///
100/// Whole cores and exact millicores use the compact `N`/`Nm` forms; any other
101/// value renders its exact nanocore remainder as fractional cores (up to 9
102/// decimals, trailing zeros trimmed), preserving sign, so a non-negative result
103/// round-trips through [`parse_cpu`] instead of collapsing to a truncated
104/// `0.000`. [`parse_cpu`] rejects negative inputs.
105#[must_use]
106pub fn format_cpu(nanocores: i64) -> String {
107    if nanocores % 1_000_000_000 == 0 {
108        return (nanocores / 1_000_000_000).to_string();
109    }
110    if nanocores % 1_000_000 == 0 {
111        return format!("{}m", nanocores / 1_000_000);
112    }
113    let sign = if nanocores.is_negative() { "-" } else { "" };
114    let magnitude = nanocores.unsigned_abs();
115    let whole = magnitude / 1_000_000_000;
116    let frac = magnitude % 1_000_000_000;
117    let decimals = format!("{frac:09}");
118    format!("{sign}{whole}.{}", decimals.trim_end_matches('0'))
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use rskit_errors::ErrorCode;
125
126    #[test]
127    fn parse_memory_handles_binary_suffixes() {
128        assert_eq!(parse_memory("512").unwrap(), 512);
129        assert_eq!(parse_memory("1k").unwrap(), 1024);
130        assert_eq!(parse_memory("2Mi").unwrap(), 2 * MIB);
131        assert_eq!(parse_memory("1g").unwrap(), GIB);
132        assert_eq!(parse_memory(" 1Ti ").unwrap(), 1024_i64.pow(4));
133    }
134
135    #[test]
136    fn parse_memory_rejects_bad_input() {
137        assert_eq!(
138            parse_memory("").unwrap_err().code(),
139            ErrorCode::InvalidFormat
140        );
141        assert_eq!(
142            parse_memory("abc").unwrap_err().code(),
143            ErrorCode::InvalidFormat
144        );
145        assert_eq!(
146            parse_memory("-5").unwrap_err().code(),
147            ErrorCode::InvalidFormat
148        );
149    }
150
151    #[test]
152    fn parse_memory_rejects_overflow() {
153        assert_eq!(
154            parse_memory("9223372036854775807t").unwrap_err().code(),
155            ErrorCode::InvalidFormat
156        );
157    }
158
159    #[test]
160    fn parse_cpu_handles_cores_and_millicores() {
161        assert_eq!(parse_cpu("1").unwrap(), 1_000_000_000);
162        assert_eq!(parse_cpu("0.5").unwrap(), 500_000_000);
163        assert_eq!(parse_cpu("500m").unwrap(), 500_000_000);
164    }
165
166    #[test]
167    fn parse_cpu_rejects_bad_input() {
168        assert_eq!(parse_cpu("").unwrap_err().code(), ErrorCode::InvalidFormat);
169        assert_eq!(
170            parse_cpu("fast").unwrap_err().code(),
171            ErrorCode::InvalidFormat
172        );
173        assert_eq!(
174            parse_cpu("-1").unwrap_err().code(),
175            ErrorCode::InvalidFormat
176        );
177    }
178
179    #[test]
180    fn parse_cpu_rejects_out_of_range() {
181        assert_eq!(
182            parse_cpu("1e30").unwrap_err().code(),
183            ErrorCode::InvalidFormat
184        );
185    }
186
187    #[test]
188    fn format_memory_uses_largest_binary_unit() {
189        assert_eq!(format_memory(512), "512");
190        assert_eq!(format_memory(2048), "2k");
191        assert_eq!(format_memory(3 * MIB), "3m");
192        assert_eq!(format_memory(4 * GIB), "4g");
193        assert_eq!(format_memory(5 * TIB), "5t");
194        assert_eq!(format_memory(6 * PIB), "6p");
195        assert_eq!(format_memory(TIB), "1t");
196        // Non-multiples fall back to raw bytes rather than truncating.
197        assert_eq!(format_memory(1536), "1536");
198        assert_eq!(format_memory(GIB + MIB), "1025m");
199    }
200
201    #[test]
202    fn memory_round_trips_through_format_and_parse() {
203        for bytes in [
204            512,
205            1536,
206            2048,
207            GIB + MIB,
208            3 * MIB,
209            4 * GIB,
210            5 * TIB,
211            6 * PIB,
212        ] {
213            let formatted = format_memory(bytes);
214            assert_eq!(parse_memory(&formatted).unwrap(), bytes, "bytes {bytes}");
215        }
216    }
217
218    #[test]
219    fn format_cpu_prefers_whole_cores_then_millicores() {
220        assert_eq!(format_cpu(1_000_000_000), "1");
221        assert_eq!(format_cpu(500_000_000), "500m");
222        assert_eq!(format_cpu(2_000_000), "2m");
223    }
224
225    #[test]
226    fn format_cpu_falls_back_to_fractional_cores() {
227        assert_eq!(format_cpu(1_500_000), "0.0015");
228        assert_eq!(format_cpu(500), "0.0000005");
229        assert_eq!(format_cpu(1_000_500_000), "1.0005");
230        // Sub-core negative values keep their sign.
231        assert_eq!(format_cpu(-500), "-0.0000005");
232    }
233
234    #[test]
235    fn format_cpu_preserves_sub_millicore_precision() {
236        // A value below the millicore grid must not collapse to "0" on round-trip.
237        for nanos in [500_i64, 1_500_000, 1_250_000_000, 750, 333_000_000] {
238            let formatted = format_cpu(nanos);
239            assert_eq!(parse_cpu(&formatted).unwrap(), nanos, "nanos {nanos}");
240        }
241    }
242
243    #[test]
244    fn cpu_round_trips_through_format_and_parse() {
245        for input in ["1", "0.5", "500m", "2"] {
246            let nanos = parse_cpu(input).unwrap();
247            let formatted = format_cpu(nanos);
248            assert_eq!(parse_cpu(&formatted).unwrap(), nanos, "input {input}");
249        }
250    }
251}
rskit_workload/resources.rs

rskit_workload/
resources.rs
