Turbopack: migrate server actions to single-graph-traversal

crates/next-api/src/app.rs

@@ -1120,8 +1120,21 @@ impl AppEndpoint {
             };
 
         let server_action_manifest_loader = if process_client_components {
+            let reduced_graphs = get_reduced_graphs_for_endpoint(
+                this.app_project.project(),
+                *rsc_entry,
+                Vc::upcast(this.app_project.client_module_context()),
+            );
+            let actions = reduced_graphs.get_server_actions_for_endpoint(
+                *rsc_entry,
+                match runtime {
+                    NextRuntime::Edge => Vc::upcast(this.app_project.edge_rsc_module_context()),
+                    NextRuntime::NodeJs => Vc::upcast(this.app_project.rsc_module_context()),
+                },
+            );
+
             let server_action_manifest = create_server_actions_manifest(
-                *ResolvedVc::upcast(app_entry.rsc_entry),
+                actions,
                 this.app_project.project().project_path(),
                 node_root,
                 app_entry.original_name.clone(),

crates/next-api/src/module_graph.rs

@@ -1,9 +1,13 @@
-use std::collections::{HashMap, HashSet};
+use std::{
+    borrow::Cow,
+    collections::{HashMap, HashSet},
+};
 
 use anyhow::{Context, Result};
 use next_core::{
     mode::NextMode,
     next_client_reference::{find_server_entries, ServerEntries},
+    next_manifests::ActionLayer,
 };
 use petgraph::{
     graph::{DiGraph, NodeIndex},
@@ -23,6 +27,7 @@ use turbopack_core::{
 use crate::{
     dynamic_imports::{map_next_dynamic, DynamicImports},
     project::Project,
+    server_actions::{map_server_actions, to_rsc_context, AllActions, AllModuleActions},
 };
 
 #[turbo_tasks::value(transparent)]
@@ -267,13 +272,98 @@ impl NextDynamicGraph {
     }
 }
 
+#[turbo_tasks::value]
+pub struct ServerActionsGraph {
+    is_single_page: bool,
+    graph: ResolvedVc<SingleModuleGraph>,
+    /// (Layer, RSC or Browser module) -> list of actions
+    data: ResolvedVc<AllModuleActions>,
+}
+
+#[turbo_tasks::value_impl]
+impl ServerActionsGraph {
+    #[turbo_tasks::function]
+    pub async fn new_with_entries(
+        graph: ResolvedVc<SingleModuleGraph>,
+        is_single_page: bool,
+    ) -> Result<Vc<Self>> {
+        let mapped = map_server_actions(*graph);
+
+        // TODO shrink graph here
+
+        Ok(ServerActionsGraph {
+            is_single_page,
+            graph,
+            data: mapped.to_resolved().await?,
+        }
+        .cell())
+    }
+
+    #[turbo_tasks::function]
+    pub async fn get_server_actions_for_endpoint(
+        &self,
+        entry: ResolvedVc<Box<dyn Module>>,
+        rsc_asset_context: Vc<Box<dyn AssetContext>>,
+    ) -> Result<Vc<AllActions>> {
+        let span = tracing::info_span!("collect server actions for endpoint");
+        async move {
+            let data = &*self.data.await?;
+            let data = if self.is_single_page {
+                // The graph contains the page (= `entry`) only, no need to filter.
+                Cow::Borrowed(data)
+            } else {
+                // The graph contains the whole app, traverse and collect all reachable imports.
+                let graph = &*self.graph.await?;
+
+                let mut result = HashMap::new();
+                graph.traverse_from_entry(entry, |module| {
+                    if let Some(node_data) = data.get(&module) {
+                        result.insert(module, *node_data);
+                    }
+                })?;
+                Cow::Owned(result)
+            };
+
+            let actions = data
+                .iter()
+                .map(|(module, (layer, actions))| async move {
+                    actions
+                        .await?
+                        .iter()
+                        .map(|(hash, name)| async move {
+                            Ok((
+                                hash.to_string(),
+                                (
+                                    *layer,
+                                    name.to_string(),
+                                    if *layer == ActionLayer::Rsc {
+                                        *module
+                                    } else {
+                                        to_rsc_context(**module, rsc_asset_context).await?
+                                    },
+                                ),
+                            ))
+                        })
+                        .try_join()
+                        .await
+                })
+                .try_flat_join()
+                .await?;
+            Ok(Vc::cell(actions.into_iter().collect()))
+        }
+        .instrument(span)
+        .await
+    }
+}
+
 /// The consumers of this shoudln't need to care about the exact contents since it's abstracted away
 /// by the accessor functions, but
 /// - In dev, contains information about the modules of the current endpoint only
 /// - In prod, there is a single `ReducedGraphs` for the whole app, containing all pages
 #[turbo_tasks::value]
 pub struct ReducedGraphs {
     next_dynamic: Vec<ResolvedVc<NextDynamicGraph>>,
+    server_actions: Vec<ResolvedVc<ServerActionsGraph>>,
     // TODO add other graphs
 }
 
@@ -312,6 +402,38 @@ impl ReducedGraphs {
         .instrument(span)
         .await
     }
+
+    /// Returns the server actions for the given page.
+    #[turbo_tasks::function]
+    pub async fn get_server_actions_for_endpoint(
+        &self,
+        entry: Vc<Box<dyn Module>>,
+        rsc_asset_context: Vc<Box<dyn AssetContext>>,
+    ) -> Result<Vc<AllActions>> {
+        let span = tracing::info_span!("collect all server actions for endpoint");
+        async move {
+            if let [graph] = &self.server_actions[..] {
+                // Just a single graph, no need to merge results
+                Ok(graph.get_server_actions_for_endpoint(entry, rsc_asset_context))
+            } else {
+                let result = self
+                    .server_actions
+                    .iter()
+                    .map(|graph| async move {
+                        Ok(graph
+                            .get_server_actions_for_endpoint(entry, rsc_asset_context)
+                            .await?
+                            .clone_value())
+                    })
+                    .try_flat_join()
+                    .await?;
+
+                Ok(Vc::cell(result.into_iter().collect()))
+            }
+        }
+        .instrument(span)
+        .await
+    }
 }
 
 /// Generates a [ReducedGraph] for the given project and endpoint containing information that is
@@ -346,7 +468,7 @@ pub async fn get_reduced_graphs_for_endpoint(
         ),
     };
 
-    let next_dynamic = async move {
+    let next_dynamic = async {
         graphs
             .iter()
             .map(|graph| {
@@ -359,5 +481,21 @@ pub async fn get_reduced_graphs_for_endpoint(
     .instrument(tracing::info_span!("generating next/dynamic graphs"))
     .await?;
 
-    Ok(ReducedGraphs { next_dynamic }.cell())
+    let server_actions = async {
+        graphs
+            .iter()
+            .map(|graph| {
+                ServerActionsGraph::new_with_entries(**graph, is_single_page).to_resolved()
+            })
+            .try_join()
+            .await
+    }
+    .instrument(tracing::info_span!("generating server actions graphs"))
+    .await?;
+
+    Ok(ReducedGraphs {
+        next_dynamic,
+        server_actions,
+    }
+    .cell())
 }