From 0166c4f7fcb074e7e46e2a9de212ef8240b61243 Mon Sep 17 00:00:00 2001
From: James Liu <contact@jamessliu.com>
Date: Fri, 13 May 2022 00:33:13 +0000
Subject: [PATCH] Profile par_for_each(_mut) tasks (#4711)

# Objective
`Query::par_for_each` and it's variants do not show up when profiling using `tracy` or other profilers. Failing to show the impact of changing batch size, the overhead of scheduling tasks, overall thread utilization, etc. other than the effect on the surrounding system.

## Solution
Add a child span that is entered on every spawned task.

Example view of the results in `tracy` using a modified `parallel_query`:
![image](https://user-images.githubusercontent.com/3137680/167560036-626bd091-344b-4664-b323-b692f4f16084.png)

---

## Changelog
Added: `tracing` spans for `Query::par_for_each` and its variants. Spans should now be visible for all
---
 crates/bevy_ecs/src/query/state.rs | 36 +++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
index e01d8df5f2..ae23c1c512 100644
--- a/crates/bevy_ecs/src/query/state.rs
+++ b/crates/bevy_ecs/src/query/state.rs
@@ -11,6 +11,8 @@ use crate::{
     world::{World, WorldId},
 };
 use bevy_tasks::TaskPool;
+#[cfg(feature = "trace")]
+use bevy_utils::tracing::Instrument;
 use fixedbitset::FixedBitSet;
 use std::fmt;
 
@@ -856,7 +858,8 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
                     let mut offset = 0;
                     while offset < table.len() {
                         let func = func.clone();
-                        scope.spawn(async move {
+                        let len = batch_size.min(table.len() - offset);
+                        let task = async move {
                             let mut fetch =
                                 QF::init(world, &self.fetch_state, last_change_tick, change_tick);
                             let mut filter = <QueryFetch<F> as Fetch>::init(
@@ -869,7 +872,6 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
                             let table = &tables[*table_id];
                             fetch.set_table(&self.fetch_state, table);
                             filter.set_table(&self.filter_state, table);
-                            let len = batch_size.min(table.len() - offset);
                             for table_index in offset..offset + len {
                                 if !filter.table_filter_fetch(table_index) {
                                     continue;
@@ -877,7 +879,17 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
                                 let item = fetch.table_fetch(table_index);
                                 func(item);
                             }
-                        });
+                        };
+                        #[cfg(feature = "trace")]
+                        let span = bevy_utils::tracing::info_span!(
+                            "par_for_each",
+                            query = std::any::type_name::<Q>(),
+                            filter = std::any::type_name::<F>(),
+                            count = len,
+                        );
+                        #[cfg(feature = "trace")]
+                        let task = task.instrument(span);
+                        scope.spawn(task);
                         offset += batch_size;
                     }
                 }
@@ -888,7 +900,8 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
                     let archetype = &archetypes[*archetype_id];
                     while offset < archetype.len() {
                         let func = func.clone();
-                        scope.spawn(async move {
+                        let len = batch_size.min(archetype.len() - offset);
+                        let task = async move {
                             let mut fetch =
                                 QF::init(world, &self.fetch_state, last_change_tick, change_tick);
                             let mut filter = <QueryFetch<F> as Fetch>::init(
@@ -902,14 +915,25 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
                             fetch.set_archetype(&self.fetch_state, archetype, tables);
                             filter.set_archetype(&self.filter_state, archetype, tables);
 
-                            let len = batch_size.min(archetype.len() - offset);
                             for archetype_index in offset..offset + len {
                                 if !filter.archetype_filter_fetch(archetype_index) {
                                     continue;
                                 }
                                 func(fetch.archetype_fetch(archetype_index));
                             }
-                        });
+                        };
+
+                        #[cfg(feature = "trace")]
+                        let span = bevy_utils::tracing::info_span!(
+                            "par_for_each",
+                            query = std::any::type_name::<Q>(),
+                            filter = std::any::type_name::<F>(),
+                            count = len,
+                        );
+                        #[cfg(feature = "trace")]
+                        let task = task.instrument(span);
+
+                        scope.spawn(task);
                         offset += batch_size;
                     }
                 }