rskit_workload/
resources.rs

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

rskit_workload/
resources.rs
