From b237938f80dfaec081cdfaeb28e76b8e12dbbb7e Mon Sep 17 00:00:00 2001
From: James Liu <contact@jamessliu.com>
Date: Sat, 30 Sep 2023 01:03:35 -0700
Subject: [PATCH] Cache parallel iteration spans (#9950)

# Objective
We cached system spans in #9390, but another common span seen in most
Bevy apps when enabling tracing are Query::par_iter(_mut) related spans.

## Solution
Cache them in QueryState. The one downside to this is that we pay for
the memory for every Query(State) instantiated, not just those that are
used for parallel iteration, but this shouldn't be a significant cost
unless the app is creating hundreds of thousands of Query(State)s
regularly.

## Metrics
Tested against `cargo run --profile stress-test --features trace_tracy
--example many_cubes`. Yellow is this PR, red is main.

`sync_simple_transforms`:


![image](https://github.com/bevyengine/bevy/assets/3137680/d60f6d69-5586-4424-9d78-aac78992aacd)

`check_visibility`:


![image](https://github.com/bevyengine/bevy/assets/3137680/096a58d2-a330-4a32-b806-09cd524e6e15)

Full frame:


![image](https://github.com/bevyengine/bevy/assets/3137680/3b088cf8-9487-4bc7-a308-026e172d6672)
---
 crates/bevy_ecs/src/query/state.rs | 45 ++++++++++++------------------
 1 file changed, 18 insertions(+), 27 deletions(-)

diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
index e64c257ec04b5f..01e218b230b6ed 100644
--- a/crates/bevy_ecs/src/query/state.rs
+++ b/crates/bevy_ecs/src/query/state.rs
@@ -12,7 +12,7 @@ use crate::{
     world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
 };
 #[cfg(feature = "trace")]
-use bevy_utils::tracing::Instrument;
+use bevy_utils::tracing::Span;
 use fixedbitset::FixedBitSet;
 use std::{any::TypeId, borrow::Borrow, fmt, mem::MaybeUninit};
 
@@ -39,6 +39,8 @@ pub struct QueryState<Q: WorldQuery, F: ReadOnlyWorldQuery = ()> {
     pub(crate) matched_archetype_ids: Vec<ArchetypeId>,
     pub(crate) fetch_state: Q::State,
     pub(crate) filter_state: F::State,
+    #[cfg(feature = "trace")]
+    par_iter_span: Span,
 }
 
 impl<Q: WorldQuery, F: ReadOnlyWorldQuery> std::fmt::Debug for QueryState<Q, F> {
@@ -125,6 +127,12 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
             matched_tables: Default::default(),
             matched_archetypes: Default::default(),
             archetype_component_access: Default::default(),
+            #[cfg(feature = "trace")]
+            par_iter_span: bevy_utils::tracing::info_span!(
+                "par_for_each",
+                query = std::any::type_name::<Q>(),
+                filter = std::any::type_name::<F>(),
+            ),
         };
         state.update_archetypes(world);
         state
@@ -1205,7 +1213,9 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                     while offset < table.entity_count() {
                         let func = func.clone();
                         let len = batch_size.min(table.entity_count() - offset);
-                        let task = async move {
+                        scope.spawn(async move {
+                            #[cfg(feature = "trace")]
+                            let _span = self.par_iter_span.enter();
                             let mut fetch =
                                 Q::init_fetch(world, &self.fetch_state, last_run, this_run);
                             let mut filter =
@@ -1223,17 +1233,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                                 }
                                 func(Q::fetch(&mut fetch, *entity, row));
                             }
-                        };
-                        #[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;
                     }
                 }
@@ -1249,7 +1249,9 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                     while offset < archetype.len() {
                         let func = func.clone();
                         let len = batch_size.min(archetype.len() - offset);
-                        let task = async move {
+                        scope.spawn(async move {
+                            #[cfg(feature = "trace")]
+                            let _span = self.par_iter_span.enter();
                             let mut fetch =
                                 Q::init_fetch(world, &self.fetch_state, last_run, this_run);
                             let mut filter =
@@ -1277,19 +1279,8 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
                                     archetype_entity.table_row(),
                                 ));
                             }
-                        };
-
-                        #[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;
                     }
                 }