async_compute.rs

//! 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),
    ));
}