steel_core/chunk/
chunk_generation_task.rs

1//! `ChunkGenerationTask` handles the generation process for chunks.
2use std::{
3    cmp::max,
4    future::{Future, poll_fn},
5    mem,
6    pin::Pin,
7    sync::{
8        Arc,
9        atomic::{AtomicBool, Ordering},
10    },
11    task::Poll,
12};
13
14use rayon::ThreadPool;
15use steel_utils::{ChunkPos, locks::SyncMutex};
16use tokio_util::sync::CancellationToken;
17
18use crate::chunk::{
19    chunk_holder::{ChunkHolder, ChunkSaveDependency},
20    chunk_map::ChunkMap,
21    chunk_pyramid::{GENERATION_PYRAMID, LOADING_PYRAMID},
22    status::ChunkStatus,
23};
24
25/// A pre-filled 2D cache of elements, efficient for async creation.
26pub struct StaticCache2D<T> {
27    min_x: i32,
28    min_z: i32,
29    size: i32,
30    /// Cache stored in row-major order (Z-then-X).
31    cache: Vec<T>,
32    /// Pins generation holders until the last user of this cache finishes.
33    save_dependencies: Option<Box<[ChunkSaveDependency]>>,
34}
35
36impl<T> StaticCache2D<T> {
37    /// Creates a `StaticCache2D` by populating it via a factory.
38    pub fn create<F>(center_x: i32, center_z: i32, radius: i32, factory: F) -> Self
39    where
40        F: Fn(i32, i32) -> T + Send + Sync + 'static,
41        T: Send + 'static,
42    {
43        let size = radius * 2 + 1;
44        let min_x = center_x - radius;
45        let min_z = center_z - radius;
46        let cap = (size * size) as usize;
47        let size_usize = size as usize;
48
49        let cache: Vec<T> = (0..cap)
50            .map(|index| {
51                let x_offset = (index % size_usize) as i32;
52                let z_offset = (index / size_usize) as i32;
53                factory(min_x + x_offset, min_z + z_offset)
54            })
55            .collect();
56
57        Self {
58            min_x,
59            min_z,
60            size,
61            cache,
62            save_dependencies: None,
63        }
64    }
65
66    /// Gets a reference to an element by world coordinates.
67    ///
68    /// # Panics
69    /// Panics if coordinates are out of bounds.
70    #[must_use]
71    pub fn get(&self, x: i32, z: i32) -> &T {
72        let Some(value) = self.try_get(x, z) else {
73            panic!(
74                "Out of bounds: ({x}, {z}) vs [({}, {}) to ({}, {})]",
75                self.min_x,
76                self.min_z,
77                self.min_x + self.size - 1,
78                self.min_z + self.size - 1
79            );
80        };
81        value
82    }
83
84    /// Gets a reference to an element by world coordinates.
85    #[must_use]
86    pub fn try_get(&self, x: i32, z: i32) -> Option<&T> {
87        let rel_x = x - self.min_x;
88        let rel_z = z - self.min_z;
89
90        if rel_x >= 0 && rel_x < self.size && rel_z >= 0 && rel_z < self.size {
91            let index = (rel_z * self.size + rel_x) as usize;
92            self.cache.get(index)
93        } else {
94            None
95        }
96    }
97}
98
99impl StaticCache2D<Arc<ChunkHolder>> {
100    fn pin_holders_for_generation(&mut self) {
101        self.save_dependencies = Some(
102            self.cache
103                .iter()
104                .map(ChunkHolder::add_save_dependency)
105                .collect::<Vec<_>>()
106                .into_boxed_slice(),
107        );
108    }
109}
110
111/// A pinned future representing a neighbor's readiness.
112pub type NeighborReady = Pin<Box<dyn Future<Output = Option<()>> + Send + Sync>>;
113
114/// A task responsible for driving a chunk to a target status.
115///
116/// It works in form of layers. Imagine a pyramid, to get to the top you first need to generate the base layer. And so on.
117/// This works in the same way but with chunk dependencies.
118///
119/// This is achieved using the Generation Pyramid and Loading Pyramid.
120///
121/// To make sure a chunk is only put through a stage once it uses an atomic with a CAS operation. Loading Pyramid must also be advanced with noop functions so this atomic can be driven forward.
122pub struct ChunkGenerationTask {
123    /// The chunk map associated with this task.
124    pub chunk_map: Arc<ChunkMap>,
125    /// The chunk position.
126    pub pos: ChunkPos,
127    /// The target generation status.
128    pub target_status: ChunkStatus,
129    /// The status scheduled for generation. Protected by a mutex for safe concurrent access.
130    pub scheduled_status: SyncMutex<Option<ChunkStatus>>,
131    /// Cancellation token — cancelled when this task should stop.
132    pub cancel_token: CancellationToken,
133    /// Cheap cancellation flag for scheduler-side filtering.
134    cancelled: AtomicBool,
135    /// Futures for neighbors. Protected by a mutex.
136    pub neighbor_ready: SyncMutex<Vec<NeighborReady>>,
137    /// Cache of required chunks.
138    pub cache: Arc<StaticCache2D<Arc<ChunkHolder>>>,
139    /// Holder for the chunk this task is targeting.
140    pub center_holder: Arc<ChunkHolder>,
141    /// Whether generation is required for this task.
142    pub needs_generation: AtomicBool,
143    /// The thread pool to use for generation.
144    pub thread_pool: Arc<ThreadPool>,
145}
146
147impl ChunkGenerationTask {
148    /// Creates a new generation task.
149    #[must_use]
150    #[inline]
151    #[expect(
152        clippy::missing_panics_doc,
153        reason = "panic is unreachable: ThreadPoolBuilder::build only fails on OS thread errors"
154    )]
155    pub fn new(
156        pos: ChunkPos,
157        target_status: ChunkStatus,
158        chunk_map: Arc<ChunkMap>,
159        thread_pool: Arc<ThreadPool>,
160        cancel_token: CancellationToken,
161    ) -> Self {
162        let worst_case_radius = GENERATION_PYRAMID
163            .get_step_to(target_status)
164            .accumulated_dependencies
165            .get_radius_of(ChunkStatus::Empty) as i32;
166
167        let chunk_map_clone = chunk_map.clone();
168        let mut cache = StaticCache2D::create(pos.0.x, pos.0.y, worst_case_radius, move |x, y| {
169            chunk_map_clone
170                    .chunks
171                    .read_sync(&ChunkPos::new(x, y), |_, chunk_holder| chunk_holder.clone())
172                    .expect("The chunkholder should be created by distance manager before the generation task is scheduled. This occurring means there is a bug in the distance manager or you called this yourself.")
173        });
174        cache.pin_holders_for_generation();
175        let center_holder = Arc::clone(cache.get(pos.0.x, pos.0.y));
176
177        Self {
178            chunk_map,
179            pos,
180            target_status,
181            scheduled_status: SyncMutex::new(None),
182            cancel_token,
183            cancelled: AtomicBool::new(false),
184            neighbor_ready: SyncMutex::new(Vec::new()),
185            cache: Arc::new(cache),
186            center_holder,
187            needs_generation: AtomicBool::new(true),
188            thread_pool,
189        }
190    }
191
192    /// Cancels this task by triggering the cancellation token.
193    pub fn cancel(&self) {
194        if !self.cancelled.swap(true, Ordering::AcqRel) {
195            self.cancel_token.cancel();
196        }
197    }
198
199    /// Returns whether this task has been explicitly cancelled.
200    #[must_use]
201    pub fn is_cancelled(&self) -> bool {
202        self.cancelled.load(Ordering::Acquire)
203    }
204
205    /// Returns the holder for the chunk this task is targeting.
206    pub(crate) const fn center_holder(&self) -> &Arc<ChunkHolder> {
207        &self.center_holder
208    }
209
210    /// Schedules a chunk for a specific layer.
211    ///
212    /// # Panics
213    /// Panics if generation is required but not expected.
214    pub fn schedule_chunk_in_layer(
215        &self,
216        status: ChunkStatus,
217        needs_generation: bool,
218        chunk_holder: &Arc<ChunkHolder>,
219    ) -> bool {
220        let published_status = chunk_holder.published_status();
221
222        let generate;
223        if let Some(published_status) = published_status {
224            generate = status > published_status;
225        } else {
226            generate = true;
227        }
228
229        let pyramid = if generate {
230            &GENERATION_PYRAMID
231        } else {
232            &LOADING_PYRAMID
233        };
234
235        assert!(
236            !generate || needs_generation,
237            "Generation required but not expected for chunk load"
238        );
239
240        if let Some(future) = chunk_holder.apply_step(
241            pyramid.get_step_to(status),
242            &self.chunk_map,
243            &self.cache,
244            self.thread_pool.clone(),
245        ) {
246            self.neighbor_ready.lock().push(future);
247        } else {
248            self.cancel();
249        }
250
251        true
252    }
253
254    /// Schedules tasks for the current layer's neighbors.
255    pub fn schedule_layer(&self, status: ChunkStatus, needs_generation: bool) {
256        let radius = self.get_radius_for_layer(status, needs_generation);
257        // This for loop is inclusive, so if the radius is 0, we will only schedule the center chunk.
258        for x in (self.pos.0.x - radius)..=(self.pos.0.x + radius) {
259            for y in (self.pos.0.y - radius)..=(self.pos.0.y + radius) {
260                let chunk_holder = self.cache.get(x, y);
261                if self.is_cancelled()
262                    || !self.schedule_chunk_in_layer(status, needs_generation, chunk_holder)
263                {
264                    return;
265                }
266            }
267        }
268    }
269
270    const fn get_radius_for_layer(&self, status: ChunkStatus, needs_generation: bool) -> i32 {
271        let pyramid = if needs_generation {
272            &GENERATION_PYRAMID
273        } else {
274            &LOADING_PYRAMID
275        };
276
277        pyramid
278            .get_step_to(self.target_status)
279            .get_accumulated_radius_of(status) as i32
280    }
281
282    /// Schedules the next layer of generation dependencies.
283    ///
284    /// # Panics
285    /// Panics if the schedule is invalid.
286    pub fn schedule_next_layer(&self) {
287        let status_to_schedule = if self.scheduled_status.lock().is_none() {
288            ChunkStatus::Empty
289        } else if !self.needs_generation.load(Ordering::Relaxed)
290            && *self.scheduled_status.lock() == Some(ChunkStatus::Empty)
291            && !self.can_load_without_generation()
292        {
293            self.needs_generation.store(true, Ordering::Relaxed);
294            ChunkStatus::Empty
295        } else {
296            self.scheduled_status
297                .lock()
298                .expect("Scheduled status missing")
299                .next()
300                .expect("Next status missing")
301        };
302
303        self.schedule_layer(
304            status_to_schedule,
305            self.needs_generation.load(Ordering::Relaxed),
306        );
307        self.scheduled_status.lock().replace(status_to_schedule);
308    }
309
310    fn can_load_without_generation(&self) -> bool {
311        if self.target_status == ChunkStatus::Empty {
312            return true;
313        }
314        let center = self.cache.get(self.pos.0.x, self.pos.0.y);
315        let highest_generated_status = center.published_status();
316
317        if let Some(highest_status) = highest_generated_status {
318            if highest_status < self.target_status {
319                return false;
320            }
321
322            let dependencies = &LOADING_PYRAMID
323                .get_step_to(self.target_status)
324                .accumulated_dependencies;
325            let range = dependencies.get_radius() as i32;
326
327            for x in (self.pos.0.x - range)..=(self.pos.0.x + range) {
328                for z in (self.pos.0.y - range)..=(self.pos.0.y + range) {
329                    let distance = max((self.pos.0.x - x).abs(), (self.pos.0.y - z).abs()) as usize;
330                    if let Some(required_status) = dependencies.get(distance) {
331                        let neighbor = self.cache.get(x, z);
332                        let published = neighbor.published_status();
333                        if published < Some(required_status) {
334                            return false;
335                        }
336                    }
337                }
338            }
339            true
340        } else {
341            false
342        }
343    }
344
345    /// Runs the generation task loop.
346    pub async fn run(self: Arc<Self>) {
347        loop {
348            tokio::select! {
349                () = self.cancel_token.cancelled() => break,
350                () = self.wait_for_scheduled_layers() => {}
351            }
352
353            if *self.scheduled_status.lock() == Some(self.target_status) {
354                break;
355            }
356
357            self.schedule_next_layer();
358        }
359    }
360
361    /// Waits for all scheduled neighbor tasks to complete.
362    pub async fn wait_for_scheduled_layers(&self) {
363        // Collect all futures first to avoid locking the mutex during await
364        let mut futures: Vec<_> = {
365            let mut lock = self.neighbor_ready.lock();
366            mem::take(&mut *lock)
367        };
368
369        if futures.is_empty() {
370            return;
371        }
372
373        let mut failed = false;
374        poll_fn(|cx| {
375            let mut index = 0;
376            while index < futures.len() {
377                match futures[index].as_mut().poll(cx) {
378                    Poll::Ready(result) => {
379                        failed |= result.is_none();
380                        drop(futures.swap_remove(index));
381                    }
382                    Poll::Pending => index += 1,
383                }
384            }
385
386            if futures.is_empty() {
387                Poll::Ready(())
388            } else {
389                Poll::Pending
390            }
391        })
392        .await;
393
394        if failed {
395            self.cancel();
396        }
397    }
398}
steel_core/chunk/chunk_generation_task.rs

steel_core/chunk/
chunk_generation_task.rs
