1use crate::error::RuntimeError;
4use crate::job::JobHandle;
5use crate::utils::{game_time, time_used};
6use crate::CURRENT;
7use async_task::Runnable;
8use std::cell::RefCell;
9use std::collections::BTreeMap;
10use std::future::Future;
11use std::rc::Rc;
12use std::sync::Mutex;
13use std::task::Waker;
14
15pub struct Builder {
17 config: Config,
18}
19
20impl Builder {
21 pub fn new() -> Self {
23 Self {
24 config: Config::default(),
25 }
26 }
27
28 pub fn tick_time_allocation(mut self, dur: f64) -> Self {
30 self.config.tick_time_allocation = dur;
31 self
32 }
33
34 pub fn apply(self) {
36 CURRENT.with_borrow_mut(|runtime| {
37 *runtime = Some(ScreepsRuntime::new(self.config));
38 })
39 }
40}
41
42impl Default for Builder {
43 fn default() -> Self {
44 Self::new()
45 }
46}
47
48pub struct Config {
50 tick_time_allocation: f64,
55}
56
57impl Default for Config {
58 fn default() -> Self {
59 Self {
60 tick_time_allocation: 0.9,
61 }
62 }
63}
64
65pub struct ScreepsRuntime {
72 scheduled: flume::Receiver<Runnable>,
76
77 sender: flume::Sender<Runnable>,
79
80 pub(crate) timers: Rc<Mutex<TimerMap>>,
83
84 config: Config,
86
87 is_blocking: Mutex<()>,
89}
90
91impl ScreepsRuntime {
92 pub(crate) fn new(config: Config) -> Self {
97 let (sender, scheduled) = flume::unbounded();
98
99 let timers = Rc::new(Mutex::new(BTreeMap::new()));
100
101 Self {
102 scheduled,
103 sender,
104 timers,
105 config,
106 is_blocking: Mutex::new(()),
107 }
108 }
109
110 pub fn spawn<F>(&self, future: F) -> JobHandle<F::Output>
112 where
113 F: Future + 'static,
114 {
115 let fut_res = Rc::new(RefCell::new(None));
116
117 let future = {
118 let fut_res = fut_res.clone();
119 async move {
120 let res = future.await;
121 let mut fut_res = fut_res.borrow_mut();
122 *fut_res = Some(res);
123 }
124 };
125
126 let sender = self.sender.clone();
127 let (runnable, task) = async_task::spawn_local(future, move |runnable| {
128 if !sender.is_disconnected() {
130 sender.send(runnable).unwrap();
131 }
132 });
133
134 runnable.schedule();
135
136 JobHandle::new(fut_res, task)
137 }
138
139 pub fn block_on<F>(&self, future: F) -> Result<F::Output, RuntimeError>
148 where
149 F: Future + 'static,
150 {
151 let _guard = self
152 .is_blocking
153 .try_lock()
154 .expect("Cannot block_on multiple futures at once. Please .await on the inner future");
155 let handle = self.spawn(future);
156
157 while !handle.is_complete() {
158 if !self.try_poll_scheduled()? {
159 return Err(RuntimeError::DeadlockDetected);
160 }
161 }
162
163 Ok(handle.fut_res.take().unwrap())
164 }
165
166 pub fn run(&self) -> Result<(), RuntimeError> {
173 self.wake_timers();
175
176 while self.try_poll_scheduled()? {}
178
179 Ok(())
180 }
181
182 pub(crate) fn try_poll_scheduled(&self) -> Result<bool, RuntimeError> {
188 if time_used() > self.config.tick_time_allocation {
189 return Err(RuntimeError::OutOfTime);
190 }
191
192 if let Ok(runnable) = self.scheduled.try_recv() {
193 runnable.run();
194 Ok(true)
195 } else {
196 Ok(false)
197 }
198 }
199
200 fn wake_timers(&self) {
201 let game_time = game_time();
202 let mut timers = self.timers.try_lock().unwrap();
203
204 let to_fire = {
205 let mut pending = timers.split_off(&(game_time + 1));
207 std::mem::swap(&mut pending, &mut timers);
209 pending
210 };
211
212 to_fire
213 .into_values()
214 .flatten()
215 .flatten()
216 .for_each(Waker::wake);
217 }
218}
219
220type TimerMap = BTreeMap<u32, Vec<Option<Waker>>>;
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225 use crate::error::RuntimeError::OutOfTime;
226 use crate::tests::*;
227 use crate::time::yield_now;
228 use crate::{spawn, with_runtime};
229 use std::cell::OnceCell;
230
231 #[test]
232 fn test_block_on() {
233 init_test();
234
235 let res = crate::block_on(async move {
236 yield_now().await;
237 1 + 2
238 })
239 .unwrap();
240
241 assert_eq!(3, res);
242 }
243
244 #[test]
245 fn test_spawn() {
246 init_test();
247
248 drop(spawn(async move {}));
249
250 with_runtime(|runtime| {
251 runtime
252 .scheduled
253 .try_recv()
254 .expect("Failed to schedule task");
255 })
256 }
257
258 #[test]
259 fn test_run() {
260 init_test();
261
262 let has_run = Rc::new(OnceCell::new());
263 {
264 let has_run = has_run.clone();
265 spawn(async move {
266 has_run.set(()).unwrap();
267 })
268 .detach();
269 }
270
271 assert!(has_run.get().is_none());
273
274 crate::run().unwrap();
275
276 assert!(has_run.get().is_some());
278 }
279
280 #[test]
281 fn test_nested_spawn() {
282 init_test();
283
284 let has_run = Rc::new(OnceCell::new());
285 {
286 let has_run = has_run.clone();
287 spawn(async move {
288 let result = spawn(async move { 1 + 2 }).await;
289
290 assert_eq!(3, result);
291
292 has_run.set(()).unwrap();
293 })
294 .detach();
295 }
296
297 assert!(has_run.get().is_none());
299
300 crate::run().unwrap();
301
302 assert!(has_run.get().is_some());
304 }
305
306 #[test]
307 fn test_respects_time_remaining() {
308 init_test();
309
310 let has_run = Rc::new(OnceCell::new());
311 {
312 let has_run = has_run.clone();
313 spawn(async move {
314 has_run.set(()).unwrap();
315 })
316 .detach();
317 }
318
319 TIME_USED.with_borrow_mut(|t| *t = 0.95);
320
321 assert!(has_run.get().is_none());
323
324 assert_eq!(Err(OutOfTime), crate::run());
325
326 assert!(has_run.get().is_none());
328 }
329}