zombienet_provider/shared/
types.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    process::ExitStatus,
5};
6
7use configuration::shared::resources::Resources;
8use serde::{Deserialize, Serialize};
9
10pub type Port = u16;
11
12pub type ExecutionResult = Result<String, (ExitStatus, String)>;
13
14#[derive(Debug, Clone, PartialEq)]
15pub struct ProviderCapabilities {
16    // default ports internal
17    /// Ensure that we have an image for each node (k8s/podman/docker)
18    pub requires_image: bool,
19    /// Allow to customize the resources through manifest (k8s).
20    pub has_resources: bool,
21    /// Used in native to prefix filepath with fullpath
22    pub prefix_with_full_path: bool,
23    /// Use default ports in node cmd/args.
24    /// NOTE: generally used in k8s/dockers since the images expose those ports.
25    pub use_default_ports_in_cmd: bool,
26}
27
28/// Role of a node in the network, exposed by providers that support labels
29/// (k8s) as the `zombie-role` label with the values zombienet v1 used
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum NodeRole {
32    /// Relay chain node, validator or not
33    Node,
34    /// Dedicated bootnode
35    Bootnode,
36    /// Collator of a parachain that is not cumulus based (e.g adder/undying)
37    Collator,
38    /// Collator of a cumulus based parachain
39    CumulusCollator,
40    /// Short-lived helper node (e.g used to generate chain-specs)
41    Temp,
42    /// A custom process declared in the network config, running next to the nodes
43    CustomProcess,
44}
45
46impl NodeRole {
47    pub fn as_label(&self) -> &'static str {
48        match self {
49            NodeRole::Node => "node",
50            NodeRole::Bootnode => "bootnode",
51            NodeRole::Collator => "collator",
52            NodeRole::CumulusCollator => "cumulus-collator",
53            NodeRole::Temp => "temp",
54            NodeRole::CustomProcess => "custom-process",
55        }
56    }
57}
58
59/// How a provider tells that a node spawned without the wrapper is ready.
60/// On kubernetes it becomes the pod's readiness probe.
61#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct ReadinessProbe {
63    pub target: ReadinessTarget,
64    /// Seconds between two attempts.
65    pub period_secs: u64,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum ReadinessTarget {
70    /// The port accepts a TCP connection.
71    Tcp(Port),
72    /// `GET <path>` on the port answers 2xx.
73    Http { port: Port, path: String },
74}
75
76/// What a provider knows about a node's process.
77#[derive(Debug, Clone, PartialEq, Eq)]
78pub enum ProcessStatus {
79    /// The process runs. `ready` is what the provider's own check says, where
80    /// it has one (kubernetes: the pod's Ready condition); `None` otherwise.
81    Running { ready: Option<bool> },
82    /// The process exited, with its code where the provider knows it (not for
83    /// a process attached by pid that is gone).
84    Exited { code: Option<i32> },
85}
86
87#[derive(Debug, Clone)]
88pub struct SpawnNodeOptions {
89    /// Name of the node
90    pub name: String,
91    /// Image of the node (IFF is supported by the provider)
92    pub image: Option<String>,
93    /// Resources to apply to the node (IFF is supported by the provider)
94    pub resources: Option<Resources>,
95    /// Main command to execute
96    pub program: String,
97    /// Arguments to pass to the main command
98    pub args: Vec<String>,
99    /// Environment to set when running the `program`
100    pub env: Vec<(String, String)>,
101    // TODO: rename startup_files
102    /// Files to inject at startup
103    pub injected_files: Vec<TransferedFile>,
104    /// Paths to create before start the node (e.g keystore)
105    /// should be created with `create_dir_all` in order
106    /// to create the full path even when we have missing parts
107    pub created_paths: Vec<PathBuf>,
108    /// Database snapshot to be injected (a local `.tgz` resolved by the
109    /// orchestrator before spawn — providers extract it directly, no
110    /// download). See `orchestrator::generators::resolve_db_snapshots`.
111    pub db_snapshot: Option<PathBuf>,
112    pub port_mapping: Option<HashMap<Port, Port>>,
113    /// Named ports the process listens on. On k8s they are the pod's container
114    /// ports and its `Service` ports (none gives a headless Service); the other
115    /// providers ignore them.
116    pub ports: Vec<(String, Port)>,
117    /// How the provider may check readiness itself; on k8s the pod's probe.
118    pub readiness: Option<ReadinessProbe>,
119    /// Optionally specify a log path for the node
120    pub node_log_path: Option<PathBuf>,
121    /// Role of the node in the network
122    pub role: Option<NodeRole>,
123    /// Run `program` through zombienet's wrapper script, which gives the
124    /// pause/resume/restart controls nodes have. The wrapper needs `bash` in
125    /// the image and replaces the image's entrypoint, so a custom process that
126    /// is just an off-the-shelf image runs without it.
127    pub wrapper: bool,
128}
129
130impl SpawnNodeOptions {
131    pub fn new<S>(name: S, program: S) -> Self
132    where
133        S: AsRef<str>,
134    {
135        Self {
136            name: name.as_ref().to_string(),
137            image: None,
138            resources: None,
139            program: program.as_ref().to_string(),
140            args: vec![],
141            env: vec![],
142            injected_files: vec![],
143            created_paths: vec![],
144            db_snapshot: None,
145            port_mapping: None,
146            ports: vec![],
147            readiness: None,
148            node_log_path: None,
149            role: None,
150            wrapper: true,
151        }
152    }
153
154    pub fn image<S>(mut self, image: S) -> Self
155    where
156        S: AsRef<str>,
157    {
158        self.image = Some(image.as_ref().to_string());
159        self
160    }
161
162    pub fn resources(mut self, resources: Resources) -> Self {
163        self.resources = Some(resources);
164        self
165    }
166
167    pub fn db_snapshot(mut self, db_snap: Option<PathBuf>) -> Self {
168        self.db_snapshot = db_snap;
169        self
170    }
171
172    pub fn args<S, I>(mut self, args: I) -> Self
173    where
174        S: AsRef<str>,
175        I: IntoIterator<Item = S>,
176    {
177        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
178        self
179    }
180
181    pub fn env<S, I>(mut self, env: I) -> Self
182    where
183        S: AsRef<str>,
184        I: IntoIterator<Item = (S, S)>,
185    {
186        self.env = env
187            .into_iter()
188            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
189            .collect();
190        self
191    }
192
193    pub fn injected_files<I>(mut self, injected_files: I) -> Self
194    where
195        I: IntoIterator<Item = TransferedFile>,
196    {
197        self.injected_files = injected_files.into_iter().collect();
198        self
199    }
200
201    pub fn created_paths<P, I>(mut self, created_paths: I) -> Self
202    where
203        P: AsRef<Path>,
204        I: IntoIterator<Item = P>,
205    {
206        self.created_paths = created_paths
207            .into_iter()
208            .map(|path| path.as_ref().into())
209            .collect();
210        self
211    }
212
213    pub fn port_mapping(mut self, ports: HashMap<Port, Port>) -> Self {
214        self.port_mapping = Some(ports);
215        self
216    }
217
218    pub fn node_log_path(mut self, path: Option<PathBuf>) -> Self {
219        self.node_log_path = path;
220        self
221    }
222
223    pub fn role(mut self, role: NodeRole) -> Self {
224        self.role = Some(role);
225        self
226    }
227
228    pub fn ports<S, I>(mut self, ports: I) -> Self
229    where
230        S: AsRef<str>,
231        I: IntoIterator<Item = (S, Port)>,
232    {
233        self.ports = ports
234            .into_iter()
235            .map(|(name, port)| (name.as_ref().to_string(), port))
236            .collect();
237        self
238    }
239
240    pub fn readiness(mut self, probe: Option<ReadinessProbe>) -> Self {
241        self.readiness = probe;
242        self
243    }
244
245    /// Run `program` directly instead of through the zombie wrapper script.
246    ///
247    /// The wrapper needs `bash` in the image and is what pause, resume and
248    /// restart talk to, so those return an error for a node spawned this way
249    /// on docker and kubernetes (native signals the process and needs no
250    /// wrapper). The wrapper also holds the program back until the injected
251    /// files are copied in; without it the program starts at once, so it must
252    /// not need any `injected_files` at startup.
253    pub fn without_wrapper(mut self) -> Self {
254        self.wrapper = false;
255        self
256    }
257}
258
259#[derive(Debug)]
260pub struct GenerateFileCommand {
261    pub program: String,
262    pub args: Vec<String>,
263    pub env: Vec<(String, String)>,
264    pub local_output_path: PathBuf,
265}
266
267impl GenerateFileCommand {
268    pub fn new<S, P>(program: S, local_output_path: P) -> Self
269    where
270        S: AsRef<str>,
271        P: AsRef<Path>,
272    {
273        Self {
274            program: program.as_ref().to_string(),
275            args: vec![],
276            env: vec![],
277            local_output_path: local_output_path.as_ref().into(),
278        }
279    }
280
281    pub fn args<S, I>(mut self, args: I) -> Self
282    where
283        S: AsRef<str>,
284        I: IntoIterator<Item = S>,
285    {
286        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
287        self
288    }
289
290    pub fn env<S, I>(mut self, env: I) -> Self
291    where
292        S: AsRef<str>,
293        I: IntoIterator<Item = (S, S)>,
294    {
295        self.env = env
296            .into_iter()
297            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
298            .collect();
299        self
300    }
301}
302
303#[derive(Debug)]
304pub struct GenerateFilesOptions {
305    pub commands: Vec<GenerateFileCommand>,
306    pub image: Option<String>,
307    pub injected_files: Vec<TransferedFile>,
308    // Allow to control the name of the node used to create the files.
309    pub temp_name: Option<String>,
310    pub expected_path: Option<PathBuf>,
311}
312
313impl GenerateFilesOptions {
314    pub fn new<I>(commands: I, image: Option<String>, expected_path: Option<PathBuf>) -> Self
315    where
316        I: IntoIterator<Item = GenerateFileCommand>,
317    {
318        Self {
319            commands: commands.into_iter().collect(),
320            injected_files: vec![],
321            image,
322            temp_name: None,
323            expected_path,
324        }
325    }
326
327    pub fn with_files<I>(
328        commands: I,
329        image: Option<String>,
330        injected_files: &[TransferedFile],
331        expected_path: Option<PathBuf>,
332    ) -> Self
333    where
334        I: IntoIterator<Item = GenerateFileCommand>,
335    {
336        Self {
337            commands: commands.into_iter().collect(),
338            injected_files: injected_files.into(),
339            image,
340            temp_name: None,
341            expected_path,
342        }
343    }
344
345    pub fn image<S>(mut self, image: S) -> Self
346    where
347        S: AsRef<str>,
348    {
349        self.image = Some(image.as_ref().to_string());
350        self
351    }
352
353    pub fn injected_files<I>(mut self, injected_files: I) -> Self
354    where
355        I: IntoIterator<Item = TransferedFile>,
356    {
357        self.injected_files = injected_files.into_iter().collect();
358        self
359    }
360
361    pub fn temp_name(mut self, name: impl Into<String>) -> Self {
362        self.temp_name = Some(name.into());
363        self
364    }
365}
366
367#[derive(Debug)]
368pub struct RunCommandOptions {
369    pub program: String,
370    pub args: Vec<String>,
371    pub env: Vec<(String, String)>,
372}
373
374impl RunCommandOptions {
375    pub fn new<S>(program: S) -> Self
376    where
377        S: AsRef<str>,
378    {
379        Self {
380            program: program.as_ref().to_string(),
381            args: vec![],
382            env: vec![],
383        }
384    }
385
386    pub fn args<S, I>(mut self, args: I) -> Self
387    where
388        S: AsRef<str>,
389        I: IntoIterator<Item = S>,
390    {
391        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
392        self
393    }
394
395    pub fn env<S, I>(mut self, env: I) -> Self
396    where
397        S: AsRef<str>,
398        I: IntoIterator<Item = (S, S)>,
399    {
400        self.env = env
401            .into_iter()
402            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
403            .collect();
404        self
405    }
406}
407
408pub struct RunScriptOptions {
409    pub local_script_path: PathBuf,
410    pub args: Vec<String>,
411    pub env: Vec<(String, String)>,
412}
413
414impl RunScriptOptions {
415    pub fn new<P>(local_script_path: P) -> Self
416    where
417        P: AsRef<Path>,
418    {
419        Self {
420            local_script_path: local_script_path.as_ref().into(),
421            args: vec![],
422            env: vec![],
423        }
424    }
425
426    pub fn args<S, I>(mut self, args: I) -> Self
427    where
428        S: AsRef<str>,
429        I: IntoIterator<Item = S>,
430    {
431        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
432        self
433    }
434
435    pub fn env<S, I>(mut self, env: I) -> Self
436    where
437        S: AsRef<str>,
438        I: IntoIterator<Item = (S, S)>,
439    {
440        self.env = env
441            .into_iter()
442            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
443            .collect();
444        self
445    }
446}
447
448// TODO(team): I think we can rename it to FileMap?
449#[derive(Debug, Clone, Serialize, Deserialize)]
450pub struct TransferedFile {
451    pub local_path: PathBuf,
452    pub remote_path: PathBuf,
453    // TODO: Can be narrowed to have strict typing on this?
454    pub mode: String,
455}
456
457impl TransferedFile {
458    pub fn new<P>(local_path: P, remote_path: P) -> Self
459    where
460        P: AsRef<Path>,
461    {
462        Self {
463            local_path: local_path.as_ref().into(),
464            remote_path: remote_path.as_ref().into(),
465            mode: "0644".to_string(), // default to rw-r--r--
466        }
467    }
468
469    pub fn mode<S>(mut self, mode: S) -> Self
470    where
471        S: AsRef<str>,
472    {
473        self.mode = mode.as_ref().to_string();
474        self
475    }
476}
477
478impl std::fmt::Display for TransferedFile {
479    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
480        write!(
481            f,
482            "File to transfer (local: {}, remote: {})",
483            self.local_path.display(),
484            self.remote_path.display()
485        )
486    }
487}
488
489/// Result of [`ProviderNode::snapshot_db`].
490/// return the filename, hash and bytes of the generated `{node.name}.tgz`
491pub struct InnerSnapshotDb {
492    /// local filename (scoped in node base_dir)
493    pub filename: String,
494    /// Hex-encoded SHA-256 of the archive contents.
495    pub sha256: String,
496    /// Size of the archive in bytes.
497    pub size: u64,
498}
zombienet_provider/shared/types.rs

zombienet_provider/shared/
types.rs
