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

async_compute/
async_compute.rs
