1pub(in crate::server) mod domain_switch;
4pub(super) mod teleport;
5
6use std::{
7 mem,
8 sync::{Arc, Weak},
9};
10
11use crate::{
12 chunk::chunk_request::{ChunkRequestHandle, ChunkRequestState, ReadyChunks},
13 server::Server,
14};
15use steel_utils::locks::SyncMutex;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum JobPoll {
20 Pending,
22 Finished,
24}
25
26pub struct ServerJobContext {
28 server: Option<Weak<Server>>,
29 pub tick_count: u64,
31 pub runs_normally: bool,
33}
34
35impl ServerJobContext {
36 const fn for_server(server: Weak<Server>, tick_count: u64, runs_normally: bool) -> Self {
37 Self {
38 server: Some(server),
39 tick_count,
40 runs_normally,
41 }
42 }
43
44 #[must_use]
46 pub fn server(&self) -> Option<Arc<Server>> {
47 self.server.as_ref().and_then(Weak::upgrade)
48 }
49
50 #[cfg(test)]
51 const fn for_test(tick_count: u64, runs_normally: bool) -> Self {
52 Self {
53 server: None,
54 tick_count,
55 runs_normally,
56 }
57 }
58}
59
60pub trait ServerJob: Send {
62 fn poll(&mut self, context: &mut ServerJobContext) -> JobPoll;
64
65 fn cancel(&mut self) {}
67}
68
69#[derive(Default)]
71pub struct ServerJobQueue {
72 jobs: SyncMutex<Vec<Box<dyn ServerJob>>>,
73}
74
75pub(crate) struct FnServerJob<F> {
77 action: Option<F>,
78}
79
80impl<F> FnServerJob<F> {
81 pub(crate) const fn new(action: F) -> Self {
83 Self {
84 action: Some(action),
85 }
86 }
87}
88
89impl<F> ServerJob for FnServerJob<F>
90where
91 F: FnOnce(&mut ServerJobContext) + Send,
92{
93 fn poll(&mut self, context: &mut ServerJobContext) -> JobPoll {
94 if let Some(action) = self.action.take() {
95 action(context);
96 }
97 JobPoll::Finished
98 }
99}
100
101impl ServerJobQueue {
102 #[must_use]
104 pub const fn new() -> Self {
105 Self {
106 jobs: SyncMutex::new(Vec::new()),
107 }
108 }
109
110 pub fn spawn(&self, job: impl ServerJob + 'static) {
112 self.jobs.lock().push(Box::new(job));
113 }
114
115 pub fn poll_now_or_spawn(
117 &self,
118 server: Weak<Server>,
119 tick_count: u64,
120 runs_normally: bool,
121 job: impl ServerJob + 'static,
122 ) -> JobPoll {
123 let mut context = ServerJobContext::for_server(server, tick_count, runs_normally);
124 self.poll_now_or_spawn_with_context(&mut context, job)
125 }
126
127 #[must_use]
129 pub fn len(&self) -> usize {
130 self.jobs.lock().len()
131 }
132
133 #[must_use]
135 pub fn is_empty(&self) -> bool {
136 self.jobs.lock().is_empty()
137 }
138
139 pub fn cancel_all(&self) {
141 let jobs = mem::take(&mut *self.jobs.lock());
142 for mut job in jobs {
143 job.cancel();
144 }
145 }
146
147 pub fn tick(
149 &self,
150 server: Weak<Server>,
151 tick_count: u64,
152 runs_normally: bool,
153 ) -> ServerJobTickStats {
154 let mut context = ServerJobContext::for_server(server, tick_count, runs_normally);
155 self.tick_with_context(&mut context)
156 }
157
158 fn tick_with_context(&self, context: &mut ServerJobContext) -> ServerJobTickStats {
159 let jobs = mem::take(&mut *self.jobs.lock());
160 let polled = jobs.len();
161 let mut pending = Vec::with_capacity(polled);
162 let mut finished = 0;
163
164 for mut job in jobs {
165 match job.poll(context) {
166 JobPoll::Pending => pending.push(job),
167 JobPoll::Finished => finished += 1,
168 }
169 }
170
171 let pending_count = {
172 let mut queued = self.jobs.lock();
173 let spawned_during_poll = queued.len();
174 pending.reserve(spawned_during_poll);
175 pending.append(&mut *queued);
176 let pending_count = pending.len();
177 *queued = pending;
178 pending_count
179 };
180
181 ServerJobTickStats {
182 polled,
183 finished,
184 pending: pending_count,
185 }
186 }
187
188 fn poll_now_or_spawn_with_context(
189 &self,
190 context: &mut ServerJobContext,
191 mut job: impl ServerJob + 'static,
192 ) -> JobPoll {
193 match job.poll(context) {
194 JobPoll::Pending => {
195 self.spawn(job);
196 JobPoll::Pending
197 }
198 JobPoll::Finished => JobPoll::Finished,
199 }
200 }
201}
202
203#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
205pub struct ServerJobTickStats {
206 pub polled: usize,
208 pub finished: usize,
210 pub pending: usize,
212}
213
214pub struct ChunkRequestJob<F> {
216 request: ChunkRequestHandle,
217 on_ready: Option<F>,
218}
219
220impl<F> ChunkRequestJob<F> {
221 pub const fn new(request: ChunkRequestHandle, on_ready: F) -> Self {
223 Self {
224 request,
225 on_ready: Some(on_ready),
226 }
227 }
228}
229
230impl<F> ServerJob for ChunkRequestJob<F>
231where
232 F: FnOnce(&mut ServerJobContext, ReadyChunks) + Send + 'static,
233{
234 fn poll(&mut self, context: &mut ServerJobContext) -> JobPoll {
235 match self.request.poll() {
236 ChunkRequestState::Pending { .. } => JobPoll::Pending,
237 ChunkRequestState::Cancelled => JobPoll::Finished,
238 ChunkRequestState::Ready => {
239 let Some(ready) = self.request.ready_chunks() else {
240 return JobPoll::Pending;
241 };
242 if let Some(on_ready) = self.on_ready.take() {
243 on_ready(context, ready);
244 }
245 JobPoll::Finished
246 }
247 }
248 }
249
250 fn cancel(&mut self) {
251 self.request.cancel();
252 }
253}
254
255#[cfg(test)]
256mod tests {
257 use std::sync::{
258 Arc,
259 atomic::{AtomicUsize, Ordering},
260 };
261
262 use super::*;
263
264 struct CountJob {
265 polls: Arc<AtomicUsize>,
266 finish_after: usize,
267 }
268
269 impl ServerJob for CountJob {
270 fn poll(&mut self, _context: &mut ServerJobContext) -> JobPoll {
271 let polls = self.polls.fetch_add(1, Ordering::Relaxed) + 1;
272 if polls >= self.finish_after {
273 JobPoll::Finished
274 } else {
275 JobPoll::Pending
276 }
277 }
278 }
279
280 #[test]
281 fn pending_job_is_polled_until_finished() {
282 let queue = ServerJobQueue::new();
283 let polls = Arc::new(AtomicUsize::new(0));
284 queue.spawn(CountJob {
285 polls: polls.clone(),
286 finish_after: 2,
287 });
288
289 let mut context = ServerJobContext::for_test(1, true);
290 let first = queue.tick_with_context(&mut context);
291 assert_eq!(
292 first,
293 ServerJobTickStats {
294 polled: 1,
295 finished: 0,
296 pending: 1,
297 }
298 );
299
300 let second = queue.tick_with_context(&mut context);
301 assert_eq!(
302 second,
303 ServerJobTickStats {
304 polled: 1,
305 finished: 1,
306 pending: 0,
307 }
308 );
309 assert_eq!(polls.load(Ordering::Relaxed), 2);
310 }
311
312 #[test]
313 fn jobs_spawned_during_poll_wait_until_next_tick() {
314 struct SpawnJob {
315 queue: Arc<ServerJobQueue>,
316 polls: Arc<AtomicUsize>,
317 }
318
319 impl ServerJob for SpawnJob {
320 fn poll(&mut self, _context: &mut ServerJobContext) -> JobPoll {
321 self.queue.spawn(CountJob {
322 polls: self.polls.clone(),
323 finish_after: 1,
324 });
325 JobPoll::Finished
326 }
327 }
328
329 let queue = Arc::new(ServerJobQueue::new());
330 let polls = Arc::new(AtomicUsize::new(0));
331 queue.spawn(SpawnJob {
332 queue: queue.clone(),
333 polls: polls.clone(),
334 });
335
336 let mut context = ServerJobContext::for_test(1, true);
337 let first = queue.tick_with_context(&mut context);
338 assert_eq!(first.polled, 1);
339 assert_eq!(first.finished, 1);
340 assert_eq!(first.pending, 1);
341 assert_eq!(polls.load(Ordering::Relaxed), 0);
342
343 let second = queue.tick_with_context(&mut context);
344 assert_eq!(second.polled, 1);
345 assert_eq!(second.finished, 1);
346 assert_eq!(second.pending, 0);
347 assert_eq!(polls.load(Ordering::Relaxed), 1);
348 }
349
350 #[test]
351 fn poll_now_or_spawn_drops_finished_job() {
352 let queue = ServerJobQueue::new();
353 let polls = Arc::new(AtomicUsize::new(0));
354 let mut context = ServerJobContext::for_test(1, true);
355
356 let result = queue.poll_now_or_spawn_with_context(
357 &mut context,
358 CountJob {
359 polls: polls.clone(),
360 finish_after: 1,
361 },
362 );
363
364 assert_eq!(result, JobPoll::Finished);
365 assert!(queue.is_empty());
366 assert_eq!(polls.load(Ordering::Relaxed), 1);
367 }
368
369 #[test]
370 fn poll_now_or_spawn_queues_pending_job_after_first_poll() {
371 let queue = ServerJobQueue::new();
372 let polls = Arc::new(AtomicUsize::new(0));
373 let mut context = ServerJobContext::for_test(1, true);
374
375 let result = queue.poll_now_or_spawn_with_context(
376 &mut context,
377 CountJob {
378 polls: polls.clone(),
379 finish_after: 2,
380 },
381 );
382
383 assert_eq!(result, JobPoll::Pending);
384 assert_eq!(queue.len(), 1);
385 assert_eq!(polls.load(Ordering::Relaxed), 1);
386
387 let stats = queue.tick_with_context(&mut context);
388 assert_eq!(stats.finished, 1);
389 assert_eq!(stats.pending, 0);
390 assert_eq!(polls.load(Ordering::Relaxed), 2);
391 }
392}