Add support for Instrumentation and Middleware in global module ID op…

…timization
vercel · Aug 22, 2024 · 5c76d20 · 5c76d20
1 parent fea6aee
commit 5c76d20
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 31 deletions.
diff --git a/crates/next-api/src/global_module_id_strategy.rs b/crates/next-api/src/global_module_id_strategy.rs
@@ -32,6 +32,19 @@ impl GlobalModuleIdStrategyBuilder {
         preprocessed_module_ids.push(preprocess_module_ids(entrypoints.pages_app_endpoint));
         preprocessed_module_ids.push(preprocess_module_ids(entrypoints.pages_document_endpoint));
 
+        if let Some(middleware) = &entrypoints.middleware {
+            dbg!("middleware");
+            preprocessed_module_ids.push(preprocess_module_ids(middleware.endpoint));
+        }
+
+        if let Some(instrumentation) = &entrypoints.instrumentation {
+            dbg!("instrumentation");
+            let node_js = instrumentation.node_js;
+            let edge = instrumentation.edge;
+            preprocessed_module_ids.push(preprocess_module_ids(node_js));
+            preprocessed_module_ids.push(preprocess_module_ids(edge));
+        }
+
         for (_, route) in entrypoints.routes.iter() {
             match route {
                 Route::Page {

diff --git a/crates/next-api/src/instrumentation.rs b/crates/next-api/src/instrumentation.rs
@@ -65,7 +65,7 @@ impl InstrumentationEndpoint {
     }
 
     #[turbo_tasks::function]
-    async fn edge_files(&self) -> Result<Vc<OutputAssets>> {
+    async fn core_modules(&self) -> Result<Vc<InstrumentationCoreModules>> {
         let userland_module = self
             .asset_context
             .process(
@@ -74,22 +74,35 @@ impl InstrumentationEndpoint {
             )
             .module();
 
-        let module = wrap_edge_entry(
+        let edge_entry_module = wrap_edge_entry(
             self.asset_context,
             self.project.project_path(),
             userland_module,
             "instrumentation".into(),
         );
 
+        Ok(InstrumentationCoreModules {
+            userland_module,
+            edge_entry_module,
+        }
+        .cell())
+    }
+
+    #[turbo_tasks::function]
+    async fn edge_files(self: Vc<Self>) -> Result<Vc<OutputAssets>> {
+        let this = self.await?;
+
+        let module = self.core_modules().await?.edge_entry_module;
+
         let mut evaluatable_assets = get_server_runtime_entries(
             Value::new(ServerContextType::Instrumentation {
-                app_dir: self.app_dir,
-                ecmascript_client_reference_transition_name: self
+                app_dir: this.app_dir,
+                ecmascript_client_reference_transition_name: this
                     .ecmascript_client_reference_transition_name,
             }),
-            self.project.next_mode(),
+            this.project.next_mode(),
         )
-        .resolve_entries(self.asset_context)
+        .resolve_entries(this.asset_context)
         .await?
         .clone_value();
 
@@ -104,7 +117,7 @@ impl InstrumentationEndpoint {
         };
         evaluatable_assets.push(evaluatable);
 
-        let edge_chunking_context = self.project.edge_chunking_context(false);
+        let edge_chunking_context = this.project.edge_chunking_context(false);
 
         let edge_files = edge_chunking_context.evaluated_chunk_group_assets(
             module.ident(),
@@ -116,36 +129,32 @@ impl InstrumentationEndpoint {
     }
 
     #[turbo_tasks::function]
-    async fn node_chunk(&self) -> Result<Vc<Box<dyn OutputAsset>>> {
-        let chunking_context = self.project.server_chunking_context(false);
+    async fn node_chunk(self: Vc<Self>) -> Result<Vc<Box<dyn OutputAsset>>> {
+        let this = self.await?;
 
-        let userland_module = self
-            .asset_context
-            .process(
-                self.source,
-                Value::new(ReferenceType::Entry(EntryReferenceSubType::Instrumentation)),
-            )
-            .module();
+        let chunking_context = this.project.server_chunking_context(false);
+
+        let userland_module = self.core_modules().await?.userland_module;
 
         let Some(module) = Vc::try_resolve_downcast(userland_module).await? else {
             bail!("Entry module must be evaluatable");
         };
 
         let EntryChunkGroupResult { asset: chunk, .. } = *chunking_context
             .entry_chunk_group(
-                self.project
+                this.project
                     .node_root()
                     .join("server/instrumentation.js".into()),
                 module,
                 get_server_runtime_entries(
                     Value::new(ServerContextType::Instrumentation {
-                        app_dir: self.app_dir,
-                        ecmascript_client_reference_transition_name: self
+                        app_dir: this.app_dir,
+                        ecmascript_client_reference_transition_name: this
                             .ecmascript_client_reference_transition_name,
                     }),
-                    self.project.next_mode(),
+                    this.project.next_mode(),
                 )
-                .resolve_entries(self.asset_context),
+                .resolve_entries(this.asset_context),
                 Value::new(AvailabilityInfo::Root),
             )
             .await?;
@@ -199,6 +208,12 @@ impl InstrumentationEndpoint {
     }
 }
 
+#[turbo_tasks::value]
+struct InstrumentationCoreModules {
+    pub userland_module: Vc<Box<dyn Module>>,
+    pub edge_entry_module: Vc<Box<dyn Module>>,
+}
+
 #[turbo_tasks::value_impl]
 impl Endpoint for InstrumentationEndpoint {
     #[turbo_tasks::function]
@@ -238,7 +253,11 @@ impl Endpoint for InstrumentationEndpoint {
     }
 
     #[turbo_tasks::function]
-    fn root_modules(self: Vc<Self>) -> Result<Vc<Modules>> {
-        Err(anyhow::anyhow!("Not implemented."))
+    async fn root_modules(self: Vc<Self>) -> Result<Vc<Modules>> {
+        let core_modules = self.core_modules().await?;
+        Ok(Vc::cell(vec![
+            core_modules.userland_module,
+            core_modules.edge_entry_module,
+        ]))
     }
 }
diff --git a/crates/next-api/src/middleware.rs b/crates/next-api/src/middleware.rs
@@ -121,13 +121,7 @@ impl MiddlewareEndpoint {
     async fn output_assets(self: Vc<Self>) -> Result<Vc<OutputAssets>> {
         let this = self.await?;
 
-        let userland_module = this
-            .asset_context
-            .process(
-                this.source,
-                Value::new(ReferenceType::Entry(EntryReferenceSubType::Middleware)),
-            )
-            .module();
+        let userland_module = self.userland_module();
 
         let config = parse_config_from_source(userland_module);
 
@@ -192,6 +186,19 @@ impl MiddlewareEndpoint {
 
         Ok(Vc::cell(output_assets))
     }
+
+    #[turbo_tasks::function]
+    async fn userland_module(self: Vc<Self>) -> Result<Vc<Box<dyn Module>>> {
+        let this = self.await?;
+
+        Ok(this
+            .asset_context
+            .process(
+                this.source,
+                Value::new(ReferenceType::Entry(EntryReferenceSubType::Middleware)),
+            )
+            .module())
+    }
 }
 
 #[turbo_tasks::value_impl]
@@ -240,6 +247,6 @@ impl Endpoint for MiddlewareEndpoint {
 
     #[turbo_tasks::function]
     fn root_modules(self: Vc<Self>) -> Result<Vc<Modules>> {
-        Err(anyhow::anyhow!("Not implemented."))
+        Ok(Vc::cell(vec![self.userland_module()]))
     }
 }