iota_simulator/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5#[cfg(msim)]
6use std::hash::Hasher;
7use std::sync::atomic::{AtomicUsize, Ordering};
8
9// Re-export things used by iota-macros
10pub use ::rand as rand_crate;
11pub use anemo;
12pub use anemo_tower;
13pub use fastcrypto;
14pub use iota_framework;
15pub use iota_move_build;
16pub use iota_types;
17pub use lru;
18pub use move_package;
19#[cfg(msim)]
20pub use msim::*;
21pub use telemetry_subscribers;
22pub use tempfile;
23pub use tower;
24
25#[cfg(msim)]
26pub mod configs {
27    use std::{collections::HashMap, ops::Range, time::Duration};
28
29    use msim::*;
30    use tracing::info;
31
32    fn ms_to_dur(range: Range<u64>) -> Range<Duration> {
33        Duration::from_millis(range.start)..Duration::from_millis(range.end)
34    }
35
36    /// A network with constant uniform latency.
37    pub fn constant_latency_ms(latency: u64) -> SimConfig {
38        uniform_latency_ms(latency..(latency + 1))
39    }
40
41    /// A network with latency sampled uniformly from a range.
42    pub fn uniform_latency_ms(range: Range<u64>) -> SimConfig {
43        let range = ms_to_dur(range);
44        SimConfig {
45            net: NetworkConfig {
46                latency: LatencyConfig {
47                    default_latency: LatencyDistribution::uniform(range),
48                    ..Default::default()
49                },
50                ..Default::default()
51            },
52        }
53    }
54
55    /// A network with bimodal latency.
56    pub fn bimodal_latency_ms(
57        // The typical latency.
58        baseline: Range<u64>,
59        // The exceptional latency.
60        degraded: Range<u64>,
61        // The frequency (from 0.0 to 1.0) with which the exceptional distribution is sampled.
62        degraded_freq: f64,
63    ) -> SimConfig {
64        let baseline = ms_to_dur(baseline);
65        let degraded = ms_to_dur(degraded);
66        SimConfig {
67            net: NetworkConfig {
68                latency: LatencyConfig {
69                    default_latency: LatencyDistribution::bimodal(
70                        baseline,
71                        degraded,
72                        degraded_freq,
73                    ),
74                    ..Default::default()
75                },
76                ..Default::default()
77            },
78        }
79    }
80
81    /// Select from among a number of configs using the IOTA_SIM_CONFIG env var.
82    pub fn env_config(
83        // Config to use when IOTA_SIM_CONFIG is not set.
84        default: SimConfig,
85        // List of (&str, SimConfig) pairs - the SimConfig associated with the value
86        // of the IOTA_SIM_CONFIG var is chosen.
87        env_configs: impl IntoIterator<Item = (&'static str, SimConfig)>,
88    ) -> SimConfig {
89        let mut env_configs = HashMap::<&'static str, SimConfig>::from_iter(env_configs);
90        if let Ok(env) = std::env::var("IOTA_SIM_CONFIG") {
91            if let Some(cfg) = env_configs.remove(env.as_str()) {
92                info!("Using test config for IOTA_SIM_CONFIG={}", env);
93                cfg
94            } else {
95                panic!(
96                    "No config found for IOTA_SIM_CONFIG={}. Available configs are: {:?}",
97                    env,
98                    env_configs.keys()
99                );
100            }
101        } else {
102            info!("Using default test config");
103            default
104        }
105    }
106}
107
108thread_local! {
109    static NODE_COUNT: AtomicUsize = const { AtomicUsize::new(0) };
110}
111
112pub struct NodeLeakDetector(());
113
114impl NodeLeakDetector {
115    pub fn new() -> Self {
116        NODE_COUNT.with(|c| c.fetch_add(1, Ordering::SeqCst));
117        Self(())
118    }
119
120    pub fn get_current_node_count() -> usize {
121        NODE_COUNT.with(|c| c.load(Ordering::SeqCst))
122    }
123}
124
125impl Default for NodeLeakDetector {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl Drop for NodeLeakDetector {
132    fn drop(&mut self) {
133        NODE_COUNT.with(|c| c.fetch_sub(1, Ordering::SeqCst));
134    }
135}
136
137#[cfg(not(msim))]
138#[macro_export]
139macro_rules! return_if_killed {
140    () => {};
141}
142
143#[cfg(msim)]
144pub fn current_simnode_id() -> msim::task::NodeId {
145    msim::runtime::NodeHandle::current().id()
146}
147
148#[cfg(msim)]
149pub mod random {
150    use std::{cell::RefCell, collections::HashSet, hash::Hash};
151
152    use rand_crate::{Rng, SeedableRng, rngs::SmallRng, thread_rng};
153    use serde::Serialize;
154
155    use super::*;
156
157    /// Given a value, produce a random probability using the value as a seed,
158    /// with an additional seed that is constant only for the current test
159    /// thread.
160    pub fn deterministic_probability<T: Hash>(value: T, chance: f32) -> bool {
161        thread_local! {
162            // a random seed that is shared by the whole test process, so that equal `value`
163            // inputs produce different outputs when the test seed changes
164            static SEED: u64 = thread_rng().gen();
165        }
166
167        chance
168            > SEED.with(|seed| {
169                let mut hasher = std::collections::hash_map::DefaultHasher::new();
170                seed.hash(&mut hasher);
171                value.hash(&mut hasher);
172                let mut rng = SmallRng::seed_from_u64(hasher.finish());
173                rng.gen_range(0.0..1.0)
174            })
175    }
176
177    /// Like deterministic_probability, but only returns true once for each
178    /// unique value. May eventually consume all memory if there are a large
179    /// number of unique, failing values.
180    pub fn deterministic_probability_once<T: Hash + Serialize>(value: T, chance: f32) -> bool {
181        thread_local! {
182            static FAILING_VALUES: RefCell<HashSet<(msim::task::NodeId, Vec<u8>)>> = RefCell::new(HashSet::new());
183        }
184
185        let bytes = bcs::to_bytes(&value).unwrap();
186        let key = (current_simnode_id(), bytes);
187
188        FAILING_VALUES.with(|failing_values| {
189            let mut failing_values = failing_values.borrow_mut();
190            if failing_values.contains(&key) {
191                false
192            } else if deterministic_probability(value, chance) {
193                failing_values.insert(key);
194                true
195            } else {
196                false
197            }
198        })
199    }
200}
iota_simulator/lib.rs

iota_simulator/
lib.rs
