Fix Astro HMR bottleneck (#1684)

packages/astro/src/core/create-vite.ts

@@ -44,7 +44,7 @@ export async function createVite(inlineConfig: ViteConfigWithSSR, { astroConfig,
     clearScreen: false, // we want to control the output, not Vite
     logLevel: 'error', // log errors only
     optimizeDeps: {
-      entries: ['src/**/*'] // Try and scan a user’s project (won’t catch everything),
+      entries: ['src/**/*'], // Try and scan a user’s project (won’t catch everything),
     },
     plugins: [
       astroVitePlugin({ config: astroConfig, devServer }),

packages/astro/src/core/dev/index.ts

@@ -79,24 +79,7 @@ export class AstroDevServer {
     this.app.use((req, res, next) => this.renderError(req, res, next));
 
     // Listen on port (and retry if taken)
-    await new Promise<void>((resolve, reject) => {
-      const onError = (err: NodeJS.ErrnoException) => {
-        if (err.code && err.code === 'EADDRINUSE') {
-          info(this.logging, 'astro', msg.portInUse({ port: this.port }));
-          this.port++;
-        } else {
-          error(this.logging, 'astro', err.stack);
-          this.httpServer?.removeListener('error', onError);
-          reject(err);
-        }
-      };
-      this.httpServer = this.app.listen(this.port, this.hostname, () => {
-        info(this.logging, 'astro', msg.devStart({ startupTime: performance.now() - devStart }));
-        info(this.logging, 'astro', msg.devHost({ host: `http://${this.hostname}:${this.port}` }));
-        resolve();
-      });
-      this.httpServer.on('error', onError);
-    });
+    await this.listen(devStart);
   }
 
   async stop() {
@@ -158,6 +141,38 @@ export class AstroDevServer {
     }
   }
 
+  /** Expose dev server to this.port */
+  public listen(devStart: number): Promise<void> {
+    let showedPortTakenMsg = false;
+    return new Promise<void>((resolve, reject) => {
+      const listen = () => {
+        this.httpServer = this.app.listen(this.port, this.hostname, () => {
+          info(this.logging, 'astro', msg.devStart({ startupTime: performance.now() - devStart }));
+          info(this.logging, 'astro', msg.devHost({ host: `http://${this.hostname}:${this.port}` }));
+          resolve();
+        });
+        this.httpServer?.on('error', onError);
+      };
+
+      const onError = (err: NodeJS.ErrnoException) => {
+        if (err.code && err.code === 'EADDRINUSE') {
+          if (!showedPortTakenMsg) {
+            info(this.logging, 'astro', msg.portInUse({ port: this.port }));
+            showedPortTakenMsg = true; // only print this once
+          }
+          this.port++;
+          return listen(); // retry
+        } else {
+          error(this.logging, 'astro', err.stack);
+          this.httpServer?.removeListener('error', onError);
+          reject(err); // reject
+        }
+      };
+
+      listen();
+    });
+  }
+
   private async createViteServer() {
     const viteConfig = await createVite(
       {
@@ -205,16 +220,7 @@ export class AstroDevServer {
 
     let pathname = req.url || '/'; // original request
     const reqStart = performance.now();
-
-    if (pathname.startsWith('/@astro')) {
-      const spec = pathname.slice(2);
-      const url = await this.viteServer.moduleGraph.resolveUrl(spec);
-      req.url = url[1];
-      return this.viteServer.middlewares.handle(req, res, next);
-    }
-
     let filePath: URL | undefined;
-
     try {
       const route = matchRoute(pathname, this.manifest);
 

packages/astro/src/core/ssr/css.ts

@@ -33,31 +33,3 @@ export function getStylesForID(id: string, viteServer: vite.ViteDevServer): Set<
 
   return css;
 }
-
-/** add CSS <link> tags to HTML */
-export function addLinkTagsToHTML(html: string, styles: Set<string>): string {
-  let output = html;
-
-  try {
-    // get position of </head>
-    let headEndPos = -1;
-    const parser = new htmlparser2.Parser({
-      onclosetag(tagname) {
-        if (tagname === 'head') {
-          headEndPos = parser.startIndex;
-        }
-      },
-    });
-    parser.write(html);
-    parser.end();
-
-    // update html
-    if (headEndPos !== -1) {
-      output = html.substring(0, headEndPos) + [...styles].map((href) => `<link rel="stylesheet" type="text/css" href="${href}">`).join('') + html.substring(headEndPos);
-    }
-  } catch (err) {
-    // on invalid HTML, do nothing
-  }
-
-  return output;
-}

packages/astro/src/core/ssr/html.ts

@@ -0,0 +1,112 @@
+import type vite from '../vite';
+
+import htmlparser2 from 'htmlparser2';
+
+/** Inject tags into HTML (note: for best performance, group as many tags as possible into as few calls as you can) */
+export function injectTags(html: string, tags: vite.HtmlTagDescriptor[]): string {
+  // TODO: this usually takes 5ms or less, but if it becomes a bottleneck we can create a WeakMap cache
+  let output = html;
+  if (!tags.length) return output;
+
+  const pos = { 'head-prepend': -1, head: -1, 'body-prepend': -1, body: -1 };
+
+  try {
+    // parse html
+    const parser = new htmlparser2.Parser({
+      onopentag(tagname) {
+        if (tagname === 'head') pos['head-prepend'] = parser.endIndex + 1;
+        if (tagname === 'body') pos['body-prepend'] = parser.endIndex + 1;
+      },
+      onclosetag(tagname) {
+        if (tagname === 'head') pos['head'] = parser.startIndex;
+        if (tagname === 'body') pos['body'] = parser.startIndex;
+      },
+    });
+    parser.write(html);
+    parser.end();
+
+    // inject
+    const lastToFirst = Object.entries(pos).sort((a, b) => b[1] - a[1]);
+    lastToFirst.forEach(([name, i]) => {
+      if (i === -1) {
+        // TODO: warn on missing tag? Is this an HTML partial?
+        return;
+      }
+      let selected = tags.filter(({ injectTo }) => {
+        if (name === 'head-prepend' && !injectTo) {
+          return true; // "head-prepend" is the default
+        } else {
+          return injectTo === name;
+        }
+      });
+      if (!selected.length) return;
+      output = output.substring(0, i) + serializeTags(selected) + html.substring(i);
+    });
+  } catch (err) {
+    // on invalid HTML, do nothing
+  }
+
+  return output;
+}
+
+// Everything below © Vite
+// https://github.com/vitejs/vite/blob/main/packages/vite/src/node/plugins/html.ts
+
+// Vite is released under the MIT license:
+
+// MIT License
+
+// Copyright (c) 2019-present, Yuxi (Evan) You and Vite contributors
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+const unaryTags = new Set(['link', 'meta', 'base']);
+
+function serializeTag({ tag, attrs, children }: vite.HtmlTagDescriptor, indent = ''): string {
+  if (unaryTags.has(tag)) {
+    return `<${tag}${serializeAttrs(attrs)}>`;
+  } else {
+    return `<${tag}${serializeAttrs(attrs)}>${serializeTags(children, incrementIndent(indent))}</${tag}>`;
+  }
+}
+
+function serializeTags(tags: vite.HtmlTagDescriptor['children'], indent = ''): string {
+  if (typeof tags === 'string') {
+    return tags;
+  } else if (tags && tags.length) {
+    return tags.map((tag) => `${indent}${serializeTag(tag, indent)}\n`).join('');
+  }
+  return '';
+}
+
+function serializeAttrs(attrs: vite.HtmlTagDescriptor['attrs']): string {
+  let res = '';
+  for (const key in attrs) {
+    if (typeof attrs[key] === 'boolean') {
+      res += attrs[key] ? ` ${key}` : ``;
+    } else {
+      res += ` ${key}=${JSON.stringify(attrs[key])}`;
+    }
+  }
+  return res;
+}
+
+function incrementIndent(indent = '') {
+  return `${indent}${indent[0] === '\t' ? '\t' : '  '}`;
+}

packages/astro/src/core/ssr/index.ts

@@ -1,5 +1,5 @@
 import type { BuildResult } from 'esbuild';
-import type { ViteDevServer } from '../vite';
+import type vite from '../vite';
 import type { AstroConfig, ComponentInstance, GetStaticPathsResult, Params, Props, Renderer, RouteCache, RouteData, RuntimeMode, SSRError } from '../../@types/astro-core';
 import type { AstroGlobal, TopLevelAstro, SSRResult, SSRElement } from '../../@types/astro-runtime';
 import type { LogOptions } from '../logger';
@@ -9,7 +9,8 @@ import fs from 'fs';
 import path from 'path';
 import { renderPage, renderSlot } from '../../runtime/server/index.js';
 import { canonicalURL as getCanonicalURL, codeFrame, resolveDependency } from '../util.js';
-import { addLinkTagsToHTML, getStylesForID } from './css.js';
+import { getStylesForID } from './css.js';
+import { injectTags } from './html.js';
 import { generatePaginateFunction } from './paginate.js';
 import { getParams, validateGetStaticPathsModule, validateGetStaticPathsResult } from './routing.js';
 
@@ -31,13 +32,13 @@ interface SSROptions {
   /** pass in route cache because SSR can’t manage cache-busting */
   routeCache: RouteCache;
   /** Vite instance */
-  viteServer: ViteDevServer;
+  viteServer: vite.ViteDevServer;
 }
 
 const cache = new Map<string, Promise<Renderer>>();
 
 // TODO: improve validation and error handling here.
-async function resolveRenderer(viteServer: ViteDevServer, renderer: string, astroConfig: AstroConfig) {
+async function resolveRenderer(viteServer: vite.ViteDevServer, renderer: string, astroConfig: AstroConfig) {
   const resolvedRenderer: any = {};
   // We can dynamically import the renderer by itself because it shouldn't have
   // any non-standard imports, the index is just meta info.
@@ -58,7 +59,7 @@ async function resolveRenderer(viteServer: ViteDevServer, renderer: string, astr
   return completedRenderer;
 }
 
-async function resolveRenderers(viteServer: ViteDevServer, astroConfig: AstroConfig): Promise<Renderer[]> {
+async function resolveRenderers(viteServer: vite.ViteDevServer, astroConfig: AstroConfig): Promise<Renderer[]> {
   const ids: string[] = astroConfig.renderers;
   const renderers = await Promise.all(
     ids.map((renderer) => {
@@ -159,16 +160,36 @@ export async function ssr({ astroConfig, filePath, logging, mode, origin, pathna
 
     let html = await renderPage(result, Component, pageProps, null);
 
-    // run transformIndexHtml() in development to add HMR client to the page.
+    // inject tags
+    const tags: vite.HtmlTagDescriptor[] = [];
+
+    // inject Astro HMR client (dev only)
+    if (mode === 'development') {
+      tags.push({
+        tag: 'script',
+        attrs: { type: 'module' },
+        children: `import 'astro/runtime/client/hmr.js';`,
+        injectTo: 'head',
+      });
+    }
+
+    // inject CSS
+    [...getStylesForID(fileURLToPath(filePath), viteServer)].forEach((href) => {
+      tags.push({
+        tag: 'link',
+        attrs: { type: 'text/css', rel: 'stylesheet', href },
+        injectTo: 'head',
+      });
+    });
+
+    // add injected tags
+    html = injectTags(html, tags);
+
+    // run transformIndexHtml() in dev to run Vite dev transformations
     if (mode === 'development') {
       html = await viteServer.transformIndexHtml(fileURLToPath(filePath), html, pathname);
     }
 
-    // insert CSS imported from Astro and JS components
-    const styles = getStylesForID(fileURLToPath(filePath), viteServer);
-    const relativeStyles = new Set<string>([...styles].map((url) => url.replace(fileURLToPath(astroConfig.projectRoot), '/')));
-    html = addLinkTagsToHTML(html, relativeStyles);
-
     return html;
   } catch (e: any) {
     viteServer.ssrFixStacktrace(e);
@@ -193,4 +214,3 @@ ${frame}
     throw e;
   }
 }
-

packages/astro/src/runtime/client/hmr.ts

@@ -1,5 +1,3 @@
-import '@vite/client';
-
 if (import.meta.hot) {
   const parser = new DOMParser();
   import.meta.hot.on('astro:reload', async ({ html }: { html: string }) => {

packages/astro/src/vite-plugin-astro/index.ts

@@ -103,18 +103,5 @@ export default function astro({ config, devServer }: AstroPluginOptions): vite.P
         return devServer.handleHotUpdate(context);
       }
     },
-    transformIndexHtml() {
-      // note: this runs only in dev
-      return [
-        {
-          injectTo: 'head-prepend',
-          tag: 'script',
-          attrs: {
-            type: 'module',
-            src: '/@astro/runtime/client/hmr',
-          },
-        },
-      ];
-    },
   };
 }