async_compute/
async_compute.rs

1//! This example shows how to use the ECS and the [`AsyncComputeTaskPool`]
2//! to spawn, poll, and complete tasks across systems and system ticks.
3
4use bevy::{
5    ecs::{system::SystemState, world::CommandQueue},
6    prelude::*,
7    tasks::{block_on, futures_lite::future, AsyncComputeTaskPool, Task},
8};
9use rand::Rng;
10use std::time::Duration;
11
12fn main() {
13    App::new()
14        .add_plugins(DefaultPlugins)
15        .add_systems(Startup, (setup_env, add_assets, spawn_tasks))
16        .add_systems(Update, handle_tasks)
17        .run();
18}
19
20// Number of cubes to spawn across the x, y, and z axis
21const NUM_CUBES: u32 = 6;
22
23#[derive(Resource, Deref)]
24struct BoxMeshHandle(Handle<Mesh>);
25
26#[derive(Resource, Deref)]
27struct BoxMaterialHandle(Handle<StandardMaterial>);
28
29/// Startup system which runs only once and generates our Box Mesh
30/// and Box Material assets, adds them to their respective Asset
31/// Resources, and stores their handles as resources so we can access
32/// them later when we're ready to render our Boxes
33fn add_assets(
34    mut commands: Commands,
35    mut meshes: ResMut<Assets<Mesh>>,
36    mut materials: ResMut<Assets<StandardMaterial>>,
37) {
38    let box_mesh_handle = meshes.add(Cuboid::new(0.25, 0.25, 0.25));
39    commands.insert_resource(BoxMeshHandle(box_mesh_handle));
40
41    let box_material_handle = materials.add(Color::srgb(1.0, 0.2, 0.3));
42    commands.insert_resource(BoxMaterialHandle(box_material_handle));
43}
44
45#[derive(Component)]
46struct ComputeTransform(Task<CommandQueue>);
47
48/// This system generates tasks simulating computationally intensive
49/// work that potentially spans multiple frames/ticks. A separate
50/// system, [`handle_tasks`], will poll the spawned tasks on subsequent
51/// frames/ticks, and use the results to spawn cubes
52fn spawn_tasks(mut commands: Commands) {
53    let thread_pool = AsyncComputeTaskPool::get();
54    for x in 0..NUM_CUBES {
55        for y in 0..NUM_CUBES {
56            for z in 0..NUM_CUBES {
57                // Spawn new task on the AsyncComputeTaskPool; the task will be
58                // executed in the background, and the Task future returned by
59                // spawn() can be used to poll for the result
60                let entity = commands.spawn_empty().id();
61                let task = thread_pool.spawn(async move {
62                    let duration = Duration::from_secs_f32(rand::thread_rng().gen_range(0.05..5.0));
63
64                    // Pretend this is a time-intensive function. :)
65                    async_std::task::sleep(duration).await;
66
67                    // Such hard work, all done!
68                    let transform = Transform::from_xyz(x as f32, y as f32, z as f32);
69                    let mut command_queue = CommandQueue::default();
70
71                    // we use a raw command queue to pass a FnOnce(&mut World) back to be
72                    // applied in a deferred manner.
73                    command_queue.push(move |world: &mut World| {
74                        let (box_mesh_handle, box_material_handle) = {
75                            let mut system_state = SystemState::<(
76                                Res<BoxMeshHandle>,
77                                Res<BoxMaterialHandle>,
78                            )>::new(world);
79                            let (box_mesh_handle, box_material_handle) =
80                                system_state.get_mut(world);
81
82                            (box_mesh_handle.clone(), box_material_handle.clone())
83                        };
84
85                        world
86                            .entity_mut(entity)
87                            // Add our new `Mesh3d` and `MeshMaterial3d` to our tagged entity
88                            .insert((
89                                Mesh3d(box_mesh_handle),
90                                MeshMaterial3d(box_material_handle),
91                                transform,
92                            ))
93                            // Task is complete, so remove task component from entity
94                            .remove::<ComputeTransform>();
95                    });
96
97                    command_queue
98                });
99
100                // Spawn new entity and add our new task as a component
101                commands.entity(entity).insert(ComputeTransform(task));
102            }
103        }
104    }
105}
106
107/// This system queries for entities that have our Task<Transform> component. It polls the
108/// tasks to see if they're complete. If the task is complete it takes the result, adds a
109/// new [`Mesh3d`] and [`MeshMaterial3d`] to the entity using the result from the task's work, and
110/// removes the task component from the entity.
111fn handle_tasks(mut commands: Commands, mut transform_tasks: Query<&mut ComputeTransform>) {
112    for mut task in &mut transform_tasks {
113        if let Some(mut commands_queue) = block_on(future::poll_once(&mut task.0)) {
114            // append the returned command queue to have it execute later
115            commands.append(&mut commands_queue);
116        }
117    }
118}
119
120/// This system is only used to setup light and camera for the environment
121fn setup_env(mut commands: Commands) {
122    // Used to center camera on spawned cubes
123    let offset = if NUM_CUBES % 2 == 0 {
124        (NUM_CUBES / 2) as f32 - 0.5
125    } else {
126        (NUM_CUBES / 2) as f32
127    };
128
129    // lights
130    commands.spawn((PointLight::default(), Transform::from_xyz(4.0, 12.0, 15.0)));
131
132    // camera
133    commands.spawn((
134        Camera3d::default(),
135        Transform::from_xyz(offset, offset, 15.0)
136            .looking_at(Vec3::new(offset, offset, 0.0), Vec3::Y),
137    ));
138}
async_compute/async_compute.rs

async_compute/
async_compute.rs
