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}