Repository: vercel/next.js
Pull request: [turbopack] Avoid calling `find_server_entries` in the whole_app_module_graph case

## Improve performance of find_client_references in production builds

### What?
Avoid calling `find_server_entries` in the production case where there is a single module graph.  While all the turbotasks launched by `find_server_entries` are by definition already done it is still a lot of tiny tasks to get cache hits and re do work that was already done when building the graph.

So instead we can just report server components and server utilities when performing the graph traversal.  Unfortunately this does not work in development builds since the stacked layout graphs interfere with ordering, so in that case we will still call `find_server_entries`, however we can do this concurrently with the graph traversal to save a bit of latency.

###  Why?

To improve performance of course!

![image.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/AwJ29EfoPcPdLSwCZxAz/89a13c00-f4a6-4fc1-8626-619ce645600d.png)

This represents a small latency progression. Collecting client references is on the critical path for chunking so this makes sense. We have to collect all client references for all pages 


Closes PACK-5169

Changed files:
filename: crates/next-api/src/module_graph.rs
patch: @@ -2,19 +2,21 @@ use std::{borrow::Cow, collections::hash_map::Entry};
 
 use anyhow::{Ok, Result};
 use either::Either;
+use futures::join;
 use next_core::{
     next_client_reference::{
         ClientReference, ClientReferenceGraphResult, ClientReferenceType, ServerEntries,
         find_server_entries,
     },
     next_dynamic::NextDynamicEntryModule,
     next_manifests::ActionLayer,
+    next_server_utility::server_utility_module::NextServerUtilityModule,
 };
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashMap;
 use tracing::Instrument;
 use turbo_rcstr::{RcStr, rcstr};
 use turbo_tasks::{
-    CollectiblesSource, FxIndexMap, ReadRef, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt,
+    CollectiblesSource, FxIndexMap, FxIndexSet, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt,
     ValueToString, Vc,
 };
 use turbo_tasks_fs::FileSystemPath;
@@ -289,7 +291,8 @@ impl ClientReferencesGraph {
             // post_order callbacks which is the same as evaluation order
             let mut client_references = Vec::new();
             let mut client_reference_modules = Vec::new();
-            let mut server_components = FxHashSet::default();
+            let mut server_components = FxIndexSet::default();
+            let mut server_utils = FxIndexSet::default();
 
             // Track how we reached each client reference.  This way if a client reference is
             // referenced by the root and by a server component we don't only associate it with the
@@ -355,6 +358,12 @@ impl ClientReferencesGraph {
                 },
                 |_, node, state_map| {
                     let module = node.module();
+                    if let Some(server_util_module) =
+                        ResolvedVc::try_downcast_type::<NextServerUtilityModule>(module)
+                    {
+                        server_utils.insert(server_util_module);
+                    }
+
                     let Some(module_type) = data.manifest.get(&module) else {
                         return Ok(());
                     };
@@ -407,8 +416,9 @@ impl ClientReferencesGraph {
 
             Ok(ClientReferenceGraphResult {
                 client_references: client_references.into_iter().collect(),
-                server_utils: vec![],
-                server_component_entries: vec![],
+                // The order of server_utils does not matter
+                server_utils: server_utils.into_iter().collect(),
+                server_component_entries: server_components.into_iter().collect(),
             }
             .cell())
         }
@@ -717,46 +727,68 @@ impl GlobalBuildInformation {
     ) -> Result<Vc<ClientReferenceGraphResult>> {
         let span = tracing::info_span!("collect all client references for endpoint");
         async move {
-            let mut result = if let [graph] = &self.client_references[..] {
-                // Just a single graph, no need to merge results
-                graph
-                    .get_client_references_for_endpoint(entry)
-                    .owned()
-                    .await?
+            let result = if let [graph] = &self.client_references[..] {
+                // Just a single graph, no need to merge results  This also naturally aggregates
+                // server components and server utilities in the correct order
+                let result = graph.get_client_references_for_endpoint(entry);
+                #[cfg(debug_assertions)]
+                {
+                    let result = result.await?;
+                    if has_layout_segments {
+                        use rustc_hash::FxHashSet;
+
+                        let ServerEntries {
+                            server_utils,
+                            server_component_entries,
+                        } = &*find_server_entries(entry, include_traced).await?;
+                        // order of server utils doesn't matter, so just ensure that they match
+                        assert_eq!(
+                            FxHashSet::from_iter(result.server_utils.iter()),
+                            FxHashSet::from_iter(server_utils.iter())
+                        );
+                        // The order of server_components does matter, enforce it is identical
+                        assert_eq!(&result.server_component_entries, server_component_entries);
+                    }
+                }
+                result
             } else {
                 let results = self
                     .client_references
                     .iter()
                     .map(|graph| graph.get_client_references_for_endpoint(entry))
-                    .try_join()
-                    .await?;
-
-                let mut iter = results.into_iter();
-                let mut result = ReadRef::into_owned(iter.next().unwrap());
-                for r in iter {
-                    result.extend(&r);
-                }
-                result
-            };
+                    .try_join();
+                // Do this separately for now, because the aggregation of multiple graph traversals
+                // messes up the order of the server_component_entries.
+                let server_entries = async {
+                    if has_layout_segments {
+                        let server_entries = find_server_entries(entry, include_traced).await?;
+                        Ok(Some(server_entries))
+                    } else {
+                        Ok(None)
+                    }
+                };
+                // Wait for both in parallel since `find_server_entries` tends to be slower than the
+                // graph traversals
+                let (results, server_entries) = join!(results, server_entries);
 
-            // TODO(luke.sandberg): at least in the whole_app_module_graph case we should be able to
-            // collect server components and server utilities during the above traversals in the
-            // correct order.  `find_server_entries returns them in reverse topological order (root
-            // layout first, page last) but the above traversals find them in DFS post
-            // order which means we would need to reverse it.
-            // For server_utils the order is irrelevant.
-            if has_layout_segments {
-                // Do this separately for now, because the graph traversal order messes up the order
-                // of the server_component_entries.
-                let ServerEntries {
+                let mut result = ClientReferenceGraphResult {
+                    client_references: results?
+                        .iter()
+                        .flat_map(|r| r.client_references.iter().copied())
+                        .collect(),
+                    ..Default::default()
+                };
+                if let Some(ServerEntries {
                     server_utils,
                     server_component_entries,
-                } = &*find_server_entries(entry, include_traced).await?;
-                result.server_utils = server_utils.clone();
-                result.server_component_entries = server_component_entries.clone();
-            }
-
-            Ok(result.cell())
+                }) = server_entries?.as_deref()
+                {
+                    result.server_utils = server_utils.clone();
+                    result.server_component_entries = server_component_entries.clone();
+                }
+                result.cell()
+            };
+            Ok(result)
         }
         .instrument(span)
         .await

filename: crates/next-core/src/next_client_reference/visit_client_reference.rs
patch: @@ -106,17 +106,6 @@ impl ClientReferenceGraphResult {
     }
 }
 
-impl ClientReferenceGraphResult {
-    /// Merges multiple return values of client_reference_graph together.
-    pub fn extend(&mut self, other: &Self) {
-        self.client_references
-            .extend(other.client_references.iter().copied());
-        self.server_component_entries
-            .extend(other.server_component_entries.iter().copied());
-        self.server_utils.extend(other.server_utils.iter().copied());
-    }
-}
-
 #[turbo_tasks::value(shared)]
 #[derive(Clone, Debug)]
 pub struct ServerEntries {