Skip to main content

zombienet_orchestrator/network/node/
jam.rs

1//! JAM nodes.
2
3use std::{net::IpAddr, time::Duration};
4
5use anyhow::anyhow;
6use async_trait::async_trait;
7use configuration::types::{Arg, AssetLocation, JamNodeMode};
8use serde::{Deserialize, Serialize};
9use subxt::ext::jsonrpsee::ws_client::{WsClient, WsClientBuilder};
10use support::net::{wait_tcp_ready, wait_ws_ready};
11use tracing::debug;
12
13use super::{
14    core::NodeCore,
15    spawned::{NodeKind, SpawnedNode},
16};
17use crate::{
18    generators::{generate_jam_node_command, GenCmdOptions},
19    network_spec::jamnode::JamNodeSpec,
20};
21
22/// A running JAM node.
23///
24/// Same provider-generic behaviour as any other node (see [`NodeCore`]), plus
25/// the JAM specifics: the libp2p peer identity used to build bootnode
26/// addresses, the node mode and the rpc endpoint.
27///
28/// Unlike [`NetworkNode`](super::NetworkNode) it exposes no subxt client and
29/// no Prometheus metrics assertions, since JAM nodes serve neither.
30#[derive(Clone, Serialize)]
31pub struct JamNetworkNode {
32    #[serde(flatten)]
33    pub(crate) core: NodeCore,
34    pub(crate) spec: JamNodeSpec,
35    /// Ip the node is reachable at.
36    pub(crate) ip: IpAddr,
37    /// `ip:rpc_port`, only bound by nodes running in `Ordinary` mode.
38    pub(crate) rpc_uri: String,
39    /// `{peer_id}@{ip}:{port}`, the form other nodes take as `--bootnode`.
40    pub(crate) peer_addr: String,
41    // Store the options used to generate the cmd, so we can recalculate
42    // (cmd, args) from a modified spec on restart.
43    pub(crate) cmd_generator_opts: GenCmdOptions,
44}
45
46/// Deserialization counterpart used when re-attaching to a running network.
47#[derive(Deserialize)]
48pub(crate) struct RawJamNetworkNode {
49    pub(crate) name: String,
50    pub(crate) spec: JamNodeSpec,
51    pub(crate) ip: IpAddr,
52    pub(crate) rpc_uri: String,
53    pub(crate) peer_addr: String,
54    pub(crate) cmd_generator_opts: GenCmdOptions,
55    #[serde(default)]
56    pub(crate) inner: serde_json::Value,
57}
58
59impl JamNetworkNode {
60    pub(crate) fn new(
61        name: impl Into<String>,
62        inner: provider::DynNode,
63        spec: JamNodeSpec,
64        ip: IpAddr,
65        cmd_generator_opts: GenCmdOptions,
66    ) -> Self {
67        let rpc_uri = format!("{ip}:{}", spec.rpc_port.0);
68        let peer_addr = format!("{}@{ip}:{}", spec.peer_id, spec.port.0);
69
70        Self {
71            core: NodeCore::new(name, inner, NodeKind::Jam),
72            spec,
73            ip,
74            rpc_uri,
75            peer_addr,
76            cmd_generator_opts,
77        }
78    }
79
80    /// The provider-generic part of this node.
81    pub fn core(&self) -> &NodeCore {
82        &self.core
83    }
84
85    pub fn name(&self) -> &str {
86        self.core.name()
87    }
88
89    pub fn spec(&self) -> &JamNodeSpec {
90        &self.spec
91    }
92
93    /// Mode this node runs in (`ordinary`, `validator` or `proxy`).
94    pub fn mode(&self) -> &JamNodeMode {
95        &self.spec.mode
96    }
97
98    /// The node libp2p local identity.
99    pub fn peer_id(&self) -> &str {
100        &self.spec.peer_id
101    }
102
103    /// `{peer_id}@{ip}:{port}`, as passed to other nodes with `--bootnode`.
104    pub fn peer_addr(&self) -> &str {
105        &self.peer_addr
106    }
107
108    /// `ip:rpc_port`.
109    ///
110    /// NOTE: only nodes running in [`JamNodeMode::Ordinary`] are started with
111    /// `--rpc-port`, so this is not bound for validators/proxies.
112    pub fn rpc_uri(&self) -> &str {
113        &self.rpc_uri
114    }
115
116    pub fn ws_uri(&self) -> String {
117        format!("ws://{}", self.rpc_uri)
118    }
119
120    /// Address used to probe for readiness: the rpc port for ordinary nodes
121    /// (the only ones that bind it), the p2p port otherwise.
122    fn probe_addr(&self) -> String {
123        match self.spec.mode {
124            JamNodeMode::Ordinary => self.rpc_uri.clone(),
125            JamNodeMode::Validator | JamNodeMode::Proxy => {
126                format!("{}:{}", self.ip, self.spec.port.0)
127            },
128        }
129    }
130
131    /// Check if the node is responsive by attempting to connect to it.
132    ///
133    /// This performs an actual connection attempt with a short timeout (2 seconds).
134    /// Returns `true` if the node is reachable and responding, `false` otherwise.
135    ///
136    /// This is more robust than `is_running()` as it verifies the node is actually alive.
137    pub async fn is_responsive(&self) -> bool {
138        tokio::time::timeout(Duration::from_secs(2), wait_tcp_ready(&self.probe_addr()))
139            .await
140            .is_ok()
141    }
142
143    /// Restart the node using the optional provided:
144    /// - Assets: binaries to download
145    /// - cmd: program to exec.
146    /// - args: Arguments to override the ones provided in config.
147    ///
148    /// The node will be restarted with the same `env` and isolated directory
149    /// (e.g same database).
150    pub async fn restart_with(
151        &self,
152        assets: Vec<AssetLocation>,
153        program: Option<String>,
154        args: Option<Vec<Arg>>,
155        after: Option<Duration>,
156    ) -> Result<(), anyhow::Error> {
157        let mut spec_cloned = self.spec.clone();
158
159        if let Some(args) = args {
160            spec_cloned.args = args;
161        }
162        if let Some(program) = program {
163            spec_cloned.command = program.as_str().try_into()?;
164        }
165
166        let (program, args) =
167            generate_jam_node_command(&spec_cloned, self.cmd_generator_opts.clone());
168
169        self.core
170            .restart_with(&assets, &program, &args, after)
171            .await
172    }
173
174    /// Try to connect to the node.
175    ///
176    /// Most of the time you only want to use [`JamNetworkNode::wait_client`] that waits for
177    /// the node to appear before it connects to it. This function directly tries
178    /// to connect to the node and returns an error if the node is not yet available
179    /// at that point in time.
180    ///
181    /// Return a [WsClient]
182    async fn try_client(&self) -> Result<WsClient, anyhow::Error> {
183        match WsClientBuilder::default().build(self.ws_uri()).await {
184            Ok(client) => Ok(client),
185            Err(error) => Err(anyhow!(format!("Error building a wsClient: {}", error))),
186        }
187    }
188
189    /// Wait until get the [WsClient] for the node
190    pub async fn wait_client(&self) -> Result<WsClient, anyhow::Error> {
191        debug!("wait_client ws_uri: {}", self.ws_uri());
192        wait_ws_ready(&self.ws_uri())
193            .await
194            .map_err(|e| anyhow!("Error awaiting http_client to be ready, err: {e}"))?;
195
196        self.try_client()
197            .await
198            .map_err(|e| anyhow!("Can't create a wsClient, err: {e}"))
199    }
200
201    /// Wait until get the [WsClient] for the node with a defined timeout
202    pub async fn wait_client_with_timeout(
203        &self,
204        timeout_secs: impl Into<u64>,
205    ) -> Result<WsClient, anyhow::Error> {
206        debug!("waiting until client is ready");
207        tokio::time::timeout(Duration::from_secs(timeout_secs.into()), self.wait_client()).await?
208    }
209}
210
211#[async_trait]
212impl SpawnedNode for JamNetworkNode {
213    fn core(&self) -> &NodeCore {
214        &self.core
215    }
216
217    /// JAM nodes expose no Prometheus endpoint, so readiness is established by
218    /// connecting to the port the node binds for its mode.
219    async fn wait_until_is_up(&self, timeout_secs: u64) -> Result<(), anyhow::Error> {
220        // Validators and proxies speak QUIC over UDP on their p2p port; a TCP probe can
221        // never succeed there. Only the ordinary node exposes a TCP (RPC) endpoint to wait
222        // on.
223        if !matches!(self.spec.mode, JamNodeMode::Ordinary) {
224            debug!(
225                "[{}] validator/proxy p2p is UDP; skipping TCP readiness wait",
226                self.name()
227            );
228            return Ok(());
229        }
230        let addr = self.probe_addr();
231        debug!("[{}] waiting until {addr} is reachable", self.name());
232
233        tokio::time::timeout(Duration::from_secs(timeout_secs), wait_tcp_ready(&addr))
234            .await
235            .map_err(|_| {
236                anyhow!(
237                    "Timeout ({timeout_secs}), waiting for {} to be up at {addr}",
238                    self.name()
239                )
240            })?
241            .map_err(|err| anyhow!("{}: {:?}", self.name(), err))
242    }
243}
244
245impl std::fmt::Debug for JamNetworkNode {
246    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247        f.debug_struct("JamNetworkNode")
248            .field("inner", &"inner_skipped")
249            .field("spec", &self.spec)
250            .field("name", &self.name())
251            .field("mode", &self.spec.mode)
252            .field("peer_addr", &self.peer_addr)
253            .field("rpc_uri", &self.rpc_uri)
254            .finish()
255    }
256}