zombienet_configuration/shared/
resources.rs

1use std::error::Error;
2
3use lazy_static::lazy_static;
4use regex::Regex;
5use serde::{
6    de::{self},
7    ser::SerializeStruct,
8    Deserialize, Serialize,
9};
10use support::constants::{SHOULD_COMPILE, THIS_IS_A_BUG};
11
12use super::{
13    errors::{ConversionError, FieldError},
14    helpers::merge_errors,
15};
16
17/// A resource quantity used to define limits (k8s/podman only).
18/// It can be constructed from a `&str` or u64, if it fails, it returns a [`ConversionError`].
19/// Possible optional prefixes are: m, K, M, G, T, P, E, Ki, Mi, Gi, Ti, Pi, Ei
20///
21/// # Examples
22///
23/// ```
24/// use zombienet_configuration::shared::resources::ResourceQuantity;
25///
26/// let quantity1: ResourceQuantity = "100000".try_into().unwrap();
27/// let quantity2: ResourceQuantity = "1000m".try_into().unwrap();
28/// let quantity3: ResourceQuantity = "1Gi".try_into().unwrap();
29/// let quantity4: ResourceQuantity = 10_000.into();
30///
31/// assert_eq!(quantity1.as_str(), "100000");
32/// assert_eq!(quantity2.as_str(), "1000m");
33/// assert_eq!(quantity3.as_str(), "1Gi");
34/// assert_eq!(quantity4.as_str(), "10000");
35/// ```
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct ResourceQuantity(String);
38
39impl ResourceQuantity {
40    pub fn as_str(&self) -> &str {
41        &self.0
42    }
43}
44
45impl TryFrom<&str> for ResourceQuantity {
46    type Error = ConversionError;
47
48    fn try_from(value: &str) -> Result<Self, Self::Error> {
49        lazy_static! {
50            static ref RE: Regex = Regex::new(r"^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$")
51                .expect(&format!("{SHOULD_COMPILE}, {THIS_IS_A_BUG}"));
52        }
53
54        if !RE.is_match(value) {
55            return Err(ConversionError::DoesntMatchRegex {
56                value: value.to_string(),
57                regex: r"^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$".to_string(),
58            });
59        }
60
61        Ok(Self(value.to_string()))
62    }
63}
64
65impl From<u64> for ResourceQuantity {
66    fn from(value: u64) -> Self {
67        Self(value.to_string())
68    }
69}
70
71/// Resources limits used in the context of podman/k8s.
72#[derive(Debug, Default, Clone, PartialEq)]
73pub struct Resources {
74    request_memory: Option<ResourceQuantity>,
75    request_cpu: Option<ResourceQuantity>,
76    limit_memory: Option<ResourceQuantity>,
77    limit_cpu: Option<ResourceQuantity>,
78}
79
80#[derive(Serialize, Deserialize)]
81struct ResourcesField {
82    memory: Option<ResourceQuantity>,
83    cpu: Option<ResourceQuantity>,
84}
85
86impl Serialize for Resources {
87    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
88    where
89        S: serde::Serializer,
90    {
91        let mut state = serializer.serialize_struct("Resources", 2)?;
92
93        if self.request_memory.is_some() || self.request_memory.is_some() {
94            state.serialize_field(
95                "requests",
96                &ResourcesField {
97                    memory: self.request_memory.clone(),
98                    cpu: self.request_cpu.clone(),
99                },
100            )?;
101        } else {
102            state.skip_field("requests")?;
103        }
104
105        if self.limit_memory.is_some() || self.limit_memory.is_some() {
106            state.serialize_field(
107                "limits",
108                &ResourcesField {
109                    memory: self.limit_memory.clone(),
110                    cpu: self.limit_cpu.clone(),
111                },
112            )?;
113        } else {
114            state.skip_field("limits")?;
115        }
116
117        state.end()
118    }
119}
120
121struct ResourcesVisitor;
122
123impl<'de> de::Visitor<'de> for ResourcesVisitor {
124    type Value = Resources;
125
126    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
127        formatter.write_str("a resources object")
128    }
129
130    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
131    where
132        A: de::MapAccess<'de>,
133    {
134        let mut resources: Resources = Resources::default();
135
136        while let Some((key, value)) = map.next_entry::<String, ResourcesField>()? {
137            match key.as_str() {
138                "requests" => {
139                    resources.request_memory = value.memory;
140                    resources.request_cpu = value.cpu;
141                },
142                "limits" => {
143                    resources.limit_memory = value.memory;
144                    resources.limit_cpu = value.cpu;
145                },
146                _ => {
147                    return Err(de::Error::unknown_field(
148                        &key,
149                        &["requests", "limits", "cpu", "memory"],
150                    ))
151                },
152            }
153        }
154        Ok(resources)
155    }
156}
157
158impl<'de> Deserialize<'de> for Resources {
159    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
160    where
161        D: serde::Deserializer<'de>,
162    {
163        deserializer.deserialize_any(ResourcesVisitor)
164    }
165}
166
167impl Resources {
168    /// Memory limit applied to requests.
169    pub fn request_memory(&self) -> Option<&ResourceQuantity> {
170        self.request_memory.as_ref()
171    }
172
173    /// CPU limit applied to requests.
174    pub fn request_cpu(&self) -> Option<&ResourceQuantity> {
175        self.request_cpu.as_ref()
176    }
177
178    /// Overall memory limit applied.
179    pub fn limit_memory(&self) -> Option<&ResourceQuantity> {
180        self.limit_memory.as_ref()
181    }
182
183    /// Overall CPU limit applied.
184    pub fn limit_cpu(&self) -> Option<&ResourceQuantity> {
185        self.limit_cpu.as_ref()
186    }
187}
188
189/// A resources builder, used to build a [`Resources`] declaratively with fields validation.
190#[derive(Debug, Default)]
191pub struct ResourcesBuilder {
192    config: Resources,
193    errors: Vec<anyhow::Error>,
194}
195
196impl ResourcesBuilder {
197    pub fn new() -> ResourcesBuilder {
198        Self::default()
199    }
200
201    fn transition(config: Resources, errors: Vec<anyhow::Error>) -> Self {
202        Self { config, errors }
203    }
204
205    /// Set the requested memory for a pod. This is the minimum memory allocated for a pod.
206    pub fn with_request_memory<T>(self, quantity: T) -> Self
207    where
208        T: TryInto<ResourceQuantity>,
209        T::Error: Error + Send + Sync + 'static,
210    {
211        match quantity.try_into() {
212            Ok(quantity) => Self::transition(
213                Resources {
214                    request_memory: Some(quantity),
215                    ..self.config
216                },
217                self.errors,
218            ),
219            Err(error) => Self::transition(
220                self.config,
221                merge_errors(self.errors, FieldError::RequestMemory(error.into()).into()),
222            ),
223        }
224    }
225
226    /// Set the requested CPU limit for a pod. This is the minimum CPU allocated for a pod.
227    pub fn with_request_cpu<T>(self, quantity: T) -> Self
228    where
229        T: TryInto<ResourceQuantity>,
230        T::Error: Error + Send + Sync + 'static,
231    {
232        match quantity.try_into() {
233            Ok(quantity) => Self::transition(
234                Resources {
235                    request_cpu: Some(quantity),
236                    ..self.config
237                },
238                self.errors,
239            ),
240            Err(error) => Self::transition(
241                self.config,
242                merge_errors(self.errors, FieldError::RequestCpu(error.into()).into()),
243            ),
244        }
245    }
246
247    /// Set the overall memory limit for a pod. This is the maximum memory threshold for a pod.
248    pub fn with_limit_memory<T>(self, quantity: T) -> Self
249    where
250        T: TryInto<ResourceQuantity>,
251        T::Error: Error + Send + Sync + 'static,
252    {
253        match quantity.try_into() {
254            Ok(quantity) => Self::transition(
255                Resources {
256                    limit_memory: Some(quantity),
257                    ..self.config
258                },
259                self.errors,
260            ),
261            Err(error) => Self::transition(
262                self.config,
263                merge_errors(self.errors, FieldError::LimitMemory(error.into()).into()),
264            ),
265        }
266    }
267
268    /// Set the overall CPU limit for a pod. This is the maximum CPU threshold for a pod.
269    pub fn with_limit_cpu<T>(self, quantity: T) -> Self
270    where
271        T: TryInto<ResourceQuantity>,
272        T::Error: Error + Send + Sync + 'static,
273    {
274        match quantity.try_into() {
275            Ok(quantity) => Self::transition(
276                Resources {
277                    limit_cpu: Some(quantity),
278                    ..self.config
279                },
280                self.errors,
281            ),
282            Err(error) => Self::transition(
283                self.config,
284                merge_errors(self.errors, FieldError::LimitCpu(error.into()).into()),
285            ),
286        }
287    }
288
289    /// Seals the builder and returns a [`Resources`] if there are no validation errors, else returns errors.
290    pub fn build(self) -> Result<Resources, Vec<anyhow::Error>> {
291        if !self.errors.is_empty() {
292            return Err(self.errors);
293        }
294
295        Ok(self.config)
296    }
297}
298
299#[cfg(test)]
300#[allow(non_snake_case)]
301mod tests {
302    use super::*;
303    use crate::NetworkConfig;
304
305    macro_rules! impl_resources_quantity_unit_test {
306        ($val:literal) => {{
307            let resources = ResourcesBuilder::new()
308                .with_request_memory($val)
309                .build()
310                .unwrap();
311
312            assert_eq!(resources.request_memory().unwrap().as_str(), $val);
313            assert_eq!(resources.request_cpu(), None);
314            assert_eq!(resources.limit_cpu(), None);
315            assert_eq!(resources.limit_memory(), None);
316        }};
317    }
318
319    #[test]
320    fn converting_a_string_a_resource_quantity_without_unit_should_succeeds() {
321        impl_resources_quantity_unit_test!("1000");
322    }
323
324    #[test]
325    fn converting_a_str_with_m_unit_into_a_resource_quantity_should_succeeds() {
326        impl_resources_quantity_unit_test!("100m");
327    }
328
329    #[test]
330    fn converting_a_str_with_K_unit_into_a_resource_quantity_should_succeeds() {
331        impl_resources_quantity_unit_test!("50K");
332    }
333
334    #[test]
335    fn converting_a_str_with_M_unit_into_a_resource_quantity_should_succeeds() {
336        impl_resources_quantity_unit_test!("100M");
337    }
338
339    #[test]
340    fn converting_a_str_with_G_unit_into_a_resource_quantity_should_succeeds() {
341        impl_resources_quantity_unit_test!("1G");
342    }
343
344    #[test]
345    fn converting_a_str_with_T_unit_into_a_resource_quantity_should_succeeds() {
346        impl_resources_quantity_unit_test!("0.01T");
347    }
348
349    #[test]
350    fn converting_a_str_with_P_unit_into_a_resource_quantity_should_succeeds() {
351        impl_resources_quantity_unit_test!("0.00001P");
352    }
353
354    #[test]
355    fn converting_a_str_with_E_unit_into_a_resource_quantity_should_succeeds() {
356        impl_resources_quantity_unit_test!("0.000000001E");
357    }
358
359    #[test]
360    fn converting_a_str_with_Ki_unit_into_a_resource_quantity_should_succeeds() {
361        impl_resources_quantity_unit_test!("50Ki");
362    }
363
364    #[test]
365    fn converting_a_str_with_Mi_unit_into_a_resource_quantity_should_succeeds() {
366        impl_resources_quantity_unit_test!("100Mi");
367    }
368
369    #[test]
370    fn converting_a_str_with_Gi_unit_into_a_resource_quantity_should_succeeds() {
371        impl_resources_quantity_unit_test!("1Gi");
372    }
373
374    #[test]
375    fn converting_a_str_with_Ti_unit_into_a_resource_quantity_should_succeeds() {
376        impl_resources_quantity_unit_test!("0.01Ti");
377    }
378
379    #[test]
380    fn converting_a_str_with_Pi_unit_into_a_resource_quantity_should_succeeds() {
381        impl_resources_quantity_unit_test!("0.00001Pi");
382    }
383
384    #[test]
385    fn converting_a_str_with_Ei_unit_into_a_resource_quantity_should_succeeds() {
386        impl_resources_quantity_unit_test!("0.000000001Ei");
387    }
388
389    #[test]
390    fn resources_config_builder_should_succeeds_and_returns_a_resources_config() {
391        let resources = ResourcesBuilder::new()
392            .with_request_memory("200M")
393            .with_request_cpu("1G")
394            .with_limit_cpu("500M")
395            .with_limit_memory("2G")
396            .build()
397            .unwrap();
398
399        assert_eq!(resources.request_memory().unwrap().as_str(), "200M");
400        assert_eq!(resources.request_cpu().unwrap().as_str(), "1G");
401        assert_eq!(resources.limit_cpu().unwrap().as_str(), "500M");
402        assert_eq!(resources.limit_memory().unwrap().as_str(), "2G");
403    }
404
405    #[test]
406    fn resources_config_toml_import_should_succeeds_and_returns_a_resources_config() {
407        let load_from_toml =
408            NetworkConfig::load_from_toml("./testing/snapshots/0001-big-network.toml").unwrap();
409
410        let resources = load_from_toml.relaychain().default_resources().unwrap();
411        assert_eq!(resources.request_memory().unwrap().as_str(), "500M");
412        assert_eq!(resources.request_cpu().unwrap().as_str(), "100000");
413        assert_eq!(resources.limit_cpu().unwrap().as_str(), "10Gi");
414        assert_eq!(resources.limit_memory().unwrap().as_str(), "4000M");
415    }
416
417    #[test]
418    fn resources_config_builder_should_fails_and_returns_an_error_if_couldnt_parse_request_memory()
419    {
420        let resources_builder = ResourcesBuilder::new().with_request_memory("invalid");
421
422        let errors = resources_builder.build().err().unwrap();
423
424        assert_eq!(errors.len(), 1);
425        assert_eq!(
426            errors.first().unwrap().to_string(),
427            r"request_memory: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
428        );
429    }
430
431    #[test]
432    fn resources_config_builder_should_fails_and_returns_an_error_if_couldnt_parse_request_cpu() {
433        let resources_builder = ResourcesBuilder::new().with_request_cpu("invalid");
434
435        let errors = resources_builder.build().err().unwrap();
436
437        assert_eq!(errors.len(), 1);
438        assert_eq!(
439            errors.first().unwrap().to_string(),
440            r"request_cpu: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
441        );
442    }
443
444    #[test]
445    fn resources_config_builder_should_fails_and_returns_an_error_if_couldnt_parse_limit_memory() {
446        let resources_builder = ResourcesBuilder::new().with_limit_memory("invalid");
447
448        let errors = resources_builder.build().err().unwrap();
449
450        assert_eq!(errors.len(), 1);
451        assert_eq!(
452            errors.first().unwrap().to_string(),
453            r"limit_memory: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
454        );
455    }
456
457    #[test]
458    fn resources_config_builder_should_fails_and_returns_an_error_if_couldnt_parse_limit_cpu() {
459        let resources_builder = ResourcesBuilder::new().with_limit_cpu("invalid");
460
461        let errors = resources_builder.build().err().unwrap();
462
463        assert_eq!(errors.len(), 1);
464        assert_eq!(
465            errors.first().unwrap().to_string(),
466            r"limit_cpu: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
467        );
468    }
469
470    #[test]
471    fn resources_config_builder_should_fails_and_returns_multiple_error_if_couldnt_parse_multiple_fields(
472    ) {
473        let resources_builder = ResourcesBuilder::new()
474            .with_limit_cpu("invalid")
475            .with_request_memory("invalid");
476
477        let errors = resources_builder.build().err().unwrap();
478
479        assert_eq!(errors.len(), 2);
480        assert_eq!(
481            errors.first().unwrap().to_string(),
482            r"limit_cpu: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
483        );
484        assert_eq!(
485            errors.get(1).unwrap().to_string(),
486            r"request_memory: 'invalid' doesn't match regex '^\d+(.\d+)?(m|K|M|G|T|P|E|Ki|Mi|Gi|Ti|Pi|Ei)?$'"
487        );
488    }
489}
zombienet_configuration/shared/resources.rs

zombienet_configuration/shared/
resources.rs
