Skip to main content

moonpool_sim/runner/
context.rs

1//! Simulation context for workloads.
2//!
3//! [`SimContext`] is the single entry point for workloads to access simulation
4//! infrastructure: providers, topology, shared state, and shutdown signaling.
5//!
6//! # Usage
7//!
8//! ```ignore
9//! use moonpool_sim::SimContext;
10//!
11//! async fn my_workload(ctx: &SimContext) -> SimulationResult<()> {
12//!     let server_ip = ctx.peer("server").expect("server not found");
13//!     let stream = ctx.network().connect(&server_ip).await?;
14//!     ctx.state().publish("connected", true);
15//!     Ok(())
16//! }
17//! ```
18
19use crate::chaos::state_handle::StateHandle;
20use crate::network::SimNetworkProvider;
21use crate::observability::SimulationLayerHandle;
22use crate::providers::{SimProviders, SimRandomProvider, SimTimeProvider};
23use crate::storage::SimStorageProvider;
24
25use moonpool_core::Providers;
26
27use super::topology::WorkloadTopology;
28
29/// Simulation context provided to workloads.
30///
31/// Wraps all simulation infrastructure into a single, non-generic struct.
32/// For code generic over `P: Providers`, pass `ctx.providers()`.
33pub struct SimContext {
34    providers: SimProviders,
35    topology: WorkloadTopology,
36    state: StateHandle,
37    obs: SimulationLayerHandle,
38}
39
40impl SimContext {
41    /// Create a new simulation context.
42    #[must_use]
43    pub fn new(
44        providers: SimProviders,
45        topology: WorkloadTopology,
46        state: StateHandle,
47        obs: SimulationLayerHandle,
48    ) -> Self {
49        Self {
50            providers,
51            topology,
52            state,
53            obs,
54        }
55    }
56
57    /// Get the full providers bundle for passing to generic code.
58    #[must_use]
59    pub fn providers(&self) -> &SimProviders {
60        &self.providers
61    }
62
63    /// Get the simulated network provider.
64    #[must_use]
65    pub fn network(&self) -> &SimNetworkProvider {
66        self.providers.network()
67    }
68
69    /// Get the simulated time provider.
70    #[must_use]
71    pub fn time(&self) -> &SimTimeProvider {
72        self.providers.time()
73    }
74
75    /// Get the task provider.
76    #[must_use]
77    pub fn task(&self) -> &crate::providers::SimTaskProvider {
78        self.providers.task()
79    }
80
81    /// Get the seeded random provider.
82    #[must_use]
83    pub fn random(&self) -> &SimRandomProvider {
84        self.providers.random()
85    }
86
87    /// Get the simulated storage provider.
88    #[must_use]
89    pub fn storage(&self) -> &SimStorageProvider {
90        self.providers.storage()
91    }
92
93    /// Get this workload's IP address.
94    #[must_use]
95    pub fn my_ip(&self) -> &str {
96        &self.topology.my_ip
97    }
98
99    /// Get this workload's client ID.
100    ///
101    /// Assigned by the builder's [`ClientId`](crate::ClientId) strategy.
102    /// Defaults to sequential IDs starting from 0 (FDB-style).
103    #[must_use]
104    pub fn client_id(&self) -> usize {
105        self.topology.client_id
106    }
107
108    /// Get the total number of workload instances sharing this entry.
109    ///
110    /// For single `.workload()` entries this is 1.
111    /// For `.workloads(count, factory)` entries this is the resolved count.
112    #[must_use]
113    pub fn client_count(&self) -> usize {
114        self.topology.client_count
115    }
116
117    /// Find a peer's IP address by workload name.
118    #[must_use]
119    pub fn peer(&self, name: &str) -> Option<String> {
120        self.topology.peer_by_name(name)
121    }
122
123    /// Get all peers as (name, ip) pairs.
124    #[must_use]
125    pub fn peers(&self) -> Vec<(String, String)> {
126        self.topology
127            .peer_names
128            .iter()
129            .zip(self.topology.peer_ips.iter())
130            .map(|(name, ip)| (name.clone(), ip.clone()))
131            .collect()
132    }
133
134    /// Get the shutdown cancellation token.
135    #[must_use]
136    pub fn shutdown(&self) -> &tokio_util::sync::CancellationToken {
137        &self.topology.shutdown_signal
138    }
139
140    /// Get the workload topology (peer IPs, process IPs, tags, etc.).
141    #[must_use]
142    pub fn topology(&self) -> &WorkloadTopology {
143        &self.topology
144    }
145
146    /// Get the shared state handle for cross-workload communication.
147    #[must_use]
148    pub fn state(&self) -> &StateHandle {
149        &self.state
150    }
151
152    /// Get a clonable handle to the observability layer.
153    ///
154    /// The handle implements [`crate::TraceQuery`], so workloads can read
155    /// the captured timeline (e.g. in `check()`) the same way invariants do.
156    /// To get events INTO the timeline, emit plain `tracing` events — e.g.
157    /// `tracing::info!(term, leader = %ip, "leader_elected")` — from inside
158    /// a process or workload task; the orchestrator's actor spans attribute
159    /// them automatically.
160    #[must_use]
161    pub fn observability(&self) -> &SimulationLayerHandle {
162        &self.obs
163    }
164}