sui_swarm/memory/
swarm.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::Node;
5use anyhow::Result;
6use futures::future::try_join_all;
7use rand::rngs::OsRng;
8use std::collections::HashMap;
9use std::net::SocketAddr;
10use std::num::NonZeroUsize;
11use std::time::Duration;
12use std::{
13    ops,
14    path::{Path, PathBuf},
15};
16use sui_types::traffic_control::{PolicyConfig, RemoteFirewallConfig};
17
18#[cfg(msim)]
19use sui_config::node::ExecutionTimeObserverConfig;
20use sui_config::node::{AuthorityOverloadConfig, DBCheckpointConfig, RunWithRange};
21use sui_config::{ExecutionCacheConfig, NodeConfig};
22use sui_macros::nondeterministic;
23use sui_node::SuiNodeHandle;
24use sui_protocol_config::{Chain, ProtocolVersion};
25use sui_swarm_config::genesis_config::{AccountConfig, GenesisConfig, ValidatorGenesisConfig};
26use sui_swarm_config::network_config::NetworkConfig;
27use sui_swarm_config::network_config_builder::{
28    CommitteeConfig, ConfigBuilder, GlobalStateHashV2EnabledConfig, ProtocolVersionsConfig,
29    SupportedProtocolVersionsCallback,
30};
31use sui_swarm_config::node_config_builder::FullnodeConfigBuilder;
32use sui_types::base_types::AuthorityName;
33use sui_types::object::Object;
34use sui_types::supported_protocol_versions::SupportedProtocolVersions;
35use tempfile::TempDir;
36use tracing::info;
37
38pub struct SwarmBuilder<R = OsRng> {
39    rng: R,
40    // template: NodeConfig,
41    dir: Option<PathBuf>,
42    committee: CommitteeConfig,
43    genesis_config: Option<GenesisConfig>,
44    network_config: Option<NetworkConfig>,
45    chain_override: Option<Chain>,
46    additional_objects: Vec<Object>,
47    fullnode_count: usize,
48    fullnode_rpc_port: Option<u16>,
49    fullnode_rpc_addr: Option<SocketAddr>,
50    fullnode_rpc_config: Option<sui_config::RpcConfig>,
51    supported_protocol_versions_config: ProtocolVersionsConfig,
52    // Default to supported_protocol_versions_config, but can be overridden.
53    fullnode_supported_protocol_versions_config: Option<ProtocolVersionsConfig>,
54    db_checkpoint_config: DBCheckpointConfig,
55    jwk_fetch_interval: Option<Duration>,
56    num_unpruned_validators: Option<usize>,
57    authority_overload_config: Option<AuthorityOverloadConfig>,
58    execution_cache_config: Option<ExecutionCacheConfig>,
59    data_ingestion_dir: Option<PathBuf>,
60    fullnode_run_with_range: Option<RunWithRange>,
61    fullnode_policy_config: Option<PolicyConfig>,
62    fullnode_fw_config: Option<RemoteFirewallConfig>,
63    max_submit_position: Option<usize>,
64    submit_delay_step_override_millis: Option<u64>,
65    global_state_hash_v2_enabled_config: GlobalStateHashV2EnabledConfig,
66    disable_fullnode_pruning: bool,
67    #[cfg(msim)]
68    execution_time_observer_config: Option<ExecutionTimeObserverConfig>,
69}
70
71impl SwarmBuilder {
72    #[allow(clippy::new_without_default)]
73    pub fn new() -> Self {
74        Self {
75            rng: OsRng,
76            dir: None,
77            committee: CommitteeConfig::Size(NonZeroUsize::new(1).unwrap()),
78            genesis_config: None,
79            network_config: None,
80            chain_override: None,
81            additional_objects: vec![],
82            fullnode_count: 0,
83            fullnode_rpc_port: None,
84            fullnode_rpc_addr: None,
85            fullnode_rpc_config: None,
86            supported_protocol_versions_config: ProtocolVersionsConfig::Default,
87            fullnode_supported_protocol_versions_config: None,
88            db_checkpoint_config: DBCheckpointConfig::default(),
89            jwk_fetch_interval: None,
90            num_unpruned_validators: None,
91            authority_overload_config: None,
92            execution_cache_config: None,
93            data_ingestion_dir: None,
94            fullnode_run_with_range: None,
95            fullnode_policy_config: None,
96            fullnode_fw_config: None,
97            max_submit_position: None,
98            submit_delay_step_override_millis: None,
99            global_state_hash_v2_enabled_config: GlobalStateHashV2EnabledConfig::Global(true),
100            disable_fullnode_pruning: false,
101            #[cfg(msim)]
102            execution_time_observer_config: None,
103        }
104    }
105}
106
107impl<R> SwarmBuilder<R> {
108    pub fn rng<N: rand::RngCore + rand::CryptoRng>(self, rng: N) -> SwarmBuilder<N> {
109        SwarmBuilder {
110            rng,
111            dir: self.dir,
112            committee: self.committee,
113            genesis_config: self.genesis_config,
114            network_config: self.network_config,
115            chain_override: self.chain_override,
116            additional_objects: self.additional_objects,
117            fullnode_count: self.fullnode_count,
118            fullnode_rpc_port: self.fullnode_rpc_port,
119            fullnode_rpc_addr: self.fullnode_rpc_addr,
120            fullnode_rpc_config: self.fullnode_rpc_config.clone(),
121            supported_protocol_versions_config: self.supported_protocol_versions_config,
122            fullnode_supported_protocol_versions_config: self
123                .fullnode_supported_protocol_versions_config,
124            db_checkpoint_config: self.db_checkpoint_config,
125            jwk_fetch_interval: self.jwk_fetch_interval,
126            num_unpruned_validators: self.num_unpruned_validators,
127            authority_overload_config: self.authority_overload_config,
128            execution_cache_config: self.execution_cache_config,
129            data_ingestion_dir: self.data_ingestion_dir,
130            fullnode_run_with_range: self.fullnode_run_with_range,
131            fullnode_policy_config: self.fullnode_policy_config,
132            fullnode_fw_config: self.fullnode_fw_config,
133            max_submit_position: self.max_submit_position,
134            submit_delay_step_override_millis: self.submit_delay_step_override_millis,
135            global_state_hash_v2_enabled_config: self.global_state_hash_v2_enabled_config,
136            disable_fullnode_pruning: self.disable_fullnode_pruning,
137            #[cfg(msim)]
138            execution_time_observer_config: self.execution_time_observer_config,
139        }
140    }
141
142    /// Set the directory that should be used by the Swarm for any on-disk data.
143    ///
144    /// If a directory is provided, it will not be cleaned up when the Swarm is dropped.
145    ///
146    /// Defaults to using a temporary directory that will be cleaned up when the Swarm is dropped.
147    pub fn dir<P: Into<PathBuf>>(mut self, dir: P) -> Self {
148        self.dir = Some(dir.into());
149        self
150    }
151
152    /// Set the committee size (the number of validators in the validator set).
153    ///
154    /// Defaults to 1.
155    pub fn committee_size(mut self, committee_size: NonZeroUsize) -> Self {
156        self.committee = CommitteeConfig::Size(committee_size);
157        self
158    }
159
160    pub fn with_validators(mut self, validators: Vec<ValidatorGenesisConfig>) -> Self {
161        self.committee = CommitteeConfig::Validators(validators);
162        self
163    }
164
165    pub fn with_genesis_config(mut self, genesis_config: GenesisConfig) -> Self {
166        assert!(self.network_config.is_none() && self.genesis_config.is_none());
167        self.genesis_config = Some(genesis_config);
168        self
169    }
170
171    pub fn with_chain_override(mut self, chain: Chain) -> Self {
172        assert!(self.chain_override.is_none());
173        self.chain_override = Some(chain);
174        self
175    }
176
177    pub fn with_num_unpruned_validators(mut self, n: usize) -> Self {
178        assert!(self.network_config.is_none());
179        self.num_unpruned_validators = Some(n);
180        self
181    }
182
183    pub fn with_jwk_fetch_interval(mut self, i: Duration) -> Self {
184        self.jwk_fetch_interval = Some(i);
185        self
186    }
187
188    pub fn with_network_config(mut self, network_config: NetworkConfig) -> Self {
189        assert!(self.network_config.is_none() && self.genesis_config.is_none());
190        self.network_config = Some(network_config);
191        self
192    }
193
194    pub fn with_accounts(mut self, accounts: Vec<AccountConfig>) -> Self {
195        self.get_or_init_genesis_config().accounts = accounts;
196        self
197    }
198
199    pub fn with_objects<I: IntoIterator<Item = Object>>(mut self, objects: I) -> Self {
200        self.additional_objects.extend(objects);
201        self
202    }
203
204    pub fn with_fullnode_count(mut self, fullnode_count: usize) -> Self {
205        self.fullnode_count = fullnode_count;
206        self
207    }
208
209    pub fn with_fullnode_rpc_port(mut self, fullnode_rpc_port: u16) -> Self {
210        assert!(self.fullnode_rpc_addr.is_none());
211        self.fullnode_rpc_port = Some(fullnode_rpc_port);
212        self
213    }
214
215    pub fn with_fullnode_rpc_addr(mut self, fullnode_rpc_addr: SocketAddr) -> Self {
216        assert!(self.fullnode_rpc_port.is_none());
217        self.fullnode_rpc_addr = Some(fullnode_rpc_addr);
218        self
219    }
220
221    pub fn with_fullnode_rpc_config(mut self, fullnode_rpc_config: sui_config::RpcConfig) -> Self {
222        self.fullnode_rpc_config = Some(fullnode_rpc_config);
223        self
224    }
225
226    pub fn with_epoch_duration_ms(mut self, epoch_duration_ms: u64) -> Self {
227        self.get_or_init_genesis_config()
228            .parameters
229            .epoch_duration_ms = epoch_duration_ms;
230        self
231    }
232
233    pub fn with_protocol_version(mut self, v: ProtocolVersion) -> Self {
234        self.get_or_init_genesis_config()
235            .parameters
236            .protocol_version = v;
237        self
238    }
239
240    pub fn with_supported_protocol_versions(mut self, c: SupportedProtocolVersions) -> Self {
241        self.supported_protocol_versions_config = ProtocolVersionsConfig::Global(c);
242        self
243    }
244
245    pub fn with_supported_protocol_version_callback(
246        mut self,
247        func: SupportedProtocolVersionsCallback,
248    ) -> Self {
249        self.supported_protocol_versions_config = ProtocolVersionsConfig::PerValidator(func);
250        self
251    }
252
253    pub fn with_supported_protocol_versions_config(mut self, c: ProtocolVersionsConfig) -> Self {
254        self.supported_protocol_versions_config = c;
255        self
256    }
257
258    pub fn with_global_state_hash_v2_enabled_config(
259        mut self,
260        c: GlobalStateHashV2EnabledConfig,
261    ) -> Self {
262        self.global_state_hash_v2_enabled_config = c;
263        self
264    }
265
266    #[cfg(msim)]
267    pub fn with_execution_time_observer_config(mut self, c: ExecutionTimeObserverConfig) -> Self {
268        self.execution_time_observer_config = Some(c);
269        self
270    }
271
272    pub fn with_fullnode_supported_protocol_versions_config(
273        mut self,
274        c: ProtocolVersionsConfig,
275    ) -> Self {
276        self.fullnode_supported_protocol_versions_config = Some(c);
277        self
278    }
279
280    pub fn with_db_checkpoint_config(mut self, db_checkpoint_config: DBCheckpointConfig) -> Self {
281        self.db_checkpoint_config = db_checkpoint_config;
282        self
283    }
284
285    pub fn with_authority_overload_config(
286        mut self,
287        authority_overload_config: AuthorityOverloadConfig,
288    ) -> Self {
289        assert!(self.network_config.is_none());
290        self.authority_overload_config = Some(authority_overload_config);
291        self
292    }
293
294    pub fn with_execution_cache_config(
295        mut self,
296        execution_cache_config: ExecutionCacheConfig,
297    ) -> Self {
298        self.execution_cache_config = Some(execution_cache_config);
299        self
300    }
301
302    pub fn with_data_ingestion_dir(mut self, path: PathBuf) -> Self {
303        self.data_ingestion_dir = Some(path);
304        self
305    }
306
307    pub fn with_fullnode_run_with_range(mut self, run_with_range: Option<RunWithRange>) -> Self {
308        if let Some(run_with_range) = run_with_range {
309            self.fullnode_run_with_range = Some(run_with_range);
310        }
311        self
312    }
313
314    pub fn with_fullnode_policy_config(mut self, config: Option<PolicyConfig>) -> Self {
315        self.fullnode_policy_config = config;
316        self
317    }
318
319    pub fn with_fullnode_fw_config(mut self, config: Option<RemoteFirewallConfig>) -> Self {
320        self.fullnode_fw_config = config;
321        self
322    }
323
324    fn get_or_init_genesis_config(&mut self) -> &mut GenesisConfig {
325        if self.genesis_config.is_none() {
326            assert!(self.network_config.is_none());
327            self.genesis_config = Some(GenesisConfig::for_local_testing());
328        }
329        self.genesis_config.as_mut().unwrap()
330    }
331
332    pub fn with_max_submit_position(mut self, max_submit_position: usize) -> Self {
333        self.max_submit_position = Some(max_submit_position);
334        self
335    }
336
337    pub fn with_disable_fullnode_pruning(mut self) -> Self {
338        self.disable_fullnode_pruning = true;
339        self
340    }
341
342    pub fn with_submit_delay_step_override_millis(
343        mut self,
344        submit_delay_step_override_millis: u64,
345    ) -> Self {
346        self.submit_delay_step_override_millis = Some(submit_delay_step_override_millis);
347        self
348    }
349}
350
351impl<R: rand::RngCore + rand::CryptoRng> SwarmBuilder<R> {
352    /// Create the configured Swarm.
353    pub fn build(self) -> Swarm {
354        let dir = if let Some(dir) = self.dir {
355            SwarmDirectory::Persistent(dir)
356        } else {
357            SwarmDirectory::new_temporary()
358        };
359
360        let ingest_data = self.data_ingestion_dir.clone();
361
362        let network_config = self.network_config.unwrap_or_else(|| {
363            let mut config_builder = ConfigBuilder::new(dir.as_ref());
364
365            if let Some(genesis_config) = self.genesis_config {
366                config_builder = config_builder.with_genesis_config(genesis_config);
367            }
368
369            if let Some(chain_override) = self.chain_override {
370                config_builder = config_builder.with_chain_override(chain_override);
371            }
372
373            if let Some(num_unpruned_validators) = self.num_unpruned_validators {
374                config_builder =
375                    config_builder.with_num_unpruned_validators(num_unpruned_validators);
376            }
377
378            if let Some(jwk_fetch_interval) = self.jwk_fetch_interval {
379                config_builder = config_builder.with_jwk_fetch_interval(jwk_fetch_interval);
380            }
381
382            if let Some(authority_overload_config) = self.authority_overload_config {
383                config_builder =
384                    config_builder.with_authority_overload_config(authority_overload_config);
385            }
386
387            if let Some(execution_cache_config) = self.execution_cache_config {
388                config_builder = config_builder.with_execution_cache_config(execution_cache_config);
389            }
390
391            if let Some(path) = self.data_ingestion_dir {
392                config_builder = config_builder.with_data_ingestion_dir(path);
393            }
394
395            if let Some(max_submit_position) = self.max_submit_position {
396                config_builder = config_builder.with_max_submit_position(max_submit_position);
397            }
398
399            if let Some(submit_delay_step_override_millis) = self.submit_delay_step_override_millis
400            {
401                config_builder = config_builder
402                    .with_submit_delay_step_override_millis(submit_delay_step_override_millis);
403            }
404
405            #[allow(unused_mut)]
406            let mut final_builder = config_builder
407                .committee(self.committee)
408                .rng(self.rng)
409                .with_objects(self.additional_objects)
410                .with_supported_protocol_versions_config(
411                    self.supported_protocol_versions_config.clone(),
412                )
413                .with_global_state_hash_v2_enabled_config(
414                    self.global_state_hash_v2_enabled_config.clone(),
415                );
416
417            #[cfg(msim)]
418            if let Some(execution_time_observer_config) = self.execution_time_observer_config {
419                final_builder = final_builder
420                    .with_execution_time_observer_config(execution_time_observer_config);
421            }
422
423            final_builder.build()
424        });
425
426        let mut nodes: HashMap<_, _> = network_config
427            .validator_configs()
428            .iter()
429            .map(|config| {
430                info!(
431                    "SwarmBuilder configuring validator with name {}",
432                    config.protocol_public_key()
433                );
434                (config.protocol_public_key(), Node::new(config.to_owned()))
435            })
436            .collect();
437
438        let mut fullnode_config_builder = FullnodeConfigBuilder::new()
439            .with_config_directory(dir.as_ref().into())
440            .with_db_checkpoint_config(self.db_checkpoint_config.clone())
441            .with_run_with_range(self.fullnode_run_with_range)
442            .with_policy_config(self.fullnode_policy_config)
443            .with_data_ingestion_dir(ingest_data)
444            .with_fw_config(self.fullnode_fw_config)
445            .with_disable_pruning(self.disable_fullnode_pruning);
446
447        if let Some(chain) = self.chain_override {
448            fullnode_config_builder = fullnode_config_builder.with_chain_override(chain);
449        }
450
451        if let Some(spvc) = &self.fullnode_supported_protocol_versions_config {
452            let supported_versions = match spvc {
453                ProtocolVersionsConfig::Default => SupportedProtocolVersions::SYSTEM_DEFAULT,
454                ProtocolVersionsConfig::Global(v) => *v,
455                ProtocolVersionsConfig::PerValidator(func) => func(0, None),
456            };
457            fullnode_config_builder =
458                fullnode_config_builder.with_supported_protocol_versions(supported_versions);
459        }
460
461        if self.fullnode_count > 0 {
462            (0..self.fullnode_count).for_each(|idx| {
463                let mut builder = fullnode_config_builder.clone();
464                if idx == 0 {
465                    // Only the first fullnode is used as the rpc fullnode, we can only use the
466                    // same address once.
467                    if let Some(rpc_addr) = self.fullnode_rpc_addr {
468                        builder = builder.with_rpc_addr(rpc_addr);
469                    }
470                    if let Some(rpc_port) = self.fullnode_rpc_port {
471                        builder = builder.with_rpc_port(rpc_port);
472                    }
473                    if let Some(rpc_config) = &self.fullnode_rpc_config {
474                        builder = builder.with_rpc_config(rpc_config.clone());
475                    }
476                }
477                let config = builder.build(&mut OsRng, &network_config);
478                info!(
479                    "SwarmBuilder configuring full node with name {}",
480                    config.protocol_public_key()
481                );
482                nodes.insert(config.protocol_public_key(), Node::new(config));
483            });
484        }
485        Swarm {
486            dir,
487            network_config,
488            nodes,
489            fullnode_config_builder,
490        }
491    }
492}
493
494/// A handle to an in-memory Sui Network.
495#[derive(Debug)]
496pub struct Swarm {
497    dir: SwarmDirectory,
498    network_config: NetworkConfig,
499    nodes: HashMap<AuthorityName, Node>,
500    // Save a copy of the fullnode config builder to build future fullnodes.
501    fullnode_config_builder: FullnodeConfigBuilder,
502}
503
504impl Drop for Swarm {
505    fn drop(&mut self) {
506        self.nodes_iter_mut().for_each(|node| node.stop());
507    }
508}
509
510impl Swarm {
511    fn nodes_iter_mut(&mut self) -> impl Iterator<Item = &mut Node> {
512        self.nodes.values_mut()
513    }
514
515    /// Return a new Builder
516    pub fn builder() -> SwarmBuilder {
517        SwarmBuilder::new()
518    }
519
520    /// Start all nodes associated with this Swarm
521    pub async fn launch(&mut self) -> Result<()> {
522        try_join_all(self.nodes_iter_mut().map(|node| node.start())).await?;
523        tracing::info!("Successfully launched Swarm");
524        Ok(())
525    }
526
527    /// Return the path to the directory where this Swarm's on-disk data is kept.
528    pub fn dir(&self) -> &Path {
529        self.dir.as_ref()
530    }
531
532    /// Return a reference to this Swarm's `NetworkConfig`.
533    pub fn config(&self) -> &NetworkConfig {
534        &self.network_config
535    }
536
537    /// Return a mutable reference to this Swarm's `NetworkConfig`.
538    // TODO: It's not ideal to mutate network config. We should consider removing this.
539    pub fn config_mut(&mut self) -> &mut NetworkConfig {
540        &mut self.network_config
541    }
542
543    pub fn all_nodes(&self) -> impl Iterator<Item = &Node> {
544        self.nodes.values()
545    }
546
547    pub fn node(&self, name: &AuthorityName) -> Option<&Node> {
548        self.nodes.get(name)
549    }
550
551    pub fn node_mut(&mut self, name: &AuthorityName) -> Option<&mut Node> {
552        self.nodes.get_mut(name)
553    }
554
555    /// Return an iterator over shared references of all nodes that are set up as validators.
556    /// This means that they have a consensus config. This however doesn't mean this validator is
557    /// currently active (i.e. it's not necessarily in the validator set at the moment).
558    pub fn validator_nodes(&self) -> impl Iterator<Item = &Node> {
559        self.nodes
560            .values()
561            .filter(|node| node.config().consensus_config.is_some())
562    }
563
564    pub fn validator_node_handles(&self) -> Vec<SuiNodeHandle> {
565        self.validator_nodes()
566            .map(|node| node.get_node_handle().unwrap())
567            .collect()
568    }
569
570    /// Returns an iterator over all currently active validators.
571    pub fn active_validators(&self) -> impl Iterator<Item = &Node> {
572        self.validator_nodes().filter(|node| {
573            node.get_node_handle().is_some_and(|handle| {
574                let state = handle.state();
575                state.is_validator(&state.epoch_store_for_testing())
576            })
577        })
578    }
579
580    /// Return an iterator over shared references of all Fullnodes.
581    pub fn fullnodes(&self) -> impl Iterator<Item = &Node> {
582        self.nodes
583            .values()
584            .filter(|node| node.config().consensus_config.is_none())
585    }
586
587    pub async fn spawn_new_node(&mut self, config: NodeConfig) -> SuiNodeHandle {
588        let name = config.protocol_public_key();
589        let node = Node::new(config);
590        node.start().await.unwrap();
591        let handle = node.get_node_handle().unwrap();
592        self.nodes.insert(name, node);
593        handle
594    }
595
596    pub fn get_fullnode_config_builder(&self) -> FullnodeConfigBuilder {
597        self.fullnode_config_builder.clone()
598    }
599}
600
601#[derive(Debug)]
602enum SwarmDirectory {
603    Persistent(PathBuf),
604    Temporary(TempDir),
605}
606
607impl SwarmDirectory {
608    fn new_temporary() -> Self {
609        SwarmDirectory::Temporary(nondeterministic!(TempDir::new().unwrap()))
610    }
611}
612
613impl ops::Deref for SwarmDirectory {
614    type Target = Path;
615
616    fn deref(&self) -> &Self::Target {
617        match self {
618            SwarmDirectory::Persistent(dir) => dir.deref(),
619            SwarmDirectory::Temporary(dir) => dir.path(),
620        }
621    }
622}
623
624impl AsRef<Path> for SwarmDirectory {
625    fn as_ref(&self) -> &Path {
626        match self {
627            SwarmDirectory::Persistent(dir) => dir.as_ref(),
628            SwarmDirectory::Temporary(dir) => dir.as_ref(),
629        }
630    }
631}
632
633#[cfg(test)]
634mod test {
635    use super::Swarm;
636    use std::num::NonZeroUsize;
637
638    #[tokio::test]
639    async fn launch() {
640        telemetry_subscribers::init_for_testing();
641        let mut swarm = Swarm::builder()
642            .committee_size(NonZeroUsize::new(4).unwrap())
643            .with_fullnode_count(1)
644            .build();
645
646        swarm.launch().await.unwrap();
647
648        for validator in swarm.validator_nodes() {
649            validator.health_check(true).await.unwrap();
650        }
651
652        for fullnode in swarm.fullnodes() {
653            fullnode.health_check(false).await.unwrap();
654        }
655
656        println!("hello");
657    }
658}
sui_swarm/memory/swarm.rs

sui_swarm/memory/
swarm.rs
