fix(node): Improve chunk flushing (#1985)

ScriptedAlchemy · web-flow · commit 73eb07e93c6e · 2024-01-19T20:38:17.000-08:00
diff --git a/.changeset/popular-spoons-grin.md b/.changeset/popular-spoons-grin.md
@@ -0,0 +1,6 @@
+---
+'@module-federation/nextjs-mf': patch
+'@module-federation/node': patch
+---
+
+Rewrite chunk flushing and hot reloading to use federation runtime apis
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
@@ -66,4 +66,4 @@ jobs:
         run: lsof -ti tcp:3005,3006,3007 | xargs kill
 
       - name: Build Next.js Apps in Production Mode
-        run: pnpm app:next:prod
+        run: pnpm app:next:build
diff --git a/apps/3000-home/project.json b/apps/3000-home/project.json
@@ -8,11 +8,11 @@
       "executor": "@nx/next:build",
       "defaultConfiguration": "production",
       "options": {
-        "outputPath": "dist/apps/3000-home"
+        "outputPath": "apps/3000-home"
       },
       "configurations": {
         "development": {
-          "outputPath": "dist/apps/3000-home"
+          "outputPath": "apps/3000-home"
         },
         "production": {}
       },
diff --git a/apps/3001-shop/pages/_document.js b/apps/3001-shop/pages/_document.js
@@ -8,14 +8,20 @@ import {
 
 class MyDocument extends Document {
   static async getInitialProps(ctx) {
+    await revalidate().then((shouldUpdate) => {
+      if (shouldUpdate) {
+        ctx.res.writeHead(307, { Location: ctx.req.url });
+        ctx.res.end();
+      }
+    });
     const initialProps = await Document.getInitialProps(ctx);
     const chunks = await flushChunks();
     ctx?.res?.on('finish', () => {
-      revalidate().then((shouldUpdate) => {
-        if (shouldUpdate) {
-          console.log('should HMR', shouldUpdate);
-        }
-      });
+      // revalidate().then((shouldUpdate) => {
+      //   if (shouldUpdate) {
+      //     console.log('should HMR', shouldUpdate);
+      //   }
+      // });
     });
 
     return {
diff --git a/apps/3001-shop/project.json b/apps/3001-shop/project.json
@@ -8,11 +8,11 @@
       "executor": "@nx/next:build",
       "defaultConfiguration": "production",
       "options": {
-        "outputPath": "apps/3001-shop/dist"
+        "outputPath": "apps/3001-shop"
       },
       "configurations": {
         "development": {
-          "outputPath": "apps/3001-shop/dist"
+          "outputPath": "apps/3001-shop"
         },
         "production": {}
       },
diff --git a/apps/3002-checkout/project.json b/apps/3002-checkout/project.json
@@ -8,11 +8,11 @@
       "executor": "@nx/next:build",
       "defaultConfiguration": "production",
       "options": {
-        "outputPath": "{options.outputPath}"
+        "outputPath": "apps/3002-checkout"
       },
       "configurations": {
         "development": {
-          "outputPath": "apps/3002-checkout/dist"
+          "outputPath": "apps/3002-checkout"
         },
         "production": {}
       },
diff --git a/package.json b/package.json
@@ -31,7 +31,8 @@
     "extract-i18n:website": "nx run website:extract-i18n",
     "sync:pullMFTypes": "concurrently \"node ./packages/enhanced/pullts.js\"",
     "app:next:dev": "nx run-many --target=serve --configuration=development -p 3000-home,3001-shop,3002-checkout",
-    "app:next:prod": "nx run-many --target=build --configuration=production -p 3000-home,3001-shop,3002-checkout",
+    "app:next:build": "nx run-many --target=build --configuration=production -p 3000-home,3001-shop,3002-checkout",
+    "app:next:prod": "nx run-many --target=serve --configuration=production -p 3000-home,3001-shop,3002-checkout",
     "app:node:dev": "nx run-many --target=serve --configuration=development -p node-host,node-local-remote,node-remote",
     "app:runtime:dev": "nx run-many --target=serve -p 3005-runtime-host,3006-runtime-remote,3007-runtime-remote",
     "commitlint": "commitlint --edit",
diff --git a/packages/nextjs-mf/src/federation-noop.ts b/packages/nextjs-mf/src/federation-noop.ts
@@ -9,5 +9,5 @@ require('next/amp');
 require('styled-jsx');
 require('styled-jsx/style');
 require('next/image');
-require('react/jsx-dev-runtime');
+// require('react/jsx-dev-runtime');
 require('react/jsx-runtime');
diff --git a/packages/nextjs-mf/src/plugins/NextFederationPlugin/index.ts b/packages/nextjs-mf/src/plugins/NextFederationPlugin/index.ts
@@ -67,6 +67,16 @@ export class NextFederationPlugin {
     // ContainerPlugin will get NextFederationPlugin._options, so NextFederationPlugin._options should be the same as normalFederationPluginOptions
     this._options = normalFederationPluginOptions;
     new ModuleFederationPlugin(normalFederationPluginOptions).apply(compiler);
+
+    const runtimeESMPath = require.resolve(
+      '@module-federation/runtime/dist/index.esm.js',
+    );
+    compiler.hooks.afterPlugins.tap('PatchAliasWebpackPlugin', () => {
+      compiler.options.resolve.alias = {
+        ...compiler.options.resolve.alias,
+        '@module-federation/runtime$': runtimeESMPath,
+      };
+    });
   }
 
   private validateOptions(compiler: Compiler): boolean {
diff --git a/packages/nextjs-mf/src/plugins/container/runtimePlugin.ts b/packages/nextjs-mf/src/plugins/container/runtimePlugin.ts
@@ -32,7 +32,7 @@ export default function (): FederationRuntimePlugin {
     beforeInit(args) {
       const { userOptions, shareInfo } = args;
       const { shared } = userOptions;
-
+      if (!globalThis.usedChunks) globalThis.usedChunks = new Set();
       if (shared) {
         Object.keys(shared || {}).forEach((sharedKey) => {
           if (!shared[sharedKey].strategy) {
@@ -50,7 +50,7 @@ export default function (): FederationRuntimePlugin {
 
       // if (__webpack_runtime_id__ && !__webpack_runtime_id__.startsWith('webpack')) return args;
       const { moduleCache, name } = args.origin;
-      const gs = (globalThis as any) || new Function('return globalThis')();
+      const gs = new Function('return globalThis')();
       const attachedRemote = gs[name];
       if (attachedRemote) {
         moduleCache.set(name, attachedRemote);
@@ -70,9 +70,35 @@ export default function (): FederationRuntimePlugin {
     afterResolve(args) {
       return args;
     },
-    // onLoad(args) {
-    //   return args;
-    // },
+    onLoad(args) {
+      const { exposeModuleFactory, exposeModule, id } = args;
+
+      const moduleOrFactory = exposeModuleFactory || exposeModule;
+      const exposedModuleExports = moduleOrFactory();
+      const handler = {
+        //@ts-ignore
+        get: function (target, prop, receiver) {
+          const origMethod = target[prop];
+          if (typeof origMethod === 'function') {
+            //@ts-ignore
+            return function (...args) {
+              globalThis.usedChunks.add(
+                //@ts-ignore
+                id,
+              );
+
+              // console.log(`function as called to ${prop}`, id);
+              //@ts-ignore
+              return origMethod.apply(this, args);
+            };
+          } else {
+            return Reflect.get(target, prop, receiver);
+          }
+        },
+      };
+
+      return () => new Proxy(exposedModuleExports, handler);
+    },
     resolveShare(args) {
       if (
         args.pkgName !== 'react' &&
diff --git a/packages/nextjs-mf/utils/index.ts b/packages/nextjs-mf/utils/index.ts
@@ -33,13 +33,16 @@ export type { FlushedChunksProps } from './flushedChunks';
  * If the function is called on the server side, it imports the revalidate function from the module federation node utilities and returns the result of calling that function.
  * @returns {Promise<boolean>} A promise that resolves with a boolean.
  */
-export const revalidate = () => {
+export const revalidate = (
+  fetchModule: any = undefined,
+  force: boolean = false,
+) => {
   if (typeof window !== 'undefined') {
     console.error('revalidate should only be called server-side');
     return Promise.resolve(false);
   }
   // @ts-ignore
   return import('@module-federation/node/utils').then((utils) => {
-    return utils.revalidate();
+    return utils.revalidate(fetchModule, force);
   });
 };
diff --git a/packages/node/global.d.ts b/packages/node/global.d.ts
@@ -33,6 +33,7 @@ declare global {
       };
     }
   }
+  var usedChunks: Set<string>;
 
   var __FEDERATION__: {
     __INSTANCES__: Array<{
diff --git a/packages/node/src/utils/flush-chunks.ts b/packages/node/src/utils/flush-chunks.ts
@@ -91,104 +91,80 @@ const createShareMap = () => {
  */
 // @ts-ignore
 const processChunk = async (chunk, shareMap, hostStats) => {
+  const chunks = new Set();
+  const [remote, req] = chunk.split('/');
+  const request = './' + req;
+  const knownRemotes = getAllKnownRemotes();
+  //@ts-ignore
+  if (!knownRemotes[remote]) {
+    console.error(
+      `flush chunks: Remote ${remote} is not defined in the global config`,
+    );
+    return;
+  }
+
   try {
-    // Create a set to store the chunks
-    const chunks = new Set();
+    //@ts-ignore
+    const remoteName = new URL(knownRemotes[remote].entry).pathname
+      .split('/')
+      .pop();
+    //@ts-ignore
 
-    // Split the chunk string into remote and request
-    const [remote, request] = chunk.split('->');
-    const knownRemotes = getAllKnownRemotes();
+    const statsFile = knownRemotes[remote].entry
+      .replace(remoteName, 'federated-stats.json')
+      .replace('ssr', 'chunks');
+    let stats = {};
 
-    // If the remote is not defined in the global config, return
-    //@ts-ignore
-    if (!knownRemotes[remote]) {
-      console.error(
-        `flush chunks:`,
-        `Remote ${remote} is not defined in the global config`,
-      );
-      return;
+    try {
+      stats = await fetch(statsFile).then((res) => res.json());
+    } catch (e) {
+      console.error('flush error', e);
     }
+    //@ts-ignore
 
-    try {
-      // Extract the remote name from the URL
-      //@ts-ignore
-      const remoteName = new URL(
-        //@ts-ignore
-        globalThis.__remote_scope__._config[remote],
-      ).pathname
-        .split('/')
-        .pop();
+    const [prefix] = knownRemotes[remote].entry.split('static/');
+    //@ts-ignore
 
-      // Construct the stats file URL from the remote config
+    if (stats.federatedModules) {
       //@ts-ignore
-      const statsFile = globalThis.__remote_scope__._config[remote]
-        .replace(remoteName, 'federated-stats.json')
-        .replace('ssr', 'chunks');
-
-      let stats = {};
-      try {
-        // Fetch the remote config and stats file
-        stats = await fetch(statsFile).then((res) => res.json());
-      } catch (e) {
-        console.error('flush error', e);
-      }
 
-      // Add the main chunk to the chunks set
-      //TODO: ensure host doesnt embed its own remote in ssr, this causes crash
-      // chunks.add(
-      //   global.__remote_scope__._config[remote].replace('ssr', 'chunks')
-      // );
+      stats.federatedModules.forEach((modules) => {
+        if (modules.exposes?.[request]) {
+          //@ts-ignore
 
-      // Extract the prefix from the remote config
-      const [prefix] =
-        //@ts-ignore
-        globalThis.__remote_scope__._config[remote].split('static/');
+          modules.exposes[request].forEach((chunk) => {
+            chunks.add([prefix, chunk].join(''));
 
-      // Process federated modules from the stats object
-      // @ts-ignore
-      if (stats.federatedModules) {
-        // @ts-ignore
-        stats.federatedModules.forEach((modules) => {
-          // Process exposed modules
-          if (modules.exposes?.[request]) {
-            // @ts-ignore
-            modules.exposes[request].forEach((chunk) => {
-              chunks.add([prefix, chunk].join(''));
-
-              //TODO: reimplement this
-              Object.values(chunk).forEach((chunk) => {
-                // Add files to the chunks set
-                // @ts-ignore
-                if (chunk.files) {
-                  // @ts-ignore
-                  chunk.files.forEach((file) => {
-                    chunks.add(prefix + file);
-                  });
-                }
-                // Process required modules
-                // @ts-ignore
-                if (chunk.requiredModules) {
-                  // @ts-ignore
-                  chunk.requiredModules.forEach((module) => {
-                    // Check if the module is in the shareMap
-                    if (shareMap[module]) {
-                      // If the module is from the host, log the host stats
-                    }
-                  });
-                }
-              });
-            });
-          }
-        });
-      }
+            Object.values(chunk).forEach((chunk) => {
+              //@ts-ignore
 
-      // Return the array of chunks
-      return Array.from(chunks);
-    } catch (e) {
-      console.error('flush error:', e);
+              if (chunk.files) {
+                //@ts-ignore
+
+                chunk.files.forEach((file) => {
+                  chunks.add(prefix + file);
+                });
+              }
+              //@ts-ignore
+
+              if (chunk.requiredModules) {
+                //@ts-ignore
+
+                chunk.requiredModules.forEach((module) => {
+                  if (shareMap[module]) {
+                    // If the module is from the host, log the host stats
+                  }
+                });
+              }
+            });
+          });
+        }
+      });
     }
+
+    return Array.from(chunks);
   } catch (e) {
-    // catch just in case
+    console.error('flush error:', e);
   }
 };
 
diff --git a/packages/node/src/utils/hot-reload.ts b/packages/node/src/utils/hot-reload.ts
