improve aggregation tree tracing, remove verbose tracing (#8148)

### Description

Make the graph aggregation tracing more useful and less verbose

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->

crates/turbo-tasks-memory/src/aggregation/aggregation_data.rs

@@ -3,7 +3,7 @@ use std::ops::{Deref, DerefMut};
 use super::{
     increase_aggregation_number_internal, AggregationContext, AggregationNode, AggregationNodeGuard,
 };
-use crate::aggregation::balance_queue::BalanceQueue;
+use crate::aggregation::{balance_queue::BalanceQueue, increase::IncreaseReason};
 
 /// Gives an reference to the aggregated data for a given item. This will
 /// convert the item to a fully aggregated node.
@@ -26,6 +26,7 @@ where
             node_id,
             u32::MAX,
             u32::MAX,
+            IncreaseReason::AggregationData,
         );
         balance_queue.process(ctx);
         let guard = ctx.node(node_id);
@@ -45,6 +46,7 @@ pub fn prepare_aggregation_data<C: AggregationContext>(ctx: &C, node_id: &C::Nod
         node_id,
         u32::MAX,
         u32::MAX,
+        IncreaseReason::AggregationData,
     );
     balance_queue.process(ctx);
 }

crates/turbo-tasks-memory/src/aggregation/balance_edge.rs

@@ -7,6 +7,7 @@ use super::{
         RemovePositveFollowerCountResult,
     },
     in_progress::is_in_progress,
+    increase::IncreaseReason,
     increase_aggregation_number_internal,
     uppers::{
         add_upper_count, remove_positive_upper_count, remove_upper_count,
@@ -62,6 +63,7 @@ pub(super) fn balance_edge<C: AggregationContext>(
                             target_id,
                             target_aggregation_number + 1,
                             target_aggregation_number + 1,
+                            IncreaseReason::EqualAggregationNumberOnBalance,
                         );
                     }
                 }

crates/turbo-tasks-memory/src/aggregation/increase.rs

@@ -6,6 +6,20 @@ use super::{
 };
 pub(super) const LEAF_NUMBER: u32 = 64;
 
+#[derive(Debug)]
+pub enum IncreaseReason {
+    Upgraded,
+    AggregationData,
+    EqualAggregationNumberOnBalance,
+    EqualAggregationNumberOnNewFollower,
+    OptimizeForUppers,
+    OptimizeForFollowers,
+    LeafEdge,
+    LeafEdgeAfterIncrease,
+    #[cfg(test)]
+    Test,
+}
+
 impl<I: Clone + Eq + Hash, D> AggregationNode<I, D> {
     /// Increase the aggregation number of a node. This might temporarily
     /// violate the graph invariants between uppers and followers of that node.
@@ -21,6 +35,7 @@ impl<I: Clone + Eq + Hash, D> AggregationNode<I, D> {
         node_id: &C::NodeRef,
         min_aggregation_number: u32,
         target_aggregation_number: u32,
+        reason: IncreaseReason,
     ) -> Option<PreparedInternalIncreaseAggregationNumber<C>> {
         if self.aggregation_number() >= min_aggregation_number {
             return None;
@@ -30,6 +45,7 @@ impl<I: Clone + Eq + Hash, D> AggregationNode<I, D> {
             uppers: self.uppers_mut().iter().cloned().collect(),
             min_aggregation_number,
             target_aggregation_number,
+            reason,
         })
     }
 
@@ -47,6 +63,7 @@ impl<I: Clone + Eq + Hash, D> AggregationNode<I, D> {
             node_id,
             new_aggregation_number,
             new_aggregation_number,
+            IncreaseReason::Test,
         )
         .map(PreparedIncreaseAggregationNumber)
     }
@@ -64,10 +81,17 @@ pub(super) fn increase_aggregation_number_immediately<C: AggregationContext>(
     node_id: C::NodeRef,
     min_aggregation_number: u32,
     target_aggregation_number: u32,
+    reason: IncreaseReason,
 ) -> Option<PreparedInternalIncreaseAggregationNumber<C>> {
     if node.aggregation_number() >= min_aggregation_number {
         return None;
     }
+
+    let _span = tracing::trace_span!(
+        "increase_aggregation_number_immediately",
+        reason = debug(&reason)
+    )
+    .entered();
     let children = matches!(**node, AggregationNode::Leaf { .. })
         .then(|| node.children().collect::<StackVec<_>>());
     match &mut **node {
@@ -128,6 +152,7 @@ pub enum PreparedInternalIncreaseAggregationNumber<C: AggregationContext> {
         uppers: StackVec<C::NodeRef>,
         min_aggregation_number: u32,
         target_aggregation_number: u32,
+        reason: IncreaseReason,
     },
     Leaf {
         children: StackVec<C::NodeRef>,
@@ -152,6 +177,7 @@ impl<C: AggregationContext> PreparedInternalOperation<C>
                 mut target_aggregation_number,
                 node_id,
                 uppers,
+                reason,
             } => {
                 let mut need_to_run = true;
                 while need_to_run {
@@ -178,6 +204,9 @@ impl<C: AggregationContext> PreparedInternalOperation<C>
                 if node.aggregation_number() >= min_aggregation_number {
                     return;
                 }
+                let _span =
+                    tracing::trace_span!("increase_aggregation_number", reason = debug(&reason))
+                        .entered();
                 let children = matches!(*node, AggregationNode::Leaf { .. })
                     .then(|| node.children().collect::<StackVec<_>>());
                 let (uppers, followers) = match &mut *node {
@@ -197,6 +226,7 @@ impl<C: AggregationContext> PreparedInternalOperation<C>
                                     &child_id,
                                     target_aggregation_number + 1,
                                     target_aggregation_number + 1,
+                                    IncreaseReason::LeafEdgeAfterIncrease,
                                 );
                             }
                             return;
@@ -253,6 +283,7 @@ impl<C: AggregationContext> PreparedInternalOperation<C>
                         &child_id,
                         target_aggregation_number + 1,
                         target_aggregation_number + 1,
+                        IncreaseReason::LeafEdgeAfterIncrease,
                     );
                 }
             }
@@ -285,12 +316,14 @@ pub fn increase_aggregation_number_internal<C: AggregationContext>(
     node_id: &C::NodeRef,
     min_aggregation_number: u32,
     target_aggregation_number: u32,
+    reason: IncreaseReason,
 ) {
     let prepared = node.increase_aggregation_number_internal(
         ctx,
         node_id,
         min_aggregation_number,
         target_aggregation_number,
+        reason,
     );
     drop(node);
     prepared.apply(ctx, balance_queue);

crates/turbo-tasks-memory/src/aggregation/new_edge.rs

@@ -2,8 +2,8 @@ use super::{
     balance_queue::BalanceQueue,
     in_progress::start_in_progress_all,
     increase::{
-        increase_aggregation_number_immediately, PreparedInternalIncreaseAggregationNumber,
-        LEAF_NUMBER,
+        increase_aggregation_number_immediately, IncreaseReason,
+        PreparedInternalIncreaseAggregationNumber, LEAF_NUMBER,
     },
     increase_aggregation_number_internal, notify_new_follower,
     notify_new_follower::PreparedNotifyNewFollower,
@@ -19,7 +19,6 @@ const MAX_AFFECTED_NODES: usize = 4096;
 
 /// Handle the addition of a new edge to a node. The the edge is propagated to
 /// the uppers of that node or added a inner node.
-#[tracing::instrument(level = tracing::Level::TRACE, name = "handle_new_edge_preparation", skip_all)]
 pub fn handle_new_edge<C: AggregationContext>(
     ctx: &C,
     origin: &mut C::Guard<'_>,
@@ -43,6 +42,7 @@ pub fn handle_new_edge<C: AggregationContext>(
                     origin_id.clone(),
                     LEAF_NUMBER,
                     LEAF_NUMBER,
+                    IncreaseReason::Upgraded,
                 )
                 .unwrap();
                 Some(PreparedNewEdge::Upgraded {
@@ -93,7 +93,6 @@ enum PreparedNewEdge<C: AggregationContext> {
 
 impl<C: AggregationContext> PreparedOperation<C> for PreparedNewEdge<C> {
     type Result = ();
-    #[tracing::instrument(level = tracing::Level::TRACE, name = "handle_new_edge", skip_all)]
     fn apply(self, ctx: &C) {
         let mut balance_queue = BalanceQueue::new();
         match self {
@@ -103,20 +102,15 @@ impl<C: AggregationContext> PreparedOperation<C> for PreparedNewEdge<C> {
                 uppers,
                 target_id,
             } => {
-                let _span = tracing::trace_span!("leaf").entered();
-                {
-                    let _span =
-                        tracing::trace_span!("increase_aggregation_number_internal").entered();
-                    // TODO add to prepared
-                    increase_aggregation_number_internal(
-                        ctx,
-                        &mut balance_queue,
-                        ctx.node(&target_id),
-                        &target_id,
-                        min_aggregation_number,
-                        target_aggregation_number,
-                    );
-                }
+                increase_aggregation_number_internal(
+                    ctx,
+                    &mut balance_queue,
+                    ctx.node(&target_id),
+                    &target_id,
+                    min_aggregation_number,
+                    target_aggregation_number,
+                    IncreaseReason::LeafEdge,
+                );
                 let mut affected_nodes = 0;
                 for upper_id in uppers {
                     affected_nodes += notify_new_follower(
@@ -158,7 +152,6 @@ impl<C: AggregationContext> PreparedOperation<C> for PreparedNewEdge<C> {
                 }
             }
         }
-        let _span = tracing::trace_span!("balance_queue").entered();
         balance_queue.process(ctx);
     }
 }
@@ -171,6 +164,7 @@ fn handle_expensive_node<C: AggregationContext>(
     balance_queue: &mut BalanceQueue<C::NodeRef>,
     node_id: &C::NodeRef,
 ) {
+    let _span = tracing::trace_span!("handle_expensive_node").entered();
     let node = ctx.node(node_id);
     let uppers = node.uppers().iter().cloned().collect::<StackVec<_>>();
     let leaf = matches!(*node, AggregationNode::Leaf { .. });

crates/turbo-tasks-memory/src/aggregation/notify_new_follower.rs

@@ -4,6 +4,7 @@ use super::{
     balance_queue::BalanceQueue,
     followers::add_follower,
     in_progress::{finish_in_progress_without_node, start_in_progress},
+    increase::IncreaseReason,
     increase_aggregation_number_internal,
     optimize::optimize_aggregation_number_for_uppers,
     uppers::add_upper,
@@ -195,6 +196,7 @@ impl<C: AggregationContext> PreparedInternalOperation<C> for PreparedNotifyNewFo
                                         &follower_id,
                                         upper_aggregation_number + 1,
                                         upper_aggregation_number + 1,
+                                        IncreaseReason::EqualAggregationNumberOnNewFollower,
                                     );
                                     // retry
                                 } else {

crates/turbo-tasks-memory/src/aggregation/optimize.rs

@@ -1,6 +1,8 @@
+use tracing::Level;
+
 use super::{
     balance_queue::BalanceQueue,
-    increase::{increase_aggregation_number_internal, LEAF_NUMBER},
+    increase::{increase_aggregation_number_internal, IncreaseReason, LEAF_NUMBER},
     AggregationContext, StackVec,
 };
 
@@ -12,6 +14,7 @@ pub const MAX_FOLLOWERS: usize = 128;
 /// The goal is to reduce the number of upper nodes, so we try to find a
 /// aggregation number that is higher than some of the upper nodes.
 /// Returns true if the aggregation number was increased.
+#[tracing::instrument(level = Level::TRACE, skip(ctx, balance_queue, node_id, uppers))]
 pub fn optimize_aggregation_number_for_uppers<C: AggregationContext>(
     ctx: &C,
     balance_queue: &mut BalanceQueue<C::NodeRef>,
@@ -55,6 +58,7 @@ pub fn optimize_aggregation_number_for_uppers<C: AggregationContext>(
             node_id,
             aggregation_number,
             aggregation_number,
+            IncreaseReason::OptimizeForUppers,
         );
         return true;
     } else {
@@ -69,6 +73,7 @@ pub fn optimize_aggregation_number_for_uppers<C: AggregationContext>(
                     node_id,
                     aggregation_number,
                     aggregation_number,
+                    IncreaseReason::OptimizeForUppers,
                 );
                 return true;
             }
@@ -81,6 +86,7 @@ pub fn optimize_aggregation_number_for_uppers<C: AggregationContext>(
 /// The goal is to reduce the number of followers, so we try to find a
 /// aggregation number that is higher than some of the followers.
 /// Returns true if the aggregation number was increased.
+#[tracing::instrument(level = Level::TRACE, skip(ctx, balance_queue, node_id, followers))]
 pub fn optimize_aggregation_number_for_followers<C: AggregationContext>(
     ctx: &C,
     balance_queue: &mut BalanceQueue<C::NodeRef>,
@@ -131,6 +137,7 @@ pub fn optimize_aggregation_number_for_followers<C: AggregationContext>(
                 node_id,
                 aggregation_number,
                 aggregation_number,
+                IncreaseReason::OptimizeForFollowers,
             );
             return true;
         }