fix: React.cache to work correctly if AsyncLocalStorage is available (#…

…590) * should bundle * add a failing example and fix with a hack * prefer top-level await * revert top-level await
dai-shi · Mar 10, 2024 · 8ff7829 · 8ff7829
1 parent 87295fd
commit 8ff7829
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 8 deletions.
diff --git a/examples/08_cookies/src/components/App.tsx b/examples/08_cookies/src/components/App.tsx
@@ -1,10 +1,17 @@
-import { Suspense } from 'react';
+import { Suspense, cache } from 'react';
 import { getContext } from 'waku/server';
 
 import { Counter } from './Counter.js';
 
+const cachedFn = cache(() => Date.now());
+
 const InternalAsyncComponent = async () => {
+  const val1 = cachedFn();
   await new Promise((resolve) => setTimeout(resolve, 1000));
+  const val2 = cachedFn();
+  if (val1 !== val2) {
+    throw new Error('Cache not working');
+  }
   console.log(getContext());
   return null;
 };

diff --git a/packages/waku/src/lib/builder/build.ts b/packages/waku/src/lib/builder/build.ts
@@ -119,6 +119,7 @@ const analyzeEntries = async (
     build: {
       write: false,
       ssr: true,
+      target: 'node18',
       rollupOptions: {
         onwarn,
         input: {
@@ -242,6 +243,7 @@ const buildServerBundle = async (
     build: {
       ssr: true,
       ssrEmitAssets: true,
+      target: 'node18',
       outDir: joinPath(rootDir, config.distDir),
       rollupOptions: {
         onwarn,
@@ -299,6 +301,7 @@ const buildSsrBundle = async (
     publicDir: false,
     build: {
       ssr: true,
+      target: 'node18',
       outDir: joinPath(rootDir, config.distDir, config.ssrDir),
       rollupOptions: {
         onwarn,

diff --git a/packages/waku/src/lib/plugins/vite-plugin-rsc-entries.ts b/packages/waku/src/lib/plugins/vite-plugin-rsc-entries.ts
@@ -9,7 +9,12 @@ export function rscEntriesPlugin(opts: {
   entriesFile: string;
   moduleMap: Record<string, string>;
 }): Plugin {
-  let codeToAdd = `
+  const codeToPrepend = `
+try {
+  globalThis.AsyncLocalStorage = (await import('node:async_hooks')).AsyncLocalStorage;
+} catch (e) {}
+`;
+  let codeToAppend = `
 export function loadModule(id) {
   switch (id) {
     ${Object.entries(opts.moduleMap)
@@ -23,19 +28,25 @@ export function loadModule(id) {
     const file = normalizePath(
       path.relative(path.dirname(opts.entriesFile), path.resolve(CONFIG_FILE)),
     );
-    codeToAdd += `
+    codeToAppend += `
 export const loadConfig = async () => (await import('${file}')).default;
 `;
   } else {
-    codeToAdd += `
+    codeToAppend += `
 export const loadConfig = async () => ({});
 `;
   }
   return {
     name: 'rsc-entries-plugin',
     transform(code, id) {
+      if (
+        // FIXME this is too hacky and not the right place to patch
+        id.endsWith('/react-server-dom-webpack-server.edge.production.min.js')
+      ) {
+        return codeToPrepend + code;
+      }
       if (id === opts.entriesFile) {
-        return code + codeToAdd;
+        return code + codeToAppend;
       }
     },
   };

diff --git a/packages/waku/src/lib/renderers/dev-worker-impl.ts b/packages/waku/src/lib/renderers/dev-worker-impl.ts
@@ -3,6 +3,7 @@
 import { pathToFileURL } from 'node:url';
 import { parentPort, getEnvironmentData } from 'node:worker_threads';
 import { Server } from 'node:http';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import type { TransferListItem } from 'node:worker_threads';
 import { createServer as createViteServer } from 'vite';
 import viteReact from '@vitejs/plugin-react';
@@ -25,6 +26,9 @@ import { rscDelegatePlugin } from '../plugins/vite-plugin-rsc-delegate.js';
 import { mergeUserViteConfig } from '../utils/merge-vite-config.js';
 import { viteHot } from '../plugins/vite-plugin-rsc-hmr.js';
 
+// For react-server-dom-webpack/server.edge
+(globalThis as any).AsyncLocalStorage = AsyncLocalStorage;
+
 const { default: module } = await import('node:module');
 const HAS_MODULE_REGISTER = typeof module.register === 'function';
 if (HAS_MODULE_REGISTER) {

diff --git a/packages/waku/src/server.ts b/packages/waku/src/server.ts
@@ -75,11 +75,18 @@ type RenderStore<
   context: RscContext;
 };
 
-const DO_NOT_BUNDLE = '';
-
 let renderStorage: AsyncLocalStorageType<RenderStore> | undefined;
 
-import(/* @vite-ignore */ DO_NOT_BUNDLE + 'node:async_hooks')
+// TODO top-level await doesn't work. Let's revisit after supporting "use server"
+// try {
+//   const { AsyncLocalStorage } = await import('node:async_hooks');
+//   renderStorage = new AsyncLocalStorage();
+// } catch (e) {
+//   console.warn(
+//     'AsyncLocalStorage is not available, rerender and getContext are only available in sync.',
+//   );
+// }
+import('node:async_hooks')
   .then(({ AsyncLocalStorage }) => {
     renderStorage = new AsyncLocalStorage();
   })