zombienet_provider/kubernetes/
pod_spec_builder.rs

1use std::collections::BTreeMap;
2
3use configuration::shared::resources::{ResourceQuantity, Resources};
4use k8s_openapi::{
5    api::core::v1::{
6        ConfigMapVolumeSource, Container, EnvVar, PodSpec, ResourceRequirements, Volume,
7        VolumeMount,
8    },
9    apimachinery::pkg::api::resource::Quantity,
10};
11
12pub(super) struct PodSpecBuilder;
13
14impl PodSpecBuilder {
15    pub(super) fn build(
16        name: &str,
17        image: &str,
18        resources: Option<&Resources>,
19        program: &str,
20        args: &[String],
21        env: &[(String, String)],
22    ) -> PodSpec {
23        PodSpec {
24            hostname: Some(name.to_string()),
25            init_containers: Some(vec![Self::build_helper_binaries_setup_container()]),
26            containers: vec![Self::build_main_container(
27                name, image, resources, program, args, env,
28            )],
29            volumes: Some(Self::build_volumes()),
30            ..Default::default()
31        }
32    }
33
34    fn build_main_container(
35        name: &str,
36        image: &str,
37        resources: Option<&Resources>,
38        program: &str,
39        args: &[String],
40        env: &[(String, String)],
41    ) -> Container {
42        Container {
43            name: name.to_string(),
44            image: Some(image.to_string()),
45            image_pull_policy: Some("Always".to_string()),
46            command: Some(
47                [
48                    vec!["/zombie-wrapper.sh".to_string(), program.to_string()],
49                    args.to_vec(),
50                ]
51                .concat(),
52            ),
53            env: Some(
54                env.iter()
55                    .map(|(name, value)| EnvVar {
56                        name: name.clone(),
57                        value: Some(value.clone()),
58                        value_from: None,
59                    })
60                    .collect(),
61            ),
62            volume_mounts: Some(Self::build_volume_mounts(vec![VolumeMount {
63                name: "zombie-wrapper-volume".to_string(),
64                mount_path: "/zombie-wrapper.sh".to_string(),
65                sub_path: Some("zombie-wrapper.sh".to_string()),
66                ..Default::default()
67            }])),
68            resources: Self::build_resources_requirements(resources),
69            ..Default::default()
70        }
71    }
72
73    fn build_helper_binaries_setup_container() -> Container {
74        Container {
75            name: "helper-binaries-setup".to_string(),
76            image: Some("europe-west3-docker.pkg.dev/parity-zombienet/zombienet-public-images/alpine:latest".to_string()),
77            image_pull_policy: Some("IfNotPresent".to_string()),
78            volume_mounts: Some(Self::build_volume_mounts(vec![VolumeMount {
79                name: "helper-binaries-downloader-volume".to_string(),
80                mount_path: "/helper-binaries-downloader.sh".to_string(),
81                sub_path: Some("helper-binaries-downloader.sh".to_string()),
82                ..Default::default()
83            }])),
84            command: Some(vec![
85                "ash".to_string(),
86                "/helper-binaries-downloader.sh".to_string(),
87            ]),
88            ..Default::default()
89        }
90    }
91
92    fn build_volumes() -> Vec<Volume> {
93        vec![
94            Volume {
95                name: "cfg".to_string(),
96                ..Default::default()
97            },
98            Volume {
99                name: "data".to_string(),
100                ..Default::default()
101            },
102            Volume {
103                name: "relay-data".to_string(),
104                ..Default::default()
105            },
106            Volume {
107                name: "zombie-wrapper-volume".to_string(),
108                config_map: Some(ConfigMapVolumeSource {
109                    name: Some("zombie-wrapper".to_string()),
110                    default_mode: Some(0o755),
111                    ..Default::default()
112                }),
113                ..Default::default()
114            },
115            Volume {
116                name: "helper-binaries-downloader-volume".to_string(),
117                config_map: Some(ConfigMapVolumeSource {
118                    name: Some("helper-binaries-downloader".to_string()),
119                    default_mode: Some(0o755),
120                    ..Default::default()
121                }),
122                ..Default::default()
123            },
124        ]
125    }
126
127    fn build_volume_mounts(non_default_mounts: Vec<VolumeMount>) -> Vec<VolumeMount> {
128        [
129            vec![
130                VolumeMount {
131                    name: "cfg".to_string(),
132                    mount_path: "/cfg".to_string(),
133                    read_only: Some(false),
134                    ..Default::default()
135                },
136                VolumeMount {
137                    name: "data".to_string(),
138                    mount_path: "/data".to_string(),
139                    read_only: Some(false),
140                    ..Default::default()
141                },
142                VolumeMount {
143                    name: "relay-data".to_string(),
144                    mount_path: "/relay-data".to_string(),
145                    read_only: Some(false),
146                    ..Default::default()
147                },
148            ],
149            non_default_mounts,
150        ]
151        .concat()
152    }
153
154    fn build_resources_requirements(resources: Option<&Resources>) -> Option<ResourceRequirements> {
155        resources.map(|resources| ResourceRequirements {
156            limits: Self::build_resources_requirements_quantities(
157                resources.limit_cpu(),
158                resources.limit_memory(),
159            ),
160            requests: Self::build_resources_requirements_quantities(
161                resources.request_cpu(),
162                resources.request_memory(),
163            ),
164            ..Default::default()
165        })
166    }
167
168    fn build_resources_requirements_quantities(
169        cpu: Option<&ResourceQuantity>,
170        memory: Option<&ResourceQuantity>,
171    ) -> Option<BTreeMap<String, Quantity>> {
172        let mut quantities = BTreeMap::new();
173
174        if let Some(cpu) = cpu {
175            quantities.insert("cpu".to_string(), Quantity(cpu.as_str().to_string()));
176        }
177
178        if let Some(memory) = memory {
179            quantities.insert("memory".to_string(), Quantity(memory.as_str().to_string()));
180        }
181
182        if !quantities.is_empty() {
183            Some(quantities)
184        } else {
185            None
186        }
187    }
188}
zombienet_provider/kubernetes/pod_spec_builder.rs

zombienet_provider/kubernetes/
pod_spec_builder.rs
