From ae9dc254306b6cedec336f6a0b8faf5a44bae026 Mon Sep 17 00:00:00 2001
From: MadCcc <1075746765@qq.com>
Date: Mon, 1 Apr 2024 14:24:08 +0800
Subject: [PATCH] feat: SSR support useServerInsertedHTML (#12247)

* feat: SSR support useServerInsertedHTML

* feat: ssr insert html

* fix: string template

* chore: update lock
---
 examples/ssr-demo/package.json                |   2 +
 examples/ssr-demo/src/layouts/index.tsx       |  23 ++++
 examples/ssr-demo/src/pages/index.tsx         |  14 ++-
 .../routePreloadOnLoad/routePreloadOnLoad.ts  |   2 +-
 packages/server/src/ssr.ts                    | 106 +++++++++++++-----
 pnpm-lock.yaml                                |  38 ++++---
 6 files changed, 137 insertions(+), 48 deletions(-)
 create mode 100644 examples/ssr-demo/src/layouts/index.tsx

diff --git a/examples/ssr-demo/package.json b/examples/ssr-demo/package.json
index 564ea9a69734..3b81a184639b 100644
--- a/examples/ssr-demo/package.json
+++ b/examples/ssr-demo/package.json
@@ -9,6 +9,8 @@
     "start:prod": "node ./production-server.js"
   },
   "dependencies": {
+    "@ant-design/cssinjs": "^1.18.5",
+    "antd": "^5",
     "express": "4.18.2",
     "umi": "workspace:*"
   }
diff --git a/examples/ssr-demo/src/layouts/index.tsx b/examples/ssr-demo/src/layouts/index.tsx
new file mode 100644
index 000000000000..06c6f7ea086c
--- /dev/null
+++ b/examples/ssr-demo/src/layouts/index.tsx
@@ -0,0 +1,23 @@
+import { createCache, extractStyle, StyleProvider } from '@ant-design/cssinjs';
+import { useState } from 'react';
+import { Outlet, useServerInsertedHTML } from 'umi';
+
+export default function Layout() {
+  const [cssCache] = useState(() => createCache());
+
+  useServerInsertedHTML(() => {
+    const style = extractStyle(cssCache, { plain: true });
+    return (
+      <style
+        id="antd-cssinjs"
+        dangerouslySetInnerHTML={{ __html: style }}
+      ></style>
+    );
+  });
+
+  return (
+    <StyleProvider cache={cssCache}>
+      <Outlet />
+    </StyleProvider>
+  );
+}
diff --git a/examples/ssr-demo/src/pages/index.tsx b/examples/ssr-demo/src/pages/index.tsx
index 279b00b9d5c1..c12ea05b1ed6 100644
--- a/examples/ssr-demo/src/pages/index.tsx
+++ b/examples/ssr-demo/src/pages/index.tsx
@@ -1,3 +1,5 @@
+import { Input } from 'antd';
+import { useId } from 'react';
 import {
   Link,
   MetadataLoader,
@@ -22,16 +24,26 @@ export default function HomePage() {
   const serverLoaderData = useServerLoaderData<typeof serverLoader>();
 
   useServerInsertedHTML(() => {
-    return <div>inserted html</div>;
+    return (
+      <style
+        dangerouslySetInnerHTML={{
+          __html: `.server_inserted_style { color: #1677ff }`,
+        }}
+      ></style>
+    );
   });
 
+  const id = useId();
+
   return (
     <div>
       <h1 className="title">Hello~</h1>
+      <p className="server_inserted_style">id: {id}</p>
       <p className={styles.blue}>This is index.tsx</p>
       <p className={cssStyle.title}>I should be pink</p>
       <p className={cssStyle.blue}>I should be cyan</p>
       <Button />
+      <Input placeholder="这个样式不应该闪烁" />
       <img src={bigImage} alt="" />
       <img src={umiLogo} alt="umi" />
       <Link to="/users/user">/users/user</Link>
diff --git a/packages/preset-umi/src/features/routePreloadOnLoad/routePreloadOnLoad.ts b/packages/preset-umi/src/features/routePreloadOnLoad/routePreloadOnLoad.ts
index a42eebe7ffce..04d623b3570f 100644
--- a/packages/preset-umi/src/features/routePreloadOnLoad/routePreloadOnLoad.ts
+++ b/packages/preset-umi/src/features/routePreloadOnLoad/routePreloadOnLoad.ts
@@ -256,7 +256,7 @@ export default (api: IApi) => {
               },
             ];
       }
-  
+
       return [];
     },
     stage: Infinity,
diff --git a/packages/server/src/ssr.ts b/packages/server/src/ssr.ts
index d25f0c7ceb30..bdc65802dfef 100644
--- a/packages/server/src/ssr.ts
+++ b/packages/server/src/ssr.ts
@@ -50,10 +50,9 @@ interface IExecMetaLoaderOpts extends IExecLoaderOpts {
   serverLoaderData?: any;
 }
 
-const createJSXProvider = (
-  Provider: any,
-  serverInsertedHTMLCallbacks: Set<() => React.ReactNode>,
-) => {
+const createJSXProvider = (Provider: any) => {
+  const serverInsertedHTMLCallbacks: Set<() => React.ReactNode> = new Set();
+
   const JSXProvider = (props: any) => {
     const addInsertedHtml = React.useCallback(
       (handler: () => React.ReactNode) => {
@@ -67,7 +66,7 @@ const createJSXProvider = (
       value: addInsertedHtml,
     });
   };
-  return JSXProvider;
+  return [JSXProvider, serverInsertedHTMLCallbacks] as const;
 };
 
 function createJSXGenerator(opts: CreateRequestHandlerOptions) {
@@ -163,16 +162,27 @@ function createJSXGenerator(opts: CreateRequestHandlerOptions) {
   };
 }
 
+const SERVER_INSERTED_HTML = 'umi-server-inserted-html';
 const getGenerateStaticHTML = (
-  serverInsertedHTMLCallbacks?: Set<() => React.ReactNode>,
+  serverInsertedHTMLCallbacks: Set<() => React.ReactNode>,
+  opts?: {
+    wrapper?: boolean;
+  },
 ) => {
+  const children = React.createElement(React.Fragment, {
+    children: Array.from(serverInsertedHTMLCallbacks || []).map((callback) =>
+      callback(),
+    ),
+  });
   return (
     ReactDomServer.renderToString(
-      React.createElement(React.Fragment, {
-        children: Array.from(serverInsertedHTMLCallbacks || []).map(
-          (callback) => callback(),
-        ),
-      }),
+      opts?.wrapper
+        ? React.createElement(
+            'div',
+            { id: SERVER_INSERTED_HTML, hidden: true },
+            children,
+          )
+        : children,
     ) || ''
   );
 };
@@ -184,11 +194,8 @@ export function createMarkupGenerator(opts: CreateRequestHandlerOptions) {
     const jsx = await jsxGeneratorDeferrer(url);
     if (jsx) {
       return new Promise(async (resolve, reject) => {
-        const serverInsertedHTMLCallbacks: Set<() => React.ReactNode> =
-          new Set();
-        const JSXProvider = createJSXProvider(
+        const [JSXProvider, serverInsertedHTMLCallbacks] = createJSXProvider(
           opts.ServerInsertedHTMLContext.Provider,
-          serverInsertedHTMLCallbacks,
         );
 
         let chunks: Buffer[] = [];
@@ -200,7 +207,10 @@ export function createMarkupGenerator(opts: CreateRequestHandlerOptions) {
         };
         writable.on('finish', async () => {
           let html = Buffer.concat(chunks).toString('utf8');
-          html += await getGenerateStaticHTML(serverInsertedHTMLCallbacks);
+          const serverHTML = getGenerateStaticHTML(serverInsertedHTMLCallbacks);
+          if (serverHTML) {
+            html = html.replace(/<\/head>/, `${serverHTML}</head>`);
+          }
           // append helmet tags to head
           if (opts.helmetContext) {
             html = html.replace(
@@ -266,9 +276,11 @@ export default function createRequestHandler(
       otherwise(): Promise<void> | void;
     };
 
+    const replaceServerHTMLScript = `<script>!function(){var e=document.getElementById("${SERVER_INSERTED_HTML}");e&&(Array.from(e.children).forEach(e=>{document.head.appendChild(e)}),e.remove())}();</script>`;
+
     if (typeof FetchEvent !== 'undefined' && args[0] instanceof FetchEvent) {
       // worker mode
-      const [ev, opts] = args as IWorkerRequestHandlerArgs;
+      const [ev, workerOpts] = args as IWorkerRequestHandlerArgs;
       const { pathname, searchParams } = new URL(ev.request.url);
 
       ret = {
@@ -290,16 +302,19 @@ export default function createRequestHandler(
           });
 
           // allow modify response
-          if (opts?.modifyResponse) {
-            res = await opts.modifyResponse(res);
+          if (workerOpts?.modifyResponse) {
+            res = await workerOpts.modifyResponse(res);
           }
 
           ev.respondWith(res);
         },
         async sendPage(jsx) {
+          const [JSXProvider, serverInsertedHTMLCallbacks] = createJSXProvider(
+            opts.ServerInsertedHTMLContext.Provider,
+          );
           // handle route path request
           const stream = await ReactDomServer.renderToReadableStream(
-            jsx.element,
+            React.createElement(JSXProvider, undefined, jsx.element),
             {
               bootstrapScripts: [jsx.manifest.assets['umi.js'] || '/umi.js'],
               onError(x: any) {
@@ -307,6 +322,22 @@ export default function createRequestHandler(
               },
             },
           );
+
+          const transformStream = new TransformStream({
+            flush(controller) {
+              if (serverInsertedHTMLCallbacks.size) {
+                const serverHTML = getGenerateStaticHTML(
+                  serverInsertedHTMLCallbacks,
+                  { wrapper: true },
+                );
+                controller.enqueue(serverHTML);
+                controller.enqueue(replaceServerHTMLScript);
+              }
+            },
+          });
+
+          stream.pipeThrough(transformStream);
+
           let res = new Response(stream, {
             headers: {
               'content-type': 'text/html; charset=utf-8',
@@ -315,8 +346,8 @@ export default function createRequestHandler(
           });
 
           // allow modify response
-          if (opts?.modifyResponse) {
-            res = await opts.modifyResponse(res);
+          if (workerOpts?.modifyResponse) {
+            res = await workerOpts.modifyResponse(res);
           }
 
           ev.respondWith(res);
@@ -343,6 +374,9 @@ export default function createRequestHandler(
           res.status(200).json(data);
         },
         async sendPage(jsx) {
+          const [JSXProvider, serverInsertedHTMLCallbacks] = createJSXProvider(
+            opts.ServerInsertedHTMLContext.Provider,
+          );
           const writable = new Writable();
 
           res.type('html');
@@ -353,19 +387,29 @@ export default function createRequestHandler(
           };
 
           writable.on('finish', async () => {
-            res.write(getGenerateStaticHTML());
+            if (serverInsertedHTMLCallbacks.size) {
+              res.write(
+                getGenerateStaticHTML(serverInsertedHTMLCallbacks, {
+                  wrapper: true,
+                }),
+              );
+              res.write(replaceServerHTMLScript);
+            }
             res.end();
           });
 
-          const stream = ReactDomServer.renderToPipeableStream(jsx.element, {
-            bootstrapScripts: [jsx.manifest.assets['umi.js'] || '/umi.js'],
-            onShellReady() {
-              stream.pipe(writable);
-            },
-            onError(x: any) {
-              console.error(x);
+          const stream = ReactDomServer.renderToPipeableStream(
+            React.createElement(JSXProvider, undefined, jsx.element),
+            {
+              bootstrapScripts: [jsx.manifest.assets['umi.js'] || '/umi.js'],
+              onShellReady() {
+                stream.pipe(writable);
+              },
+              onError(x: any) {
+                console.error(x);
+              },
             },
-          });
+          );
         },
         otherwise: next,
       };
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 1b0fde386c80..37e6c9e372ee 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1123,6 +1123,12 @@ importers:
 
   examples/ssr-demo:
     dependencies:
+      '@ant-design/cssinjs':
+        specifier: ^1.18.5
+        version: 1.18.5(react-dom@18.1.0)(react@18.1.0)
+      antd:
+        specifier: ^5
+        version: 5.8.4(react-dom@18.1.0)(react@18.1.0)
       express:
         specifier: 4.18.2
         version: 4.18.2
@@ -3112,8 +3118,8 @@ packages:
       stylis: 4.1.4
     dev: false
 
-  /@ant-design/cssinjs@1.17.0(react-dom@18.1.0)(react@18.1.0):
-    resolution: {integrity: sha512-MgGCZ6sfD3yQB0XW0hN4jgixMxApTlDYyct+pc7fRZNO4CaqWWm/9iXkkljNR27lyWLZmm+XiDfcIOo1bnrnMA==}
+  /@ant-design/cssinjs@1.18.5(react-dom@17.0.2)(react@17.0.2):
+    resolution: {integrity: sha512-Ub4n3d+MAX/qtE5S9PM8iOn5ocU7GUAIC4Adc2X8UCMXnsRRfpJBHsBdtQ1qoAuaQ7lU2M1BTCuJ+fkv4fOWiw==}
     peerDependencies:
       react: '>=16.0.0'
       react-dom: '>=16.0.0'
@@ -3128,14 +3134,14 @@ packages:
       '@emotion/unitless': 0.7.5
       classnames: 2.3.2
       csstype: 3.1.3
-      rc-util: 5.38.1(react-dom@18.1.0)(react@18.1.0)
-      react: 18.1.0
-      react-dom: 18.1.0(react@18.1.0)
+      rc-util: 5.38.1(react-dom@17.0.2)(react@17.0.2)
+      react: 17.0.2
+      react-dom: 17.0.2(react@17.0.2)
       stylis: 4.3.0
     dev: true
 
-  /@ant-design/cssinjs@1.18.1(react-dom@17.0.2)(react@17.0.2):
-    resolution: {integrity: sha512-1JURAPrsjK1GwpqByTq3bJ7nF7lbMKDZpehqeR2n8/IR5O58/W1U4VcOeaw5ZyTHri3tEMcom7dyP2tvxpW54g==}
+  /@ant-design/cssinjs@1.18.5(react-dom@18.1.0)(react@18.1.0):
+    resolution: {integrity: sha512-Ub4n3d+MAX/qtE5S9PM8iOn5ocU7GUAIC4Adc2X8UCMXnsRRfpJBHsBdtQ1qoAuaQ7lU2M1BTCuJ+fkv4fOWiw==}
     peerDependencies:
       react: '>=16.0.0'
       react-dom: '>=16.0.0'
@@ -3149,12 +3155,11 @@ packages:
       '@emotion/hash': 0.8.0
       '@emotion/unitless': 0.7.5
       classnames: 2.3.2
-      csstype: 3.1.2
-      rc-util: 5.38.1(react-dom@17.0.2)(react@17.0.2)
-      react: 17.0.2
-      react-dom: 17.0.2(react@17.0.2)
+      csstype: 3.1.3
+      rc-util: 5.38.1(react-dom@18.1.0)(react@18.1.0)
+      react: 18.1.0
+      react-dom: 18.1.0(react@18.1.0)
       stylis: 4.3.0
-    dev: true
 
   /@ant-design/cssinjs@1.9.1(react-dom@18.1.0)(react@18.1.0):
     resolution: {integrity: sha512-CZt1vCMs/sY7RoacYuIkZwQmb8Bhp99ReNNE9Y8lnUzik8fmCdKAQA7ecvVOFwmNFdcBHga7ye/XIRrsbkiqWw==}
@@ -17648,7 +17653,7 @@ packages:
     dependencies:
       '@ahooksjs/use-request': 2.8.15(react@17.0.2)
       '@ant-design/antd-theme-variable': 1.0.0
-      '@ant-design/cssinjs': 1.18.1(react-dom@17.0.2)(react@17.0.2)
+      '@ant-design/cssinjs': 1.18.5(react-dom@17.0.2)(react@17.0.2)
       '@ant-design/icons': 4.7.0(react-dom@17.0.2)(react@17.0.2)
       '@ant-design/moment-webpack-plugin': 0.0.3
       '@ant-design/pro-components': 2.3.20(antd@4.23.6)(react-dom@17.0.2)(react@17.0.2)
@@ -19548,7 +19553,7 @@ packages:
         optional: true
     dependencies:
       '@ant-design/colors': 7.0.0
-      '@ant-design/cssinjs': 1.17.0(react-dom@18.1.0)(react@18.1.0)
+      '@ant-design/cssinjs': 1.18.5(react-dom@18.1.0)(react@18.1.0)
       '@ant-design/icons': 5.0.1(react-dom@18.1.0)(react@18.1.0)
       '@ant-design/react-slick': 1.0.0(react@18.1.0)
       '@babel/runtime': 7.21.0
@@ -19747,7 +19752,7 @@ packages:
         optional: true
     dependencies:
       '@ant-design/colors': 7.0.0
-      '@ant-design/cssinjs': 1.16.0(react-dom@18.1.0)(react@18.1.0)
+      '@ant-design/cssinjs': 1.18.5(react-dom@18.1.0)(react@18.1.0)
       '@ant-design/icons': 5.2.5(react-dom@18.1.0)(react@18.1.0)
       '@ant-design/react-slick': 1.0.0(react@18.1.0)
       '@babel/runtime': 7.21.0
@@ -22496,6 +22501,7 @@ packages:
     dependencies:
       loose-envify: 1.4.0
       object-assign: 4.1.1
+    bundledDependencies: false
 
   /create-require@1.1.1:
     resolution: {integrity: sha512-dcKFX3jn0MpIaXjisoRvexIJVEKzaq7z2rZKxf+MSr9TkdmHmsU4m2lcLojrj/FHl8mk5VxMmYA+ftRkP/3oKQ==}
@@ -40724,6 +40730,7 @@ packages:
       react: 16.14.0
       scheduler: 0.19.1
     dev: false
+    bundledDependencies: false
 
   /react-dom@17.0.2(react@17.0.2):
     resolution: {integrity: sha512-s4h96KtLDUQlsENhMn1ar8t2bEa+q/YAtj8pPPdIjPDGBDIVNsrD9aXNWqspUe6AzKCIG0C1HZZLqLV7qpOBGA==}
@@ -41371,6 +41378,7 @@ packages:
       object-assign: 4.1.1
       prop-types: 15.8.1
     dev: false
+    bundledDependencies: false
 
   /react@17.0.2:
     resolution: {integrity: sha512-gnhPt75i/dq/z3/6q/0asP78D0u592D5L1pd7M8P+dck6Fu/jJeL6iVVK23fptSUZj8Vjf++7wXA8UNclGQcbA==}