๐Ÿ“ฆ bevyengine / bevy

๐Ÿ“„ async_compute.rs ยท 162 lines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162//! This example demonstrates how to use Bevy's ECS and the [`AsyncComputeTaskPool`]
//! to offload computationally intensive tasks to a background thread pool, process them
//! asynchronously, and apply the results across systems and ticks.
//!
//! Unlike the channel-based approach (where tasks send results directly via a communication
//! channel), this example uses the `AsyncComputeTaskPool` to run tasks in the background,
//! check for their completion, and handle results when the task is ready. This method allows
//! tasks to be processed in parallel without blocking the main thread, but requires periodically
//! checking the status of each task.
//!
//! The channel-based approach, on the other hand, detaches tasks and communicates results
//! through a channel, avoiding the need to check task statuses manually.

use bevy::{
    ecs::{system::SystemState, world::CommandQueue},
    prelude::*,
    tasks::{futures::check_ready, AsyncComputeTaskPool, Task},
};
use futures_timer::Delay;
use rand::Rng;
use std::time::Duration;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_systems(Startup, (setup_env, add_assets, spawn_tasks))
        .add_systems(Update, handle_tasks)
        .run();
}

// Number of cubes to spawn across the x, y, and z axis
const NUM_CUBES: u32 = 6;

#[derive(Resource, Deref)]
struct BoxMeshHandle(Handle<Mesh>);

#[derive(Resource, Deref)]
struct BoxMaterialHandle(Handle<StandardMaterial>);

/// Startup system which runs only once and generates our Box Mesh
/// and Box Material assets, adds them to their respective Asset
/// Resources, and stores their handles as resources so we can access
/// them later when we're ready to render our Boxes
fn add_assets(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    let box_mesh_handle = meshes.add(Cuboid::new(0.25, 0.25, 0.25));
    commands.insert_resource(BoxMeshHandle(box_mesh_handle));

    let box_material_handle = materials.add(Color::srgb(1.0, 0.2, 0.3));
    commands.insert_resource(BoxMaterialHandle(box_material_handle));
}

#[derive(Component)]
struct ComputeTransform(Task<CommandQueue>);

/// This system generates tasks simulating computationally intensive
/// work that potentially spans multiple frames/ticks. A separate
/// system, [`handle_tasks`], will track the spawned tasks on subsequent
/// frames/ticks, and use the results to spawn cubes.
///
/// The task is offloaded to the `AsyncComputeTaskPool`, allowing heavy computation
/// to be handled asynchronously, without blocking the main game thread.
fn spawn_tasks(mut commands: Commands) {
    let thread_pool = AsyncComputeTaskPool::get();
    for x in 0..NUM_CUBES {
        for y in 0..NUM_CUBES {
            for z in 0..NUM_CUBES {
                // Spawn new task on the AsyncComputeTaskPool; the task will be
                // executed in the background, and the Task future returned by
                // spawn() can be used to poll for the result
                let entity = commands.spawn_empty().id();
                let task = thread_pool.spawn(async move {
                    let duration = Duration::from_secs_f32(rand::rng().random_range(0.05..5.0));

                    // Pretend this is a time-intensive function. :)
                    Delay::new(duration).await;

                    // Such hard work, all done!
                    let transform = Transform::from_xyz(x as f32, y as f32, z as f32);
                    let mut command_queue = CommandQueue::default();

                    // we use a raw command queue to pass a FnOnce(&mut World) back to be
                    // applied in a deferred manner.
                    command_queue.push(move |world: &mut World| {
                        let (box_mesh_handle, box_material_handle) = {
                            let mut system_state = SystemState::<(
                                Res<BoxMeshHandle>,
                                Res<BoxMaterialHandle>,
                            )>::new(world);
                            let (box_mesh_handle, box_material_handle) =
                                system_state.get_mut(world);

                            (box_mesh_handle.clone(), box_material_handle.clone())
                        };

                        world
                            .entity_mut(entity)
                            // Add our new `Mesh3d` and `MeshMaterial3d` to our tagged entity
                            .insert((
                                Mesh3d(box_mesh_handle),
                                MeshMaterial3d(box_material_handle),
                                transform,
                            ));
                    });

                    command_queue
                });

                // Add our new task as a component
                commands.entity(entity).insert(ComputeTransform(task));
            }
        }
    }
}

/// This system queries for entities that have the `ComputeTransform` component.
/// It checks if the tasks associated with those entities are complete.
/// If the task is complete, it extracts the result, adds a new [`Mesh3d`] and [`MeshMaterial3d`]
/// to the entity using the result from the task, and removes the task component from the entity.
///
/// **Important Note:**
/// - Don't use `future::block_on(poll_once)` to check if tasks are completed, as it is expensive and
///   can block the main thread. Also, it leaves around a `Task<T>` which will panic if awaited again.
/// - Instead, use `check_ready` for efficient polling, which does not block the main thread.
fn handle_tasks(
    mut commands: Commands,
    mut transform_tasks: Query<(Entity, &mut ComputeTransform)>,
) {
    for (entity, mut task) in &mut transform_tasks {
        // Use `check_ready` to efficiently poll the task without blocking the main thread.
        if let Some(mut commands_queue) = check_ready(&mut task.0) {
            // Append the returned command queue to execute it later.
            commands.append(&mut commands_queue);
            // Task is complete, so remove the task component from the entity.
            commands.entity(entity).remove::<ComputeTransform>();
        }
    }
}

/// This system is only used to setup light and camera for the environment
fn setup_env(mut commands: Commands) {
    // Used to center camera on spawned cubes
    let offset = if NUM_CUBES.is_multiple_of(2) {
        (NUM_CUBES / 2) as f32 - 0.5
    } else {
        (NUM_CUBES / 2) as f32
    };

    // lights
    commands.spawn((PointLight::default(), Transform::from_xyz(4.0, 12.0, 15.0)));

    // camera
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(offset, offset, 15.0)
            .looking_at(Vec3::new(offset, offset, 0.0), Vec3::Y),
    ));
}