Skip to main content

zombienet_orchestrator/network/node/
core.rs

1//! Provider-generic node behaviour, shared by every node type.
2//!
3//! [`NodeCore`] wraps the handle to the running process (the provider's
4//! [`DynNode`]) plus the small amount of runtime state zombienet keeps about
5//! it, and implements everything that doesn't depend on _what_ the node is
6//! running: logs, lifecycle (pause/resume/restart), scripts and db snapshots.
7//!
8//! Chain-flavoured nodes ([`NetworkNode`](super::NetworkNode) for substrate,
9//! [`JamNetworkNode`](super::jam::JamNetworkNode) for JAM) embed one of these
10//! and add their own protocol-specific surface on top.
11
12use std::{
13    path::{Path, PathBuf},
14    sync::{
15        atomic::{AtomicBool, AtomicU64, Ordering},
16        Arc,
17    },
18    time::{Duration, SystemTime, UNIX_EPOCH},
19};
20
21use anyhow::anyhow;
22use configuration::types::AssetLocation;
23use fancy_regex::Regex;
24use glob_match::glob_match;
25use provider::{
26    types::{ExecutionResult, InnerSnapshotDb, RunScriptOptions},
27    DynNode,
28};
29use serde::Serialize;
30use tracing::{debug, trace};
31
32use super::{
33    serialize_provider_node, spawned::NodeKind, BoxedClosure, LogLineCount, LogLineCountOptions,
34};
35use crate::shared::types::NodeSnapshot;
36
37/// The provider handle and runtime state every spawned node has.
38///
39/// Cloning a `NodeCore` shares the same underlying process handle _and_ the
40/// same running/last-start state, so all the clones of a node stay in sync.
41#[derive(Clone, Serialize)]
42pub struct NodeCore {
43    #[serde(serialize_with = "serialize_provider_node")]
44    pub(crate) inner: DynNode,
45    pub(crate) name: String,
46    /// What flavour of node this is. Stored (and serialized to `zombie.json`)
47    /// so a node can be identified without its concrete type at hand.
48    pub(crate) kind: NodeKind,
49    #[serde(skip)]
50    is_running: Arc<AtomicBool>,
51    // Store the last timestamp when we start the node
52    #[serde(skip)]
53    last_start_ts: Arc<AtomicU64>,
54}
55
56impl NodeCore {
57    pub(crate) fn new(name: impl Into<String>, inner: DynNode, kind: NodeKind) -> Self {
58        Self {
59            inner,
60            name: name.into(),
61            kind,
62            is_running: Arc::new(AtomicBool::new(false)),
63            last_start_ts: Arc::new(AtomicU64::new(0)),
64        }
65    }
66
67    /// The provider node (the actual running process/pod/container).
68    pub fn inner(&self) -> &DynNode {
69        &self.inner
70    }
71
72    pub fn name(&self) -> &str {
73        &self.name
74    }
75
76    /// What flavour of node this is.
77    pub fn kind(&self) -> NodeKind {
78        self.kind
79    }
80
81    /// Args used for bootstrap the node.
82    /// NOTE: this may not be in sync if you restart the node with new args.
83    pub fn args(&self) -> Vec<&str> {
84        self.inner.args()
85    }
86
87    /// Check if the node is currently running (not paused).
88    ///
89    /// This returns the internal running state.
90    pub fn is_running(&self) -> bool {
91        self.is_running.load(Ordering::Acquire)
92    }
93
94    /// Get the last timestamp when the node start.
95    pub fn last_start_ts(&self) -> u64 {
96        self.last_start_ts.load(Ordering::Acquire)
97    }
98
99    pub(crate) fn set_is_running(&self, is_running: bool) {
100        self.is_running.store(is_running, Ordering::Release);
101    }
102
103    /// Set the timestamp when the node was started
104    pub(crate) fn set_last_start_ts(&self, ts: u64) {
105        self.last_start_ts.store(ts, Ordering::Release);
106    }
107
108    /// On-disk base directory of the node — root of `data/`, `relay-data/`,
109    /// `cfg/`, etc.
110    /// This will be the _base directory_ of the inner (provider) node.
111    pub fn base_dir(&self) -> &PathBuf {
112        self.inner.base_dir()
113    }
114
115    /// Pause the node, this is implemented by pausing the
116    /// actual process (e.g polkadot) with sending `SIGSTOP` signal
117    ///
118    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
119    /// with global setting `teardown_on_failure` disabled.
120    pub async fn pause(&self) -> Result<(), anyhow::Error> {
121        self.set_is_running(false);
122        self.inner.pause().await?;
123        Ok(())
124    }
125
126    /// Resume the node, this is implemented by resuming the
127    /// actual process (e.g polkadot) with sending `SIGCONT` signal
128    ///
129    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
130    /// with global setting `teardown_on_failure` disabled.
131    pub async fn resume(&self) -> Result<(), anyhow::Error> {
132        self.set_is_running(true);
133        self.inner.resume().await?;
134        Ok(())
135    }
136
137    /// Restart the node using the same `cmd`, `args` and `env` (and same isolated dir)
138    ///
139    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
140    /// with global setting `teardown_on_failure` disabled.
141    pub async fn restart(&self, after: Option<Duration>) -> Result<(), anyhow::Error> {
142        self.set_is_running(false);
143        self.inner.restart(after).await?;
144        self.set_is_running(true);
145        self.set_last_start_ts(SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs());
146        Ok(())
147    }
148
149    /// Get the logs of the node
150    /// TODO: do we need the `since` param, maybe we could be handy later for loop filtering
151    pub async fn logs(&self) -> Result<String, anyhow::Error> {
152        Ok(self.inner.logs().await?)
153    }
154
155    /// Wait until a the number of matching log lines is reach
156    pub async fn wait_log_line_count(
157        &self,
158        pattern: impl Into<String>,
159        is_glob: bool,
160        count: usize,
161    ) -> Result<(), anyhow::Error> {
162        let pattern = pattern.into();
163        let pattern_clone = pattern.clone();
164        debug!("waiting until we find pattern {pattern} {count} times");
165        let match_fn: BoxedClosure = if is_glob {
166            Box::new(move |line: &str| Ok(glob_match(&pattern, line)))
167        } else {
168            let re = Regex::new(&pattern)?;
169            Box::new(move |line: &str| re.is_match(line).map_err(|e| anyhow!(e.to_string())))
170        };
171
172        loop {
173            let mut q = 0_usize;
174            let logs = self.logs().await?;
175            for line in logs.lines() {
176                trace!("line is {line}");
177                if match_fn(line)? {
178                    trace!("pattern {pattern_clone} match in line {line}");
179                    q += 1;
180                    if q >= count {
181                        return Ok(());
182                    }
183                }
184            }
185
186            tokio::time::sleep(Duration::from_secs(2)).await;
187        }
188    }
189
190    /// Waits until the number of matching log lines satisfies a custom condition,
191    /// optionally waiting for the entire duration of the timeout.
192    ///
193    /// This method searches log lines for a given substring or glob pattern,
194    /// and evaluates the number of matching lines using a user-provided predicate function.
195    /// Optionally, it can wait for the full timeout duration to ensure the condition
196    /// holds consistently (e.g., for verifying absence of logs).
197    ///
198    /// # Arguments
199    /// * `substring` - The substring or pattern to match within log lines.
200    /// * `is_glob` - Whether to treat `substring` as a glob pattern (`true`) or a regex (`false`).
201    /// * `options` - Configuration for timeout, match count predicate, and full-duration waiting.
202    ///
203    /// # Returns
204    /// * `Ok(LogLineCount::TargetReached(n))` if the predicate was satisfied within the timeout,
205    /// * `Ok(LogLineCount::TargetFails(n))` if the predicate was not satisfied in time,
206    /// * `Err(e)` if an error occurred during log retrieval or matching.
207    ///
208    /// # Example
209    /// ```rust
210    /// # use std::{sync::Arc, time::Duration};
211    /// # use provider::NativeProvider;
212    /// # use support::{fs::local::LocalFileSystem};
213    /// # use zombienet_orchestrator::{Orchestrator, network::node::{NetworkNode, LogLineCountOptions}};
214    /// # use configuration::NetworkConfig;
215    /// # async fn example() -> Result<(), anyhow::Error> {
216    /// #   let provider = NativeProvider::new(LocalFileSystem {});
217    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
218    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
219    /// #   let network = orchestrator.spawn(config).await?;
220    /// let node = network.get_node("alice")?;
221    /// // Wait (up to 10 seconds) until pattern occurs once
222    /// let options = LogLineCountOptions {
223    ///     predicate: Arc::new(|count| count == 1),
224    ///     timeout: Duration::from_secs(10),
225    ///     wait_until_timeout_elapses: false,
226    /// };
227    /// let result = node
228    ///     .wait_log_line_count_with_timeout("error", false, options)
229    ///     .await?;
230    /// #   Ok(())
231    /// # }
232    /// ```
233    pub async fn wait_log_line_count_with_timeout(
234        &self,
235        substring: impl Into<String>,
236        is_glob: bool,
237        options: LogLineCountOptions,
238    ) -> Result<LogLineCount, anyhow::Error> {
239        let substring = substring.into();
240        debug!(
241            "waiting until match lines count within {} seconds",
242            options.timeout.as_secs_f64()
243        );
244
245        let start = tokio::time::Instant::now();
246
247        let match_fn: BoxedClosure = if is_glob {
248            Box::new(move |line: &str| Ok(glob_match(&substring, line)))
249        } else {
250            let re = Regex::new(&substring)?;
251            Box::new(move |line: &str| re.is_match(line).map_err(|e| anyhow!(e.to_string())))
252        };
253
254        if options.wait_until_timeout_elapses {
255            tokio::time::sleep(options.timeout).await;
256        }
257
258        let mut q;
259        loop {
260            q = 0_u32;
261            let logs = self.logs().await?;
262            for line in logs.lines() {
263                if match_fn(line)? {
264                    q += 1;
265
266                    // If `wait_until_timeout_elapses` is set then check the condition just once at the
267                    // end after the whole log file is processed. This is to address the cases when the
268                    // predicate becomes true and false again.
269                    // eg. expected exactly 2 matching lines are expected but 3 are present
270                    if !options.wait_until_timeout_elapses && (options.predicate)(q) {
271                        return Ok(LogLineCount::TargetReached(q));
272                    }
273                }
274            }
275
276            if start.elapsed() >= options.timeout {
277                break;
278            }
279
280            tokio::time::sleep(Duration::from_secs(2)).await;
281        }
282
283        if (options.predicate)(q) {
284            Ok(LogLineCount::TargetReached(q))
285        } else {
286            Ok(LogLineCount::TargetFailed(q))
287        }
288    }
289
290    /// Restart the node overriding the program and/or args, optionally
291    /// downloading `assets` (binaries) first.
292    ///
293    /// The node keeps the same `env` and isolated directory (e.g same database).
294    /// Callers are expected to have (re)generated `program`/`args` for their
295    /// own node flavour.
296    pub(crate) async fn restart_with(
297        &self,
298        assets: &[AssetLocation],
299        program: &str,
300        args: &[String],
301        after: Option<Duration>,
302    ) -> Result<(), anyhow::Error> {
303        self.set_is_running(false);
304        self.inner
305            .restart_with(assets, program, args, after)
306            .await?;
307        self.set_is_running(true);
308        self.set_last_start_ts(SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs());
309        Ok(())
310    }
311
312    /// Tar the node's database into `out_path` (gzipped).
313    ///
314    /// NOTE: Currently __only__ implemented in native provider. Also,
315    /// the caller is responsible for pausing the node first;
316    /// snapshotting a running node risks a torn RocksDB state.
317    pub(crate) async fn snapshot_db(
318        &self,
319        out_path: impl AsRef<Path>,
320        is_cumulus_based: bool,
321    ) -> Result<NodeSnapshot, anyhow::Error> {
322        let out_path = out_path.as_ref().to_path_buf();
323
324        let InnerSnapshotDb {
325            filename,
326            sha256,
327            size,
328        } = self.inner.snapshot_db(is_cumulus_based).await?;
329
330        // now we need to _move_ the inner file to the out_path
331        let remote_file_path = PathBuf::from(&filename);
332        self.inner
333            .receive_file(remote_file_path.as_ref(), out_path.as_ref())
334            .await?;
335
336        Ok(NodeSnapshot {
337            path: out_path,
338            sha256,
339            size,
340            node_name: self.name().into(),
341        })
342    }
343
344    /// Run a script inside the node's container/environment
345    ///
346    /// The script will be uploaded to the node, made executable, and executed with
347    /// the provided arguments and environment variables.
348    ///
349    /// Returns `Ok(stdout)` on success, or `Err((exit_status, stderr))` on failure.
350    pub async fn run_script(
351        &self,
352        options: RunScriptOptions,
353    ) -> Result<ExecutionResult, anyhow::Error> {
354        self.inner
355            .run_script(options)
356            .await
357            .map_err(|e| anyhow!("Failed to run script: {e}"))
358    }
359}
360
361impl std::fmt::Debug for NodeCore {
362    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
363        f.debug_struct("NodeCore")
364            .field("inner", &"inner_skipped")
365            .field("name", &self.name)
366            .field("kind", &self.kind)
367            .field("is_running", &self.is_running())
368            .finish()
369    }
370}