WIP

unstubbable · Jun 27, 2024 · df07778 · df07778
1 parent c340be2
commit df07778
Show file tree

Hide file tree

Showing 14 changed files with 201 additions and 146 deletions.
diff --git a/apps/aws-app/src/handler/index.tsx b/apps/aws-app/src/handler/index.tsx
@@ -1,4 +1,4 @@
-import {routerLocationAsyncLocalStorage} from '@mfng/core/router-location-async-local-storage';
+import {requestContextAsyncLocalStorage} from '@mfng/core/request-context-async-local-storage';
 import {
   createRscActionStream,
   createRscAppStream,
@@ -38,35 +38,38 @@ async function renderApp(
 ): Promise<Response> {
   const {pathname, search} = new URL(request.url);
 
-  return routerLocationAsyncLocalStorage.run({pathname, search}, async () => {
-    const rscAppStream = createRscAppStream(<App />, {
-      reactClientManifest,
-      mainCssHref: cssManifest[`main.css`]!,
-      formState,
-    });
+  return requestContextAsyncLocalStorage.run(
+    {routerLocation: {pathname, search}},
+    async () => {
+      const rscAppStream = createRscAppStream(<App />, {
+        reactClientManifest,
+        mainCssHref: cssManifest[`main.css`]!,
+        formState,
+      });
+
+      if (request.headers.get(`accept`) === `text/x-component`) {
+        return new Response(rscAppStream, {
+          headers: {
+            'Content-Type': `text/x-component; charset=utf-8`,
+            'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
+          },
+        });
+      }
+
+      const htmlStream = await createHtmlStream(rscAppStream, {
+        reactSsrManifest,
+        bootstrapScripts: [jsManifest[`main.js`]!],
+        formState,
+      });
 
-    if (request.headers.get(`accept`) === `text/x-component`) {
-      return new Response(rscAppStream, {
+      return new Response(htmlStream, {
         headers: {
-          'Content-Type': `text/x-component; charset=utf-8`,
+          'Content-Type': `text/html; charset=utf-8`,
           'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
         },
       });
-    }
-
-    const htmlStream = await createHtmlStream(rscAppStream, {
-      reactSsrManifest,
-      bootstrapScripts: [jsManifest[`main.js`]!],
-      formState,
-    });
-
-    return new Response(htmlStream, {
-      headers: {
-        'Content-Type': `text/html; charset=utf-8`,
-        'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
-      },
-    });
-  });
+    },
+  );
 }
 
 async function handleGet(request: Request): Promise<Response> {

diff --git a/apps/cloudflare-app/src/worker/index.tsx b/apps/cloudflare-app/src/worker/index.tsx
@@ -1,18 +1,18 @@
 import type {Request} from '@cloudflare/workers-types';
-import {routerLocationAsyncLocalStorage} from '@mfng/core/router-location-async-local-storage';
+import {requestContextAsyncLocalStorage} from '@mfng/core/request-context-async-local-storage';
 import {
   createRscActionStream,
   createRscAppStream,
   createRscFormState,
 } from '@mfng/core/server/rsc';
 import {
-  type PartialPrerenderResult,
-  createHtmlStream,
-  partiallyPrerender,
+  type PrerenderResult,
+  prerender,
   resumePartialPrerender,
 } from '@mfng/core/server/ssr';
+import type {RouterLocation} from '@mfng/core/use-router-location';
 import * as React from 'react';
-import type {PostponedState, ReactFormState} from 'react-dom/server';
+import type {ReactFormState} from 'react-dom/server';
 import {App} from './app.js';
 import {
   cssManifest,
@@ -23,8 +23,8 @@ import {
 } from './manifests.js';
 
 interface PrerenderCache {
-  get(url: string): PartialPrerenderResult | undefined;
-  set(url: string, result: PartialPrerenderResult): void;
+  get(url: string): PrerenderResult | undefined;
+  set(url: string, result: PrerenderResult): void;
 }
 
 // TODO: Use a persistent cache; honor cache control headers.
@@ -69,69 +69,71 @@ async function renderApp(
   formState?: ReactFormState,
 ): Promise<Response> {
   const {pathname, search} = new URL(request.url);
-
-  // TODO: Refactor to requestAsyncLocalStorage that contains location and
-  // prerender/postpone status.
-  return routerLocationAsyncLocalStorage.run({pathname, search}, async () => {
-    const rscStream = createRscAppStream(<App />, {
-      reactClientManifest,
-      mainCssHref: cssManifest[`main.css`]!,
-      formState,
-    });
-
-    const [prerenderRscStream, finalRscStream] = rscStream.tee();
-
-    // TODO: Client-side navigation request should not get prerender stream
-    // (with postponed elements).
-    if (request.headers.get(`accept`) === `text/x-component`) {
-      return new Response(prerenderRscStream, {
-        headers: {'Content-Type': `text/x-component; charset=utf-8`},
-      });
-    }
-
-    let cachedResult = prerenderCache.get(request.url);
-
-    if (!cachedResult) {
-      cachedResult = await partiallyPrerender(prerenderRscStream, {
-        reactSsrManifest,
-        bootstrapScripts: [jsManifest[`main.js`]!],
-        formState,
-      });
-
-      prerenderCache.set(request.url, cachedResult);
-    }
-
-    const postponedState: PostponedState | null = JSON.parse(
-      cachedResult.postponed,
+  const routerLocation: RouterLocation = {pathname, search};
+
+  if (request.headers.get(`accept`) === `text/x-component`) {
+    return requestContextAsyncLocalStorage.run(
+      {routerLocation, isPrerender: false},
+      () =>
+        new Response(
+          createRscAppStream(<App />, {
+            reactClientManifest,
+            mainCssHref: cssManifest[`main.css`]!,
+            formState,
+          }),
+          {headers: {'Content-Type': `text/x-component; charset=utf-8`}},
+        ),
     );
+  }
 
-    if (postponedState) {
-      const resumedRscStream = createRscAppStream(<App />, {
-        reactClientManifest,
-        mainCssHref: cssManifest[`main.css`]!,
-        formState,
-      });
-
-      const resumedHtmlStream = await resumePartialPrerender(resumedRscStream, {
-        postponedState: JSON.parse(cachedResult.postponed),
-        reactSsrManifest,
-      });
-
-      return new Response(
-        prependString(resumedHtmlStream, cachedResult.prelude),
-        {headers: {'Content-Type': `text/html; charset=utf-8`}},
-      );
-    }
+  const prerenderResult =
+    prerenderCache.get(request.url) ??
+    (await requestContextAsyncLocalStorage.run(
+      {routerLocation, isPrerender: true},
+      async () => {
+        const prerenderRscStream = createRscAppStream(<App />, {
+          reactClientManifest,
+          mainCssHref: cssManifest[`main.css`]!,
+          formState,
+        });
+
+        const result = await prerender(prerenderRscStream, {
+          reactSsrManifest,
+          bootstrapScripts: [jsManifest[`main.js`]!],
+          formState,
+        });
+
+        prerenderCache.set(request.url, result);
+
+        return result;
+      },
+    ));
+
+  if (prerenderResult.didPostpone) {
+    const resumedHtmlStream = await requestContextAsyncLocalStorage.run(
+      {routerLocation, isPrerender: false},
+      async () => {
+        const resumedRscStream = createRscAppStream(<App />, {
+          reactClientManifest,
+          mainCssHref: cssManifest[`main.css`]!,
+          formState,
+        });
+
+        return resumePartialPrerender(resumedRscStream, {
+          postponedState: prerenderResult.postponedState,
+          reactSsrManifest,
+        });
+      },
+    );
 
-    const htmlStream = await createHtmlStream(finalRscStream, {
-      reactSsrManifest,
-      bootstrapScripts: [jsManifest[`main.js`]!],
-      formState,
-    });
+    return new Response(
+      prependString(resumedHtmlStream, prerenderResult.prelude),
+      {headers: {'Content-Type': `text/html; charset=utf-8`}},
+    );
+  }
 
-    return new Response(htmlStream, {
-      headers: {'Content-Type': `text/html; charset=utf-8`},
-    });
+  return new Response(prerenderResult.html, {
+    headers: {'Content-Type': `text/html; charset=utf-8`},
   });
 }
 

diff --git a/apps/shared-app/src/server/postponed.tsx b/apps/shared-app/src/server/postponed.tsx
@@ -1,3 +1,4 @@
+import {usePostpone} from '@mfng/core/server/rsc';
 import * as React from 'react';
 import 'server-only';
 import {Markdown} from './markdown.js';
@@ -9,12 +10,8 @@ async function fetchContent(): Promise<string> {
   return `This is a postponed server component.`;
 }
 
-let count = 0;
-
 export async function Postponed(): Promise<React.ReactElement> {
-  if (count++ % 2 === 0) {
-    React.unstable_postpone();
-  }
+  usePostpone();
 
   const content = await fetchContent();
 

diff --git a/apps/vercel-app/src/edge-function-handler/index.tsx b/apps/vercel-app/src/edge-function-handler/index.tsx
@@ -1,4 +1,4 @@
-import {routerLocationAsyncLocalStorage} from '@mfng/core/router-location-async-local-storage';
+import {requestContextAsyncLocalStorage} from '@mfng/core/request-context-async-local-storage';
 import {
   createRscActionStream,
   createRscAppStream,
@@ -44,35 +44,38 @@ async function renderApp(
 ): Promise<Response> {
   const {pathname, search} = new URL(request.url);
 
-  return routerLocationAsyncLocalStorage.run({pathname, search}, async () => {
-    const rscAppStream = createRscAppStream(<App />, {
-      reactClientManifest,
-      mainCssHref: cssManifest[`main.css`]!,
-      formState,
-    });
+  return requestContextAsyncLocalStorage.run(
+    {routerLocation: {pathname, search}},
+    async () => {
+      const rscAppStream = createRscAppStream(<App />, {
+        reactClientManifest,
+        mainCssHref: cssManifest[`main.css`]!,
+        formState,
+      });
+
+      if (request.headers.get(`accept`) === `text/x-component`) {
+        return new Response(rscAppStream, {
+          headers: {
+            'Content-Type': `text/x-component; charset=utf-8`,
+            'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
+          },
+        });
+      }
+
+      const htmlStream = await createHtmlStream(rscAppStream, {
+        reactSsrManifest,
+        bootstrapScripts: [jsManifest[`main.js`]!],
+        formState,
+      });
 
-    if (request.headers.get(`accept`) === `text/x-component`) {
-      return new Response(rscAppStream, {
+      return new Response(htmlStream, {
         headers: {
-          'Content-Type': `text/x-component; charset=utf-8`,
+          'Content-Type': `text/html; charset=utf-8`,
           'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
         },
       });
-    }
-
-    const htmlStream = await createHtmlStream(rscAppStream, {
-      reactSsrManifest,
-      bootstrapScripts: [jsManifest[`main.js`]!],
-      formState,
-    });
-
-    return new Response(htmlStream, {
-      headers: {
-        'Content-Type': `text/html; charset=utf-8`,
-        'Cache-Control': `s-maxage=60, stale-while-revalidate=${oneDay}`,
-      },
-    });
-  });
+    },
+  );
 }
 
 async function handleGet(request: Request): Promise<Response> {

diff --git a/packages/core/README.md b/packages/core/README.md
@@ -40,14 +40,15 @@ npm install @mfng/core react@canary react-dom@canary react-server-dom-webpack@ca
 - `createRscAppStream`
 - `createRscActionStream`
 - `createRscFormState`
+- `usePostpone`
 
 #### `@mfng/core/server/ssr`
 
 - `createHtmlStream`
 
-#### `@mfng/core/router-location-async-local-storage`
+#### `@mfng/core/request-context-async-local-storage`
 
-- `routerLocationAsyncLocalStorage`
+- `requestContextAsyncLocalStorage`
 
 ### Client
 

diff --git a/packages/core/package.json b/packages/core/package.json
@@ -34,9 +34,9 @@
       "node": "./lib/server/shared/use-router-location.js",
       "default": "./lib/client/use-router-location.js"
     },
-    "./router-location-async-local-storage": {
-      "@mfng:internal": "./src/server/shared/router-location-async-local-storage.ts",
-      "default": "./lib/server/shared/router-location-async-local-storage.js"
+    "./request-context-async-local-storage": {
+      "@mfng:internal": "./src/server/shared/request-context-async-local-storage.ts",
+      "default": "./lib/server/shared/request-context-async-local-storage.js"
     }
   },
   "files": [

diff --git a/packages/core/src/server/rsc/index.ts b/packages/core/src/server/rsc/index.ts
@@ -1,3 +1,4 @@
 export * from './create-rsc-action-stream.js';
 export * from './create-rsc-app-stream.js';
 export * from './create-rsc-form-state.js';
+export * from './use-postpone.js';
diff --git a/packages/core/src/server/rsc/use-postpone.ts b/packages/core/src/server/rsc/use-postpone.ts
@@ -0,0 +1,16 @@
+import * as React from 'react';
+import {requestContextAsyncLocalStorage} from '../shared/request-context-async-local-storage.js';
+
+export function usePostpone(): void {
+  const requestContext = requestContextAsyncLocalStorage.getStore();
+
+  if (!requestContext) {
+    throw new Error(
+      `usePostpone() was called outside of an asynchronous context initialized by calling requestContextAsyncLocalStorage.run()`,
+    );
+  }
+
+  if (requestContext.isPrerender) {
+    return React.unstable_postpone();
+  }
+}
diff --git a/packages/core/src/server/shared/request-context-async-local-storage.ts b/packages/core/src/server/shared/request-context-async-local-storage.ts
@@ -0,0 +1,10 @@
+import {AsyncLocalStorage} from 'node:async_hooks';
+import type {RouterLocation} from '../../use-router-location.js';
+
+export interface RequestContext {
+  readonly routerLocation: RouterLocation;
+  readonly isPrerender?: boolean;
+}
+
+export const requestContextAsyncLocalStorage =
+  new AsyncLocalStorage<RequestContext>();
diff --git a/packages/core/src/server/shared/router-location-async-local-storage.ts b/packages/core/src/server/shared/router-location-async-local-storage.ts