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 std::any::Any;
20
21use crate::chaos::state_handle::{StateHandle, Timeline};
22use crate::network::SimNetworkProvider;
23use crate::providers::{SimProviders, SimRandomProvider, SimTimeProvider};
24use crate::storage::SimStorageProvider;
25
26use moonpool_core::{Providers, TimeProvider, TokioTaskProvider};
27
28use super::topology::WorkloadTopology;
29
30/// Simulation context provided to workloads.
31///
32/// Wraps all simulation infrastructure into a single, non-generic struct.
33/// For code generic over `P: Providers`, pass `ctx.providers()`.
34pub struct SimContext {
35    providers: SimProviders,
36    topology: WorkloadTopology,
37    state: StateHandle,
38}
39
40impl SimContext {
41    /// Create a new simulation context.
42    pub fn new(providers: SimProviders, topology: WorkloadTopology, state: StateHandle) -> Self {
43        Self {
44            providers,
45            topology,
46            state,
47        }
48    }
49
50    /// Get the full providers bundle for passing to generic code.
51    pub fn providers(&self) -> &SimProviders {
52        &self.providers
53    }
54
55    /// Get the simulated network provider.
56    pub fn network(&self) -> &SimNetworkProvider {
57        self.providers.network()
58    }
59
60    /// Get the simulated time provider.
61    pub fn time(&self) -> &SimTimeProvider {
62        self.providers.time()
63    }
64
65    /// Get the task provider.
66    pub fn task(&self) -> &TokioTaskProvider {
67        self.providers.task()
68    }
69
70    /// Get the seeded random provider.
71    pub fn random(&self) -> &SimRandomProvider {
72        self.providers.random()
73    }
74
75    /// Get the simulated storage provider.
76    pub fn storage(&self) -> &SimStorageProvider {
77        self.providers.storage()
78    }
79
80    /// Get this workload's IP address.
81    pub fn my_ip(&self) -> &str {
82        &self.topology.my_ip
83    }
84
85    /// Get this workload's client ID.
86    ///
87    /// Assigned by the builder's [`ClientId`](crate::ClientId) strategy.
88    /// Defaults to sequential IDs starting from 0 (FDB-style).
89    pub fn client_id(&self) -> usize {
90        self.topology.client_id
91    }
92
93    /// Get the total number of workload instances sharing this entry.
94    ///
95    /// For single `.workload()` entries this is 1.
96    /// For `.workloads(count, factory)` entries this is the resolved count.
97    pub fn client_count(&self) -> usize {
98        self.topology.client_count
99    }
100
101    /// Find a peer's IP address by workload name.
102    pub fn peer(&self, name: &str) -> Option<String> {
103        self.topology.peer_by_name(name)
104    }
105
106    /// Get all peers as (name, ip) pairs.
107    pub fn peers(&self) -> Vec<(String, String)> {
108        self.topology
109            .peer_names
110            .iter()
111            .zip(self.topology.peer_ips.iter())
112            .map(|(name, ip)| (name.clone(), ip.clone()))
113            .collect()
114    }
115
116    /// Get the shutdown cancellation token.
117    pub fn shutdown(&self) -> &tokio_util::sync::CancellationToken {
118        &self.topology.shutdown_signal
119    }
120
121    /// Get the workload topology (peer IPs, process IPs, tags, etc.).
122    pub fn topology(&self) -> &WorkloadTopology {
123        &self.topology
124    }
125
126    /// Get the shared state handle for cross-workload communication.
127    pub fn state(&self) -> &StateHandle {
128        &self.state
129    }
130
131    /// Emit an event to a named timeline.
132    ///
133    /// Automatically captures simulation time and this process's IP address.
134    pub fn emit<T: Any + 'static>(&self, key: &str, event: T) {
135        let time_ms = self.time().now().as_millis() as u64;
136        let source = self.my_ip();
137        self.state.emit_raw(key, event, time_ms, source);
138    }
139
140    /// Get a typed timeline by name.
141    ///
142    /// Returns `None` if no events have been emitted to this key,
143    /// or if the type doesn't match.
144    pub fn timeline<T: Any + 'static>(&self, key: &str) -> Option<Timeline<T>> {
145        self.state.timeline(key)
146    }
147}