docs/bevy/examples/async_tasks/async_compute.rs 6.7 K · 161 lines · raw · history

1 //! This example demonstrates how to use Bevy's ECS and the [`AsyncComputeTaskPool`]
2 //! to offload computationally intensive tasks to a background thread pool, process them
3 //! asynchronously, and apply the results across systems and ticks.
4 //!
5 //! Unlike the channel-based approach (where tasks send results directly via a communication
6 //! channel), this example uses the `AsyncComputeTaskPool` to run tasks in the background,
7 //! check for their completion, and handle results when the task is ready. This method allows
8 //! tasks to be processed in parallel without blocking the main thread, but requires periodically
9 //! checking the status of each task.
10 //!
11 //! The channel-based approach, on the other hand, detaches tasks and communicates results
12 //! through a channel, avoiding the need to check task statuses manually.
13
14 use bevy::{
15 ecs::{system::SystemState, world::CommandQueue},
16 prelude::*,
17 tasks::{futures::check_ready, AsyncComputeTaskPool, Task},
18 };
19 use futures_timer::Delay;
20 use rand::RngExt;
21 use std::time::Duration;
22
23 fn main() {
24 App::new()
25 .add_plugins(DefaultPlugins)
26 .add_systems(Startup, (setup_env, add_assets, spawn_tasks))
27 .add_systems(Update, handle_tasks)
28 .run();
29 }
30
31 // Number of cubes to spawn across the x, y, and z axis
32 const NUM_CUBES: u32 = 6;
33
34 #[derive(Resource, Deref)]
35 struct BoxMeshHandle(Handle<Mesh>);
36
37 #[derive(Resource, Deref)]
38 struct BoxMaterialHandle(Handle<StandardMaterial>);
39
40 /// Startup system which runs only once and generates our Box Mesh
41 /// and Box Material assets, adds them to their respective Asset
42 /// Resources, and stores their handles as resources so we can access
43 /// them later when we're ready to render our Boxes
44 fn add_assets(
45 mut commands: Commands,
46 mut meshes: ResMut<Assets<Mesh>>,
47 mut materials: ResMut<Assets<StandardMaterial>>,
48 ) {
49 let box_mesh_handle = meshes.add(Cuboid::new(0.25, 0.25, 0.25));
50 commands.insert_resource(BoxMeshHandle(box_mesh_handle));
51
52 let box_material_handle = materials.add(Color::srgb(1.0, 0.2, 0.3));
53 commands.insert_resource(BoxMaterialHandle(box_material_handle));
54 }
55
56 #[derive(Component)]
57 struct ComputeTransform(Task<CommandQueue>);
58
59 /// This system generates tasks simulating computationally intensive
60 /// work that potentially spans multiple frames/ticks. A separate
61 /// system, [`handle_tasks`], will track the spawned tasks on subsequent
62 /// frames/ticks, and use the results to spawn cubes.
63 ///
64 /// The task is offloaded to the `AsyncComputeTaskPool`, allowing heavy computation
65 /// to be handled asynchronously, without blocking the main game thread.
66 fn spawn_tasks(mut commands: Commands) {
67 let thread_pool = AsyncComputeTaskPool::get();
68 for x in 0..NUM_CUBES {
69 for y in 0..NUM_CUBES {
70 for z in 0..NUM_CUBES {
71 // Spawn new task on the AsyncComputeTaskPool; the task will be
72 // executed in the background, and the Task future returned by
73 // spawn() can be used to poll for the result
74 let entity = commands.spawn_empty().id();
75 let task = thread_pool.spawn(async move {
76 let duration = Duration::from_secs_f32(rand::rng().random_range(0.05..5.0));
77
78 // Pretend this is a time-intensive function. :)
79 Delay::new(duration).await;
80
81 // Such hard work, all done!
82 let transform = Transform::from_xyz(x as f32, y as f32, z as f32);
83 let mut command_queue = CommandQueue::default();
84
85 // we use a raw command queue to pass a FnOnce(&mut World) back to be
86 // applied in a deferred manner.
87 command_queue.push(move |world: &mut World| {
88 let (box_mesh_handle, box_material_handle) = {
89 let mut system_state = SystemState::<(
90 Res<BoxMeshHandle>,
91 Res<BoxMaterialHandle>,
92 )>::new(world);
93 let (box_mesh_handle, box_material_handle) =
94 system_state.get_mut(world).unwrap();
95
96 (box_mesh_handle.clone(), box_material_handle.clone())
97 };
98
99 world
100 .entity_mut(entity)
101 // Add our new `Mesh3d` and `MeshMaterial3d` to our tagged entity
102 .insert((
103 Mesh3d(box_mesh_handle),
104 MeshMaterial3d(box_material_handle),
105 transform,
106 ));
107 });
108
109 command_queue
110 });
111
112 // Add our new task as a component
113 commands.entity(entity).insert(ComputeTransform(task));
114 }
115 }
116 }
117 }
118
119 /// This system queries for entities that have the `ComputeTransform` component.
120 /// It checks if the tasks associated with those entities are complete.
121 /// If the task is complete, it extracts the result, adds a new [`Mesh3d`] and [`MeshMaterial3d`]
122 /// to the entity using the result from the task, and removes the task component from the entity.
123 ///
124 /// **Important Note:**
125 /// - Don't use `future::block_on(poll_once)` to check if tasks are completed, as it is expensive and
126 /// can block the main thread. Also, it leaves around a `Task<T>` which will panic if awaited again.
127 /// - Instead, use `check_ready` for efficient polling, which does not block the main thread.
128 fn handle_tasks(
129 mut commands: Commands,
130 mut transform_tasks: Query<(Entity, &mut ComputeTransform)>,
131 ) {
132 for (entity, mut task) in &mut transform_tasks {
133 // Use `check_ready` to efficiently poll the task without blocking the main thread.
134 if let Some(mut commands_queue) = check_ready(&mut task.0) {
135 // Append the returned command queue to execute it later.
136 commands.append(&mut commands_queue);
137 // Task is complete, so remove the task component from the entity.
138 commands.entity(entity).remove::<ComputeTransform>();
139 }
140 }
141 }
142
143 /// This system is only used to setup light and camera for the environment
144 fn setup_env(mut commands: Commands) {
145 // Used to center camera on spawned cubes
146 let offset = if NUM_CUBES.is_multiple_of(2) {
147 (NUM_CUBES / 2) as f32 - 0.5
148 } else {
149 (NUM_CUBES / 2) as f32
150 };
151
152 // lights
153 commands.spawn((PointLight::default(), Transform::from_xyz(4.0, 12.0, 15.0)));
154
155 // camera
156 commands.spawn((
157 Camera3d::default(),
158 Transform::from_xyz(offset, offset, 15.0)
159 .looking_at(Vec3::new(offset, offset, 0.0), Vec3::Y),
160 ));
161 }