sui_swarm/memory/
node.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::Result;
5use anyhow::anyhow;
6use std::sync::Mutex;
7use std::sync::MutexGuard;
8use sui_config::NodeConfig;
9use sui_node::SuiNodeHandle;
10use sui_types::base_types::AuthorityName;
11use sui_types::base_types::ConciseableName;
12use sui_types::crypto::KeypairTraits;
13use tap::TapFallible;
14use tracing::{error, info};
15
16use super::container::Container;
17
18/// A handle to an in-memory Sui Node.
19///
20/// Each Node is attempted to run in isolation from each other by running them in their own tokio
21/// runtime in a separate thread. By doing this we can ensure that all asynchronous tasks
22/// associated with a Node are able to be stopped when desired (either when a Node is dropped or
23/// explicitly stopped by calling [`Node::stop`]) by simply dropping that Node's runtime.
24#[derive(Debug)]
25pub struct Node {
26    container: Mutex<Option<Container>>,
27    config: Mutex<NodeConfig>,
28    runtime_type: RuntimeType,
29}
30
31impl Node {
32    /// Create a new Node from the provided `NodeConfig`.
33    ///
34    /// The Node is returned without being started. See [`Node::spawn`] or [`Node::start`] for how to
35    /// start the node.
36    ///
37    /// [`NodeConfig`]: sui_config::NodeConfig
38    pub fn new(config: NodeConfig) -> Self {
39        Self {
40            container: Default::default(),
41            config: config.into(),
42            runtime_type: RuntimeType::SingleThreaded,
43        }
44    }
45
46    /// Return the `name` of this Node
47    pub fn name(&self) -> AuthorityName {
48        self.config().protocol_public_key()
49    }
50
51    pub fn config(&self) -> MutexGuard<'_, NodeConfig> {
52        self.config.lock().unwrap()
53    }
54
55    pub fn json_rpc_address(&self) -> std::net::SocketAddr {
56        self.config().json_rpc_address
57    }
58
59    /// Start this Node
60    pub async fn spawn(&self) -> Result<()> {
61        info!(name =% self.name().concise(), "starting in-memory node");
62        let config = self.config().clone();
63        *self.container.lock().unwrap() = Some(Container::spawn(config, self.runtime_type).await);
64        Ok(())
65    }
66
67    /// Start this Node, waiting until its completely started up.
68    pub async fn start(&self) -> Result<()> {
69        self.spawn().await
70    }
71
72    /// Stop this Node
73    pub fn stop(&self) {
74        info!(name =% self.name().concise(), "stopping in-memory node");
75        *self.container.lock().unwrap() = None;
76        info!(name =% self.name().concise(), "node stopped");
77    }
78
79    /// Stop this Node the way a SIGTERM stops a production node: the tokio runtime is dropped,
80    /// so for up to `poll_budget` polls (or `max_duration`) async tasks keep running while
81    /// `spawn_blocking` returns cancelled handles. Returns after the node is stopped.
82    #[cfg(msim)]
83    pub async fn graceful_stop(&self, poll_budget: usize, max_duration: std::time::Duration) {
84        let Some(node_id) = self
85            .container
86            .lock()
87            .unwrap()
88            .as_ref()
89            .and_then(|c| c.sim_node_id())
90        else {
91            return;
92        };
93        info!(name =% self.name().concise(), poll_budget, "gracefully stopping in-memory node");
94        sui_simulator::runtime::Handle::current().begin_graceful_shutdown(
95            node_id,
96            None,
97            poll_budget,
98            max_duration,
99        );
100        // The simulator kills the node by the deadline.
101        tokio::time::sleep(max_duration).await;
102        self.stop();
103    }
104
105    /// If this Node is currently running
106    pub fn is_running(&self) -> bool {
107        self.container
108            .lock()
109            .unwrap()
110            .as_ref()
111            .is_some_and(|c| c.is_alive())
112    }
113
114    pub fn get_node_handle(&self) -> Option<SuiNodeHandle> {
115        self.container
116            .lock()
117            .unwrap()
118            .as_ref()
119            .and_then(|c| c.get_node_handle())
120    }
121
122    /// Perform a health check on this Node by:
123    /// * Checking that the node is running
124    /// * Calling the Node's gRPC Health service if it's a validator.
125    pub async fn health_check(&self, is_validator: bool) -> Result<(), HealthCheckError> {
126        {
127            let lock = self.container.lock().unwrap();
128            let container = lock.as_ref().ok_or(HealthCheckError::NotRunning)?;
129            if !container.is_alive() {
130                return Err(HealthCheckError::NotRunning);
131            }
132        }
133
134        if is_validator {
135            let network_address = self
136                .config()
137                .network_address()
138                .clone()
139                .rewrite_http_to_https();
140            let tls_config = sui_tls::create_rustls_client_config(
141                self.config().network_key_pair().public().to_owned(),
142                sui_tls::SUI_VALIDATOR_SERVER_NAME.to_string(),
143                None,
144            );
145            let channel = mysten_network::client::connect(&network_address, tls_config)
146                .await
147                .map_err(|err| anyhow!(err.to_string()))
148                .map_err(HealthCheckError::Failure)
149                .tap_err(|e| error!("error connecting to {}: {e}", self.name().concise()))?;
150
151            let mut client = tonic_health::pb::health_client::HealthClient::new(channel);
152            client
153                .check(tonic_health::pb::HealthCheckRequest::default())
154                .await
155                .map_err(|e| HealthCheckError::Failure(e.into()))
156                .tap_err(|e| {
157                    error!(
158                        "error performing health check on {}: {e}",
159                        self.name().concise()
160                    )
161                })?;
162        }
163
164        Ok(())
165    }
166}
167
168#[derive(Debug)]
169pub enum HealthCheckError {
170    NotRunning,
171    Failure(anyhow::Error),
172    Unknown(anyhow::Error),
173}
174
175impl std::fmt::Display for HealthCheckError {
176    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
177        write!(f, "{:?}", self)
178    }
179}
180
181impl std::error::Error for HealthCheckError {}
182
183/// The type of tokio runtime that should be used for a particular Node
184#[derive(Clone, Copy, Debug)]
185pub enum RuntimeType {
186    SingleThreaded,
187    MultiThreaded,
188}
189
190#[cfg(test)]
191mod test {
192    use crate::memory::Swarm;
193
194    #[tokio::test]
195    async fn start_and_stop() {
196        telemetry_subscribers::init_for_testing();
197        let swarm = Swarm::builder().build();
198
199        let validator = swarm.validator_nodes().next().unwrap();
200
201        validator.start().await.unwrap();
202        validator.health_check(true).await.unwrap();
203        validator.stop();
204        validator.health_check(true).await.unwrap_err();
205
206        validator.start().await.unwrap();
207        validator.health_check(true).await.unwrap();
208    }
209}
sui_swarm/memory/node.rs

sui_swarm/memory/
node.rs
