Skip to main content

steel_core/server/jobs/
mod.rs

1//! Tick-polled server jobs.
2
3pub(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/// Result of polling a server job.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum JobPoll {
20    /// Poll the job again on a later server job tick.
21    Pending,
22    /// The job is complete and should be removed.
23    Finished,
24}
25
26/// Context passed to jobs when they are polled.
27pub struct ServerJobContext {
28    server: Option<Weak<Server>>,
29    /// Current server tick count.
30    pub tick_count: u64,
31    /// Whether normal world ticking is currently running.
32    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    /// Returns the server if it is still alive.
45    #[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
60/// A unit of server work resumed from a known tick stage.
61pub trait ServerJob: Send {
62    /// Polls this job.
63    fn poll(&mut self, context: &mut ServerJobContext) -> JobPoll;
64
65    /// Cancels the job before it finishes.
66    fn cancel(&mut self) {}
67}
68
69/// Tick-owned job queue.
70#[derive(Default)]
71pub struct ServerJobQueue {
72    jobs: SyncMutex<Vec<Box<dyn ServerJob>>>,
73}
74
75/// One tick-stage callback represented as a server job.
76pub(crate) struct FnServerJob<F> {
77    action: Option<F>,
78}
79
80impl<F> FnServerJob<F> {
81    /// Creates a one-shot tick-stage callback.
82    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    /// Creates an empty job queue.
103    #[must_use]
104    pub const fn new() -> Self {
105        Self {
106            jobs: SyncMutex::new(Vec::new()),
107        }
108    }
109
110    /// Adds a job to be polled on the next server job tick.
111    pub fn spawn(&self, job: impl ServerJob + 'static) {
112        self.jobs.lock().push(Box::new(job));
113    }
114
115    /// Polls a job immediately from a known safe point, queuing it only if it remains pending.
116    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    /// Returns the number of queued jobs.
128    #[must_use]
129    pub fn len(&self) -> usize {
130        self.jobs.lock().len()
131    }
132
133    /// Returns true if no jobs are queued.
134    #[must_use]
135    pub fn is_empty(&self) -> bool {
136        self.jobs.lock().is_empty()
137    }
138
139    /// Cancels and removes all queued jobs.
140    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    /// Polls queued jobs from the server game tick.
148    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/// Job polling counts for diagnostics.
204#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
205pub struct ServerJobTickStats {
206    /// Jobs polled this tick.
207    pub polled: usize,
208    /// Jobs completed this tick.
209    pub finished: usize,
210    /// Jobs queued after this tick.
211    pub pending: usize,
212}
213
214/// A chunk request that invokes a callback once all requested chunks are ready.
215pub struct ChunkRequestJob<F> {
216    request: ChunkRequestHandle,
217    on_ready: Option<F>,
218}
219
220impl<F> ChunkRequestJob<F> {
221    /// Creates a job around a chunk request and a tick-stage callback.
222    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}